perf: 依赖数据库 CASCADE/SET NULL 替代手动清理关联表,缩短删除事务

- 批量删除移除 cleanup_key_references 手动清理,改为依赖 FK CASCADE/SET NULL
- video_tasks.key_id FK 增加 ondelete="SET NULL",附带幂等迁移脚本
- _sync_delete 增加 statement_timeout 和任务级超时保护
- 批量导入在每次 await 前释放闲置 DB 连接,按批次提交写入避免长事务
- 前端轮询改为先查后等,首次查询不再多等一个间隔
This commit is contained in:
fawney19
2026-03-09 03:48:20 +08:00
parent 654ce89541
commit fa69287449
8 changed files with 468 additions and 18 deletions

View File

@@ -42,6 +42,7 @@ from src.database.database import get_db
from src.models.database import Provider, ProviderAPIKey, User
from src.services.provider.pool.config import parse_pool_config
from src.services.provider_keys.auth_type import OAUTH_AUTH_TYPES
from src.services.scheduling.utils import release_db_connection_before_await
from src.utils.auth_utils import require_admin
router = APIRouter(prefix="/api/admin/provider-oauth", tags=["Provider OAuth"])
@@ -129,6 +130,7 @@ _PROVIDER_OAUTH_BATCH_TASK_TTL_SECONDS = 24 * 3600
_PROVIDER_OAUTH_BATCH_TASK_MAX_ERROR_SAMPLES = 20
_PROVIDER_OAUTH_DEFAULT_TIMEOUT_SECONDS = 30.0
_PROVIDER_OAUTH_BATCH_IMPORT_PROXY_TIMEOUT_SECONDS = 60.0
_PROVIDER_OAUTH_BATCH_IMPORT_COMMIT_BATCH_SIZE = 25
_PROVIDER_OAUTH_BATCH_TASK_ALLOWED_STATUSES = {
"submitted",
"processing",
@@ -1639,6 +1641,27 @@ def _estimate_batch_import_total(provider_type: str, raw_credentials: str) -> in
return len(_parse_standard_oauth_import_entries(raw_credentials))
def _release_batch_import_db_connection_before_await(db: Session) -> None:
"""Best-effort 释放批量导入任务的只读 DB 连接。
批量导入会在单个后台任务里执行大量 await上游 token 校验、邮箱探测、Redis 进度更新)。
如果前面做过 Provider 查询而 Session 一直保持事务打开,连接会长时间占着不放,
在大批量导入且多数条目最终失败时尤其容易把连接池拖满。
这里复用调度器已有的 helper仅在 Session 没有挂起写入时才提前结束事务,
避免影响 flush 后尚未提交的数据。
"""
release_db_connection_before_await(db)
def _commit_batch_import_writes_if_needed(db: Session, pending_writes: int) -> int:
"""按固定批次提交导入写入,避免长事务持续占用连接。"""
if pending_writes < _PROVIDER_OAUTH_BATCH_IMPORT_COMMIT_BATCH_SIZE:
return pending_writes
db.commit()
return 0
def _apply_codex_import_hints(auth_config: dict[str, Any], import_entry: dict[str, str]) -> None:
"""将导入文件中可用的 Codex 账号信息作为兜底补全(不覆盖已有值)。"""
for field in ("account_id", "plan_type", "user_id", "email"):
@@ -1901,11 +1924,14 @@ async def _batch_import_standard_oauth_internal(
success_count = 0
failed_count = 0
processed_count = 0
pending_success_writes = 0
db_lock = asyncio.Lock()
sem = asyncio.Semaphore(max(concurrency, 1))
_release_batch_import_db_connection_before_await(db)
async def _process_entry(idx: int, import_entry: dict[str, Any]) -> None:
nonlocal success_count, failed_count, processed_count
nonlocal success_count, failed_count, processed_count, pending_success_writes
result_item: BatchImportResultItem
async with sem:
@@ -1951,6 +1977,7 @@ async def _batch_import_standard_oauth_internal(
json_body = None
try:
_release_batch_import_db_connection_before_await(db)
resp = await post_oauth_token(
provider_type=provider_type,
token_url=token_url,
@@ -1970,6 +1997,7 @@ async def _batch_import_standard_oauth_internal(
processed_count += 1
results[idx] = result_item
if progress_hook is not None:
_release_batch_import_db_connection_before_await(db)
await progress_hook(
total, processed_count, success_count, failed_count, result_item
)
@@ -1997,6 +2025,7 @@ async def _batch_import_standard_oauth_internal(
processed_count += 1
results[idx] = result_item
if progress_hook is not None:
_release_batch_import_db_connection_before_await(db)
await progress_hook(
total, processed_count, success_count, failed_count, result_item
)
@@ -2016,6 +2045,7 @@ async def _batch_import_standard_oauth_internal(
processed_count += 1
results[idx] = result_item
if progress_hook is not None:
_release_batch_import_db_connection_before_await(db)
await progress_hook(
total, processed_count, success_count, failed_count, result_item
)
@@ -2039,6 +2069,7 @@ async def _batch_import_standard_oauth_internal(
}
try:
_release_batch_import_db_connection_before_await(db)
auth_config = await enrich_auth_config(
provider_type=provider_type,
auth_config=auth_config,
@@ -2067,6 +2098,7 @@ async def _batch_import_standard_oauth_internal(
processed_count += 1
results[idx] = result_item
if progress_hook is not None:
_release_batch_import_db_connection_before_await(db)
await progress_hook(
total,
processed_count,
@@ -2108,6 +2140,11 @@ async def _batch_import_standard_oauth_internal(
proxy=key_proxy,
)
pending_success_writes += 1
pending_success_writes = _commit_batch_import_writes_if_needed(
db, pending_success_writes
)
result_item = BatchImportResultItem(
index=idx,
status="success",
@@ -2129,6 +2166,7 @@ async def _batch_import_standard_oauth_internal(
processed_count += 1
results[idx] = result_item
if progress_hook is not None:
_release_batch_import_db_connection_before_await(db)
await progress_hook(total, processed_count, success_count, failed_count, result_item)
await asyncio.gather(
@@ -2136,7 +2174,7 @@ async def _batch_import_standard_oauth_internal(
return_exceptions=True,
)
if success_count > 0:
if success_count > 0 and pending_success_writes > 0:
db.commit()
final_results = [r for r in results if r is not None]
@@ -2446,11 +2484,14 @@ async def _batch_import_kiro_internal(
success_count = 0
failed_count = 0
processed_count = 0
pending_success_writes = 0
db_lock = asyncio.Lock()
sem = asyncio.Semaphore(max(concurrency, 1))
_release_batch_import_db_connection_before_await(db)
async def _process_entry(idx: int, cred: dict[str, Any]) -> None:
nonlocal success_count, failed_count, processed_count
nonlocal success_count, failed_count, processed_count, pending_success_writes
result_item: BatchImportResultItem
async with sem:
@@ -2466,6 +2507,7 @@ async def _batch_import_kiro_internal(
processed_count += 1
results[idx] = result_item
if progress_hook is not None:
_release_batch_import_db_connection_before_await(db)
await progress_hook(
total, processed_count, success_count, failed_count, result_item
)
@@ -2475,6 +2517,7 @@ async def _batch_import_kiro_internal(
cfg.provider_type = ProviderType.KIRO.value
try:
_release_batch_import_db_connection_before_await(db)
access_token, new_cfg = await refresh_access_token(
cfg,
proxy_config=proxy_config,
@@ -2490,11 +2533,13 @@ async def _batch_import_kiro_internal(
processed_count += 1
results[idx] = result_item
if progress_hook is not None:
_release_batch_import_db_connection_before_await(db)
await progress_hook(
total, processed_count, success_count, failed_count, result_item
)
return
_release_batch_import_db_connection_before_await(db)
email = await _fetch_kiro_email(new_cfg.to_dict(), proxy_config=proxy_config)
if email and not new_cfg.email:
new_cfg.email = email
@@ -2514,6 +2559,7 @@ async def _batch_import_kiro_internal(
processed_count += 1
results[idx] = result_item
if progress_hook is not None:
_release_batch_import_db_connection_before_await(db)
await progress_hook(
total,
processed_count,
@@ -2550,6 +2596,11 @@ async def _batch_import_kiro_internal(
proxy=key_proxy,
)
pending_success_writes += 1
pending_success_writes = _commit_batch_import_writes_if_needed(
db, pending_success_writes
)
result_item = BatchImportResultItem(
index=idx,
status="success",
@@ -2572,6 +2623,7 @@ async def _batch_import_kiro_internal(
processed_count += 1
results[idx] = result_item
if progress_hook is not None:
_release_batch_import_db_connection_before_await(db)
await progress_hook(total, processed_count, success_count, failed_count, result_item)
await asyncio.gather(
@@ -2580,7 +2632,7 @@ async def _batch_import_kiro_internal(
)
# 提交所有成功的记录
if success_count > 0:
if success_count > 0 and pending_success_writes > 0:
db.commit()
final_results = [r for r in results if r is not None]

View File

@@ -1950,7 +1950,7 @@ class VideoTask(Base):
api_key_name = Column(String(200), nullable=True, comment="API Key 名称快照")
provider_id = Column(String(36), ForeignKey("providers.id"), index=True)
endpoint_id = Column(String(36), ForeignKey("provider_endpoints.id"), index=True)
key_id = Column(String(36), ForeignKey("provider_api_keys.id"), index=True)
key_id = Column(String(36), ForeignKey("provider_api_keys.id", ondelete="SET NULL"), index=True)
# 格式转换追踪
client_api_format = Column(String(50), nullable=False)

View File

@@ -8,12 +8,14 @@ from __future__ import annotations
import asyncio
import json
import time
import uuid
from collections.abc import Callable
from concurrent.futures import Future
import redis.asyncio as aioredis
from sqlalchemy import delete as sa_delete
from sqlalchemy import text
from src.clients.redis_client import get_redis_client
from src.core.logger import logger
@@ -30,6 +32,12 @@ _TASK_RETAIN_SECONDS = 600
# 批量删除时每个独立事务处理的 key 数量
_CLEANUP_BATCH_SIZE = 50
# 单个批次的数据库 statement 超时(秒)
_BATCH_STATEMENT_TIMEOUT_S = 30
# 整个任务的最大执行时间(秒)
_TASK_TIMEOUT_S = 600
# Redis key 前缀
_REDIS_KEY_PREFIX = "batch_delete_task"
@@ -175,23 +183,35 @@ def _sync_delete(
) -> int:
"""在线程中执行的同步删除逻辑,避免阻塞事件循环。
按小批次_CLEANUP_BATCH_SIZE清理关联表并删除 key
每个批次独立事务,单批失败跳过并继续,确保进度条持续推进
按小批次_CLEANUP_BATCH_SIZE直接删除 key依赖数据库 CASCADE/SET NULL 自动清理关联表。
每个批次独立事务,单批失败跳过并继续。
"""
from src.database import create_session
from src.models.database import ProviderAPIKey
from src.services.provider_keys.key_side_effects import cleanup_key_references
db = create_session()
try:
affected = 0
total_batches = (len(key_ids) + _CLEANUP_BATCH_SIZE - 1) // _CLEANUP_BATCH_SIZE
batch_idx = 0
task_start = time.monotonic()
for i in range(0, len(key_ids), _CLEANUP_BATCH_SIZE):
# 整体超时保护
if time.monotonic() - task_start > _TASK_TIMEOUT_S:
logger.warning(
"[BATCH_DELETE] task timeout after {}s, deleted {}/{}",
_TASK_TIMEOUT_S,
affected,
len(key_ids),
)
break
batch = key_ids[i : i + _CLEANUP_BATCH_SIZE]
batch_idx += 1
try:
cleanup_key_references(db, batch)
# 设置 statement_timeout防止单条 SQL 无限等锁
timeout_ms = _BATCH_STATEMENT_TIMEOUT_S * 1000
db.execute(text(f"SET LOCAL statement_timeout = '{timeout_ms}'"))
result = db.execute(
sa_delete(ProviderAPIKey).where(
ProviderAPIKey.provider_id == provider_id,
@@ -203,7 +223,9 @@ def _sync_delete(
db.commit()
except Exception as exc:
logger.warning(
"[BATCH_DELETE] batch failed (keys {}-{}): {}",
"[BATCH_DELETE] batch {}/{} failed (keys {}-{}): {}",
batch_idx,
total_batches,
i,
i + len(batch),
exc,
@@ -212,9 +234,16 @@ def _sync_delete(
db.rollback()
except Exception:
pass
# 每 5 个批次或最后一个批次上报进度,避免过于频繁写 Redis
if progress_callback is not None and (batch_idx % 5 == 0 or batch_idx == total_batches):
# 每个批次都上报一次进度
if progress_callback is not None:
progress_callback(affected)
logger.info(
"[BATCH_DELETE] sync delete finished: deleted={}/{} batches={} elapsed={:.1f}s",
affected,
len(key_ids),
total_batches,
time.monotonic() - task_start,
)
return affected
finally:
try:
@@ -230,6 +259,12 @@ async def _run_batch_delete(
) -> None:
r = await get_redis_client(require_redis=False)
await _update_task_field(task_id, r=r, status=STATUS_RUNNING)
logger.info(
"[BATCH_DELETE_TASK] started task={} provider={} total={}",
task_id,
provider_id[:8],
len(key_ids),
)
# 从工作线程安全地触发 Redis 进度更新
loop = asyncio.get_running_loop()
@@ -245,7 +280,10 @@ async def _run_batch_delete(
pass
try:
affected = await asyncio.to_thread(_sync_delete, provider_id, key_ids, on_progress)
affected = await asyncio.wait_for(
asyncio.to_thread(_sync_delete, provider_id, key_ids, on_progress),
timeout=_TASK_TIMEOUT_S + 30, # 留一点余量给 _sync_delete 内部超时
)
if progress_futures:
results = await asyncio.gather(
@@ -292,6 +330,15 @@ async def _run_batch_delete(
affected,
)
except asyncio.TimeoutError:
msg = f"task timeout after {_TASK_TIMEOUT_S + 30}s"
await _update_task_field(task_id, r=r, status=STATUS_FAILED, message=msg)
logger.error(
"[BATCH_DELETE_TASK] {} task={} provider={}",
msg,
task_id,
provider_id[:8],
)
except asyncio.CancelledError:
await _update_task_field(
task_id,

View File

@@ -31,7 +31,6 @@ from src.services.provider.fingerprint import generate_fingerprint, normalize_fi
from src.services.provider_keys.auth_type import normalize_auth_type
from src.services.provider_keys.duplicate_check import check_duplicate_key
from src.services.provider_keys.key_side_effects import (
cleanup_key_references,
run_create_key_side_effects,
run_delete_key_side_effects,
run_update_key_side_effects,
@@ -519,12 +518,10 @@ async def batch_delete_endpoint_keys_response(db: Session, key_ids: list[str]) -
# 收集受影响的 provider_id
affected_provider_ids = {key.provider_id for key in keys if key.provider_id}
# 批量 SQL DELETE一次提交
# 批量 SQL DELETE依赖数据库 CASCADE/SET NULL 自动清理关联表
success_count = 0
try:
found_id_list = list(found_ids)
# 先清理关联表,避免 CASCADE 级联删除超时
cleanup_key_references(db, found_id_list)
db.execute(sa_delete(ProviderAPIKey).where(ProviderAPIKey.id.in_(found_id_list)))
db.commit()
success_count = len(found_ids)