import { randomUUID } from "node:crypto"; import { Prisma, type TranscriptionTask } from "@prisma/client"; import { prisma } from "../prisma"; import { getAsrConfig } from "./dashscope"; import { isActive, transcriptionState } from "./types"; function configurationError(): string | null { try { getAsrConfig(); return null; } catch (error) { return (error as Error).message; } } export async function enqueueTranscription(videoId: string, force = false) { return prisma.$transaction(async (tx) => { const videos = await tx.$queryRaw< { aweme_id: string }[] >`SELECT aweme_id FROM "Video" WHERE aweme_id = ${videoId} FOR UPDATE`; if (!videos.length) return null; const [transcript, task] = await Promise.all([ tx.videoTranscript.findUnique({ where: { videoId } }), tx.transcriptionTask.findUnique({ where: { videoId } }), ]); if ((task && isActive(task.status)) || (!force && transcript)) return transcriptionState(transcript, task); const lastError = configurationError(); const data = { status: lastError ? "FAILED" : "QUEUED", lastError, providerTaskId: null, submittingAt: null, submittedAt: null, leaseUntil: null, leaseToken: null, attempts: 0, nextPollAt: new Date(), }; const saved = await tx.transcriptionTask.upsert({ where: { videoId }, create: { videoId, ...data }, update: data, }); return transcriptionState(transcript, saved); }); } export async function enqueueMissingTranscriptions() { const error = configurationError(); if (error) throw new Error(error); // One durable, atomic batch. Completed historical transcripts are never requeued. return prisma.$executeRaw` INSERT INTO "TranscriptionTask" ("videoId", status, "updatedAt") SELECT v.aweme_id, 'QUEUED', (CURRENT_TIMESTAMP AT TIME ZONE 'UTC') FROM "Video" v WHERE NOT EXISTS (SELECT 1 FROM "VideoTranscript" t WHERE t."videoId" = v.aweme_id) ON CONFLICT ("videoId") DO UPDATE SET status = 'QUEUED', "providerTaskId" = NULL, "submittingAt" = NULL, "submittedAt" = NULL, "leaseUntil" = NULL, "leaseToken" = NULL, attempts = 0, "lastError" = NULL, "nextPollAt" = (CURRENT_TIMESTAMP AT TIME ZONE 'UTC'), "updatedAt" = (CURRENT_TIMESTAMP AT TIME ZONE 'UTC') WHERE "TranscriptionTask".status = 'FAILED' `; } export async function readTranscription(videoId: string) { const video = await prisma.video.findUnique({ where: { aweme_id: videoId }, include: { transcript: true, transcriptionTask: true }, }); return video ? transcriptionState(video.transcript, video.transcriptionTask) : null; } export async function claimTask(): Promise { return prisma.$transaction(async (tx) => { // Serialize admission across all worker processes, including PM2 restarts. await tx.$queryRaw`SELECT pg_advisory_xact_lock(74196321)::text`; const running = await tx.transcriptionTask.count({ where: { status: "PROCESSING" }, }); const eligible = running < 2 ? Prisma.sql`status IN ('QUEUED', 'PROCESSING')` : Prisma.sql`status = 'PROCESSING'`; const rows = await tx.$queryRaw` SELECT * FROM "TranscriptionTask" WHERE ${eligible} AND "nextPollAt" <= (CURRENT_TIMESTAMP AT TIME ZONE 'UTC') AND ("leaseUntil" IS NULL OR "leaseUntil" < (CURRENT_TIMESTAMP AT TIME ZONE 'UTC')) ORDER BY CASE WHEN status = 'PROCESSING' THEN 0 ELSE 1 END, "nextPollAt", "createdAt" LIMIT 1 FOR UPDATE SKIP LOCKED `; if (!rows.length) return null; return tx.transcriptionTask.update({ where: { videoId: rows[0].videoId }, data: { status: "PROCESSING", leaseToken: randomUUID(), leaseUntil: new Date(Date.now() + 180_000), }, }); }); } export async function updateOwnedTask( task: TranscriptionTask, data: Prisma.TranscriptionTaskUpdateManyMutationInput, ) { const result = await prisma.transcriptionTask.updateMany({ where: { videoId: task.videoId, leaseToken: task.leaseToken, status: "PROCESSING", leaseUntil: { gt: new Date() }, }, data, }); if (!result.count) throw new Error("转写任务租约已失效"); } export async function completeTask( task: TranscriptionTask, transcript: string[], ) { await prisma.$transaction(async (tx) => { const result = await tx.transcriptionTask.updateMany({ where: { videoId: task.videoId, leaseToken: task.leaseToken, status: "PROCESSING", leaseUntil: { gt: new Date() }, }, data: { status: "SUCCEEDED", lastError: null, leaseToken: null, leaseUntil: null, attempts: 0, }, }); if (!result.count) throw new Error("转写任务租约已失效"); await tx.videoTranscript.upsert({ where: { videoId: task.videoId }, create: { videoId: task.videoId, transcript }, update: { transcript }, }); }); }