fix: harden Phase 3 task leases and metering
This commit is contained in:
parent
2e99155264
commit
9bf98522af
@ -11,6 +11,10 @@ type Voice = { id: string; provider_voice_id: string; name: string; language?: s
|
||||
type Usage = { available: number; used: number; reserved: number; plan: string };
|
||||
type Task = { id: string; status: string; text_length: number; voice_id: string; parameters: { format?: string; speed?: number }; error_code?: string | null; audio_available: boolean; created_at: string; finished_at?: string | null };
|
||||
|
||||
function normalizeText(value: string) {
|
||||
return value.normalize("NFC").replace(/\r\n?/g, "\n");
|
||||
}
|
||||
|
||||
export default function TtsPage() {
|
||||
const [csrf, setCsrf] = useState("");
|
||||
const [loggedIn, setLoggedIn] = useState(false);
|
||||
@ -72,7 +76,7 @@ export default function TtsPage() {
|
||||
method: "POST",
|
||||
credentials: "include",
|
||||
headers: { "Content-Type": "application/json", "X-CSRF-Token": csrf, "Idempotency-Key": crypto.randomUUID() },
|
||||
body: JSON.stringify({ text, voice_id: voice, parameters: { format, speed: Number(speed) } }),
|
||||
body: JSON.stringify({ text: normalizeText(text), voice_id: voice, parameters: { format, speed: Number(speed) } }),
|
||||
});
|
||||
const body = await response.json().catch(() => ({}));
|
||||
if (!response.ok) throw new Error(body.error?.message ?? "任务创建失败");
|
||||
@ -91,7 +95,7 @@ export default function TtsPage() {
|
||||
<div className="mt-10"><PreviewNotice>{loggedIn ? "当前使用测试数据库和真实任务 Worker。上游未配置时任务会明确失败并释放额度。" : "生成、历史、回放和下载需要登录。"}</PreviewNotice></div>
|
||||
{!loggedIn ? <div className="mt-8"><StatusCard title="请先登录" description="登录后可使用真实任务、额度和历史功能。"><Link className="button-primary mt-5" href="/login">前往登录 <ArrowRight /></Link></StatusCard></div> : <>
|
||||
<div className="mt-8 grid gap-6 lg:grid-cols-[1.2fr_0.8fr]">
|
||||
<section className="rounded-2xl border border-line bg-panel/60 p-5 sm:p-7" aria-labelledby="tts-input-title"><div className="flex items-center justify-between gap-4"><div><p className="text-xs font-semibold uppercase tracking-[0.18em] text-subtle">01 / Input</p><h2 className="mt-2 text-xl font-semibold" id="tts-input-title">输入文本</h2></div><span className="text-xs text-subtle">{text.length} / 当前计划限制</span></div><textarea className="field-input mt-6 min-h-64 resize-y leading-7" onChange={(event) => setText(event.target.value)} placeholder="输入需要转换的文本。" value={text} /><div className="mt-5 flex flex-col gap-4 border-t border-line pt-5 sm:flex-row sm:items-center sm:justify-between"><p className="text-sm text-muted">可用额度：{usage?.available.toLocaleString() ?? "加载中"}</p><button className="button-primary" disabled={busy || !csrf} onClick={createTask} type="button">{busy ? "生成中…" : "生成语音"} <ArrowRight /></button></div></section>
|
||||
<section className="rounded-2xl border border-line bg-panel/60 p-5 sm:p-7" aria-labelledby="tts-input-title"><div className="flex items-center justify-between gap-4"><div><p className="text-xs font-semibold uppercase tracking-[0.18em] text-subtle">01 / Input</p><h2 className="mt-2 text-xl font-semibold" id="tts-input-title">输入文本</h2></div><span className="text-xs text-subtle">{normalizeText(text).length} / 当前计划限制</span></div><textarea className="field-input mt-6 min-h-64 resize-y leading-7" onChange={(event) => setText(event.target.value)} placeholder="输入需要转换的文本。" value={text} /><div className="mt-5 flex flex-col gap-4 border-t border-line pt-5 sm:flex-row sm:items-center sm:justify-between"><p className="text-sm text-muted">可用额度：{usage?.available.toLocaleString() ?? "加载中"}</p><button className="button-primary" disabled={busy || !csrf} onClick={createTask} type="button">{busy ? "生成中…" : "生成语音"} <ArrowRight /></button></div></section>
|
||||
<aside className="space-y-6"><section className="rounded-2xl border border-line bg-panel/60 p-5 sm:p-7" aria-labelledby="tts-settings-title"><p className="text-xs font-semibold uppercase tracking-[0.18em] text-subtle">02 / Parameters</p><h2 className="mt-2 text-xl font-semibold" id="tts-settings-title">声音参数</h2><div className="mt-6 space-y-5"><div><label className="field-label" htmlFor="voice">音色</label><select className="field-input" id="voice" onChange={(event) => setVoice(event.target.value)} value={voice}><option value="">选择音色</option>{voices.map((item) => <option key={item.id} value={item.provider_voice_id}>{item.name}{item.language ? ` · ${item.language}` : ""}</option>)}</select></div><div><label className="field-label" htmlFor="speed">语速</label><select className="field-input" id="speed" onChange={(event) => setSpeed(event.target.value)} value={speed}><option value="0.5">0.5x</option><option value="0.8">0.8x</option><option value="1">1.0x</option><option value="1.2">1.2x</option><option value="1.5">1.5x</option><option value="2">2.0x</option></select></div><div><label className="field-label" htmlFor="format">格式</label><select className="field-input" id="format" onChange={(event) => setFormat(event.target.value)} value={format}><option value="wav">WAV</option><option value="mp3">MP3</option></select></div></div></section><StatusCard title="额度" description={`${usage?.plan?.toUpperCase() ?? ""} · 已用 ${usage?.used.toLocaleString() ?? "-"} · 冻结 ${usage?.reserved.toLocaleString() ?? "-"}`} tone="info" /></aside>
|
||||
</div>
|
||||
{error ? <p className="mt-6 rounded-xl border border-danger/30 bg-danger/10 px-4 py-3 text-sm leading-6 text-danger" role="alert">{error}</p> : null}
|
||||
|
||||
@ -39,6 +39,7 @@ TTS_UPSTREAM_URL=
|
||||
TTS_API_KEY=
|
||||
TTS_TIMEOUT_SECONDS=120
|
||||
AUDIO_STORAGE_DIR=/home/flym/kaotings-audio
|
||||
AUDIO_RETENTION_SECONDS=604800
|
||||
EOF
|
||||
chmod 600 /etc/kaotings/api.env
|
||||
echo "database and API environment provisioned"
|
||||
|
||||
@ -9,7 +9,7 @@ User=flym
|
||||
Group=flym
|
||||
WorkingDirectory=/home/flym/kaotings-api
|
||||
EnvironmentFile=/etc/kaotings/api.env
|
||||
ExecStart=/home/flym/kaotings-api-venv/bin/uvicorn app.main:app --host 127.0.0.1 --port 8000
|
||||
ExecStart=/home/flym/kaotings-api-venv/bin/uvicorn app.main:app --host 127.0.0.1 --port 8000 --workers 1
|
||||
Restart=always
|
||||
RestartSec=3
|
||||
NoNewPrivileges=true
|
||||
|
||||
@ -13,3 +13,4 @@ TTS_UPSTREAM_URL=
|
||||
TTS_API_KEY=
|
||||
TTS_TIMEOUT_SECONDS=120
|
||||
AUDIO_STORAGE_DIR=/home/flym/kaotings-audio
|
||||
AUDIO_RETENTION_SECONDS=604800
|
||||
|
||||
@ -25,6 +25,7 @@ class Settings:
|
||||
tts_api_key: str
|
||||
tts_timeout_seconds: int
|
||||
audio_storage_dir: str
|
||||
audio_retention_seconds: int
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> "Settings":
|
||||
@ -48,6 +49,7 @@ class Settings:
|
||||
tts_api_key=os.getenv("TTS_API_KEY", ""),
|
||||
tts_timeout_seconds=int(os.getenv("TTS_TIMEOUT_SECONDS", "120")),
|
||||
audio_storage_dir=os.getenv("AUDIO_STORAGE_DIR", "./data/audio"),
|
||||
audio_retention_seconds=int(os.getenv("AUDIO_RETENTION_SECONDS", "604800")),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@ -1,6 +1,8 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import time
|
||||
import unicodedata
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
@ -118,6 +120,10 @@ def usage_view(quota: dict) -> dict[str, Any]:
|
||||
}
|
||||
|
||||
|
||||
def normalize_tts_text(value: str) -> str:
|
||||
return unicodedata.normalize("NFC", value.replace("\r\n", "\n").replace("\r", "\n"))
|
||||
|
||||
|
||||
def task_view(connection: Connection, task: dict) -> dict[str, Any]:
|
||||
audio = connection.execute("SELECT expires_at, status FROM audio_files WHERE task_id = %s", (task["id"],)).fetchone()
|
||||
return {
|
||||
@ -134,14 +140,18 @@ def claim_next_task() -> dict | None:
|
||||
task = connection.execute("SELECT * FROM tts_tasks WHERE status = 'queued' ORDER BY created_at FOR UPDATE SKIP LOCKED LIMIT 1").fetchone()
|
||||
if not task:
|
||||
return None
|
||||
updated = connection.execute("UPDATE tts_tasks SET status = 'running', started_at = now(), updated_at = now(), attempt_count = attempt_count + 1 WHERE id = %s RETURNING *", (task["id"],)).fetchone()
|
||||
lease_token = new_token()
|
||||
updated = connection.execute("UPDATE tts_tasks SET status = 'running', lease_token = %s, lease_expires_at = now() + interval '5 minutes', started_at = now(), updated_at = now(), attempt_count = attempt_count + 1 WHERE id = %s RETURNING *", (lease_token, task["id"])).fetchone()
|
||||
connection.commit()
|
||||
return updated
|
||||
|
||||
|
||||
def finish_task_failure(task: dict, code: str) -> None:
|
||||
with psycopg.connect(settings.database_url, row_factory=psycopg.rows.dict_row) as connection:
|
||||
connection.execute("UPDATE tts_tasks SET status = 'failed', error_code = %s, finished_at = now(), updated_at = now() WHERE id = %s AND status = 'running'", (code, task["id"]))
|
||||
updated = connection.execute("UPDATE tts_tasks SET status = 'failed', error_code = %s, finished_at = now(), updated_at = now(), lease_token = NULL, lease_expires_at = NULL WHERE id = %s AND status = 'running' AND lease_token = %s", (code, task["id"], task["lease_token"]))
|
||||
if updated.rowcount != 1:
|
||||
connection.rollback()
|
||||
return
|
||||
connection.execute("UPDATE quota_accounts SET reserved = GREATEST(0, reserved - %s), version = version + 1 WHERE id = %s", (task["reserved_amount"], task["quota_account_id"]))
|
||||
connection.execute("INSERT INTO usage_records(user_id, quota_account_id, task_id, type, amount) VALUES (%s, %s, %s, 'release', %s)", (task["user_id"], task["quota_account_id"], task["id"], task["reserved_amount"]))
|
||||
connection.commit()
|
||||
@ -157,13 +167,40 @@ def finish_task_success(task: dict, audio: bytes, mime_type: str) -> None:
|
||||
target.write_bytes(audio)
|
||||
checksum = hashlib.sha256(audio).hexdigest()
|
||||
with psycopg.connect(settings.database_url, row_factory=psycopg.rows.dict_row) as connection:
|
||||
current = connection.execute("SELECT status, lease_token FROM tts_tasks WHERE id = %s FOR UPDATE", (task["id"],)).fetchone()
|
||||
if not current or current["status"] != "running" or current["lease_token"] != task["lease_token"]:
|
||||
target.unlink(missing_ok=True)
|
||||
connection.rollback()
|
||||
return
|
||||
connection.execute("INSERT INTO audio_files(task_id, owner_id, storage_key, mime_type, size_bytes, checksum) VALUES (%s, %s, %s, %s, %s, %s)", (task["id"], task["user_id"], storage_key, mime_type, len(audio), checksum))
|
||||
connection.execute("UPDATE tts_tasks SET status = 'succeeded', finished_at = now(), updated_at = now() WHERE id = %s AND status = 'running'", (task["id"],))
|
||||
connection.execute("UPDATE tts_tasks SET status = 'succeeded', finished_at = now(), updated_at = now(), lease_token = NULL, lease_expires_at = NULL WHERE id = %s AND status = 'running' AND lease_token = %s", (task["id"], task["lease_token"]))
|
||||
connection.execute("UPDATE quota_accounts SET reserved = GREATEST(0, reserved - %s), used = used + %s, version = version + 1 WHERE id = %s", (task["reserved_amount"], task["reserved_amount"], task["quota_account_id"]))
|
||||
connection.execute("INSERT INTO usage_records(user_id, quota_account_id, task_id, type, amount) VALUES (%s, %s, %s, 'consume', %s)", (task["user_id"], task["quota_account_id"], task["id"], task["reserved_amount"]))
|
||||
connection.commit()
|
||||
|
||||
|
||||
def recover_expired_tasks() -> None:
|
||||
with psycopg.connect(settings.database_url, row_factory=psycopg.rows.dict_row) as connection:
|
||||
tasks = connection.execute("SELECT * FROM tts_tasks WHERE status = 'running' AND lease_expires_at IS NOT NULL AND lease_expires_at < now() FOR UPDATE SKIP LOCKED").fetchall()
|
||||
for task in tasks:
|
||||
connection.execute("UPDATE tts_tasks SET status = 'failed', error_code = 'WORKER_LEASE_EXPIRED', finished_at = now(), updated_at = now(), lease_token = NULL, lease_expires_at = NULL WHERE id = %s AND status = 'running'", (task["id"],))
|
||||
connection.execute("UPDATE quota_accounts SET reserved = GREATEST(0, reserved - %s), version = version + 1 WHERE id = %s", (task["reserved_amount"], task["quota_account_id"]))
|
||||
connection.execute("INSERT INTO usage_records(user_id, quota_account_id, task_id, type, amount) VALUES (%s, %s, %s, 'release', %s)", (task["user_id"], task["quota_account_id"], task["id"], task["reserved_amount"]))
|
||||
connection.commit()
|
||||
|
||||
|
||||
def cleanup_orphan_audio() -> None:
|
||||
storage_dir = Path(settings.audio_storage_dir)
|
||||
if not storage_dir.exists():
|
||||
return
|
||||
with psycopg.connect(settings.database_url, row_factory=psycopg.rows.dict_row) as connection:
|
||||
known = {row["storage_key"] for row in connection.execute("SELECT storage_key FROM audio_files WHERE status = 'available'").fetchall()}
|
||||
now = time.time()
|
||||
for path in storage_dir.rglob("*"):
|
||||
if path.is_file() and path.relative_to(storage_dir).as_posix() not in known and now - path.stat().st_mtime > 3600:
|
||||
path.unlink(missing_ok=True)
|
||||
|
||||
|
||||
async def process_task(task: dict) -> None:
|
||||
if not settings.tts_upstream_url:
|
||||
await asyncio.to_thread(finish_task_failure, task, "UPSTREAM_NOT_CONFIGURED")
|
||||
@ -194,6 +231,7 @@ async def process_task(task: dict) -> None:
|
||||
|
||||
async def task_worker() -> None:
|
||||
while True:
|
||||
await asyncio.to_thread(recover_expired_tasks)
|
||||
task = await asyncio.to_thread(claim_next_task)
|
||||
if task:
|
||||
await process_task(task)
|
||||
@ -207,6 +245,7 @@ worker_task: asyncio.Task | None = None
|
||||
@app.on_event("startup")
|
||||
async def start_worker():
|
||||
global worker_task
|
||||
await asyncio.to_thread(cleanup_orphan_audio)
|
||||
worker_task = asyncio.create_task(task_worker())
|
||||
|
||||
|
||||
@ -351,7 +390,10 @@ def create_tts_task(payload: TtsTaskRequest, request: Request, connection: Conne
|
||||
raise error("IDEMPOTENCY_REQUIRED", "需要有效的 Idempotency-Key", 400)
|
||||
plan = effective_plan(connection, user)
|
||||
policy = connection.execute("SELECT * FROM plan_policies WHERE code = %s", (plan,)).fetchone()
|
||||
text_length = len(payload.text)
|
||||
text = normalize_tts_text(payload.text)
|
||||
if not text.strip():
|
||||
raise error("EMPTY_TEXT", "文本不能只有空白", 422)
|
||||
text_length = len(text)
|
||||
if not policy or text_length > policy["max_text_length"]:
|
||||
raise error("TEXT_TOO_LONG", "文本超过当前计划限制", 422)
|
||||
parameters = dict(payload.parameters)
|
||||
@ -362,12 +404,16 @@ def create_tts_task(payload: TtsTaskRequest, request: Request, connection: Conne
|
||||
voice = connection.execute("SELECT * FROM tts_voices WHERE provider_voice_id = %s AND enabled = true", (payload.voice_id,)).fetchone()
|
||||
if not voice or plan not in (voice["allowed_plans"] or ["free", "vip"]):
|
||||
raise error("VOICE_NOT_ALLOWED", "音色不可用", 422)
|
||||
request_hash = hashlib.sha256(json.dumps({"text": payload.text, "voice": payload.voice_id, "parameters": parameters}, ensure_ascii=False, sort_keys=True).encode()).hexdigest()
|
||||
request_hash = hashlib.sha256(json.dumps({"text": text, "voice": payload.voice_id, "parameters": parameters}, ensure_ascii=False, sort_keys=True).encode()).hexdigest()
|
||||
connection.execute("SELECT pg_advisory_xact_lock(hashtext(%s))", (str(user["id"]),))
|
||||
existing = connection.execute("SELECT * FROM tts_tasks WHERE user_id = %s AND idempotency_key = %s", (user["id"], idempotency_key)).fetchone()
|
||||
if existing:
|
||||
if existing["request_hash"] != request_hash:
|
||||
raise error("IDEMPOTENCY_CONFLICT", "幂等键已用于其他请求", 409)
|
||||
return task_view(connection, existing)
|
||||
active_count = connection.execute("SELECT count(*) FROM tts_tasks WHERE user_id = %s AND status IN ('queued', 'running')", (user["id"],)).fetchone()["count"]
|
||||
if active_count >= policy["max_concurrency"]:
|
||||
raise error("CONCURRENCY_LIMIT", "并发任务数已达上限", 409)
|
||||
quota = ensure_quota(connection, user["id"], plan)
|
||||
quota = connection.execute("SELECT * FROM quota_accounts WHERE id = %s FOR UPDATE", (quota["id"],)).fetchone()
|
||||
available = quota["limit_snapshot"] + quota["adjustment"] - quota["used"] - quota["reserved"]
|
||||
@ -378,7 +424,7 @@ def create_tts_task(payload: TtsTaskRequest, request: Request, connection: Conne
|
||||
INSERT INTO tts_tasks(user_id, text, text_length, voice_id, provider_voice_id, parameters, idempotency_key, request_hash, policy_version, quota_account_id, reserved_amount)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING *
|
||||
""",
|
||||
(user["id"], payload.text, text_length, voice["id"], payload.voice_id, Json(parameters), idempotency_key, request_hash, policy["version"], quota["id"], text_length),
|
||||
(user["id"], text, text_length, voice["id"], payload.voice_id, Json(parameters), idempotency_key, request_hash, policy["version"], quota["id"], text_length),
|
||||
).fetchone()
|
||||
connection.execute("UPDATE quota_accounts SET reserved = reserved + %s, version = version + 1 WHERE id = %s", (text_length, quota["id"]))
|
||||
connection.execute("INSERT INTO usage_records(user_id, quota_account_id, task_id, type, amount, idempotency_key) VALUES (%s, %s, %s, 'reserve', %s, %s)", (user["id"], quota["id"], task["id"], text_length, f"reserve:{task['id']}"))
|
||||
|
||||
3
services/api/migrations/003_tts_leases.sql
Normal file
3
services/api/migrations/003_tts_leases.sql
Normal file
@ -0,0 +1,3 @@
|
||||
ALTER TABLE tts_tasks ADD COLUMN IF NOT EXISTS lease_token TEXT;
|
||||
ALTER TABLE tts_tasks ADD COLUMN IF NOT EXISTS lease_expires_at TIMESTAMPTZ;
|
||||
CREATE INDEX IF NOT EXISTS tts_tasks_lease_idx ON tts_tasks(status, lease_expires_at);
|
||||
Loading…
Reference in New Issue
Block a user