import { Injectable } from '@nestjs/common'; import { PrismaService } from '../../infrastructure/database/prisma.service'; import { Prisma } from '@prisma/client'; import { JobLockConflictError, JobNotActiveError, JobAlreadyTerminalError, JobNotCancellableError, JobExtendLockError, } from './ai-job.errors'; import { AiJobStateMachine, LifecycleStatus, LIFECYCLE_TO_STATUS, } from './ai-job-state-machine'; /** * ── Job 创建输入 ── */ export interface CreateJobInput { userId: string; jobType: string; triggerType?: string; queueName?: string; targetType?: string; targetId?: string; parentJobId?: string; idempotencyKey?: string; priority?: number; credentialMode?: string; credentialId?: string; modelProvider?: string; modelName?: string; promptKey?: string; promptVersion?: string; outputSchemaVersion?: string; inputSchemaVersion?: string; inputRef?: string; maxAttempts?: number; timeoutMs?: number; } /** * ── Job 错误信息 ── */ export interface JobErrorInfo { errorCode: string; publicErrorMessage?: string; internalErrorMessage?: string; } /** * M-AI-03-03: AiJobLifecycleRepository * * 统一 lifecycleStatus 的唯一写入入口。 * * 约束: * - 所有状态变更使用 CAS(条件 UPDATE WHERE lifecycleStatus IN (...)) * - 禁止先读后无条件 update * - 可选 tx 参数支持在外部 Prisma Transaction 中复用事务 * - lockJob 不使用 tx(CAS 必须立即对其他 Worker 可见) */ @Injectable() export class AiJobLifecycleRepository { private readonly LOCK_TTL_MS = 60_000; // ADR-003:60s constructor( private readonly prisma: PrismaService, private readonly stateMachine: AiJobStateMachine, ) {} // ═══════════════════════════════════════════════════════════ // 创建 // ═══════════════════════════════════════════════════════════ /** * 在给定事务中创建 Job(lifecycleStatus = 'queued')。 * 调用方负责在同一 tx 中写入 OutboxEvent。 */ async createJob( tx: Prisma.TransactionClient, input: CreateJobInput, ) { const job = await tx.aiJob.create({ data: { userId: input.userId, jobType: input.jobType, triggerType: input.triggerType ?? 'user_api', queueName: input.queueName ?? 'ai-interactive', targetType: input.targetType ?? null, targetId: input.targetId ?? null, parentJobId: input.parentJobId ?? null, idempotencyKey: input.idempotencyKey ?? null, priority: input.priority ?? 0, credentialMode: input.credentialMode ?? 'platform_key', credentialId: input.credentialId ?? null, // modelProvider / modelName 不在此层设置默认值。 // 调用方(AiJobService)负责从 JobDefinition 解析后传入。 // 若未传入,依赖 Prisma schema @default 保证列非空。 ...(input.modelProvider ? { modelProvider: input.modelProvider } : {}), ...(input.modelName ? { modelName: input.modelName } : {}), promptKey: input.promptKey ?? null, promptVersion: input.promptVersion ?? null, outputSchemaVersion: input.outputSchemaVersion ?? null, inputSchemaVersion: input.inputSchemaVersion ?? null, inputRef: input.inputRef ?? null, maxAttempts: input.maxAttempts ?? 3, timeoutMs: input.timeoutMs ?? 120000, // ── 生命周期 ── lifecycleStatus: 'queued', status: LIFECYCLE_TO_STATUS['queued'], // Shadow Write: pending queuedAt: new Date(), attemptCount: 0, progress: 0, }, }); return job; } // ═══════════════════════════════════════════════════════════ // 抢锁(CAS: queued → running) // ═══════════════════════════════════════════════════════════ /** * CAS 抢锁。只有 lifecycleStatus = 'queued' 的 Job 才能被锁定。 * * 必须是独立操作(不传 tx)—— CAS 需要立即对其他 Worker 可见。 */ async lockJob(jobId: string, instanceId: string) { const now = new Date(); const lockUntil = new Date(now.getTime() + this.LOCK_TTL_MS); // CAS: queued → running,仅当 status 仍为 queued 时成功 const result = await this.prisma.aiJob.updateMany({ where: { id: jobId, lifecycleStatus: 'queued', }, data: { lifecycleStatus: 'running', status: LIFECYCLE_TO_STATUS['running'], // Shadow Write: processing lockedBy: instanceId, lockedAt: now, lockUntil, startedAt: now, attemptCount: { increment: 1 }, // 原子递增,每次 lockJob 真实执行次数 +1 }, }); if (result.count === 0) { // 抢锁失败 —— 检查原因以提供更精确的错误 const current = await this.prisma.aiJob.findUnique({ where: { id: jobId }, select: { lifecycleStatus: true }, }); if (!current) { throw new JobLockConflictError(jobId); // Job 不存在 } if (this.stateMachine.isTerminal(current.lifecycleStatus as LifecycleStatus)) { throw new JobAlreadyTerminalError(jobId, current.lifecycleStatus!); } if (current.lifecycleStatus === 'running') { throw new JobLockConflictError(jobId); // 已被其他 Worker 抢走 } throw new JobLockConflictError(jobId); } const job = await this.prisma.aiJob.findUniqueOrThrow({ where: { id: jobId }, }); return job; } // ═══════════════════════════════════════════════════════════ // 续锁(Heartbeat) // ═══════════════════════════════════════════════════════════ /** * 续约锁定期。仅当 Job 处于 'running' 状态且 lockedBy 匹配时成功。 */ async extendLock(jobId: string, instanceId: string) { const now = new Date(); const lockUntil = new Date(now.getTime() + this.LOCK_TTL_MS); const result = await this.prisma.aiJob.updateMany({ where: { id: jobId, lifecycleStatus: 'running', lockedBy: instanceId, }, data: { lockUntil }, }); if (result.count === 0) { const current = await this.prisma.aiJob.findUnique({ where: { id: jobId }, select: { lifecycleStatus: true, lockedBy: true }, }); if (!current || current.lifecycleStatus !== 'running') { throw new JobNotActiveError(jobId, current?.lifecycleStatus ?? 'unknown'); } throw new JobExtendLockError(jobId, instanceId); } // 读取当前状态以检查 cancelRequestedAt const job = await this.prisma.aiJob.findUniqueOrThrow({ where: { id: jobId }, }); return { jobId, lockUntil: lockUntil.getTime(), cancelRequested: job.cancelRequestedAt != null, }; } // ═══════════════════════════════════════════════════════════ // 标记成功 // ═══════════════════════════════════════════════════════════ /** * 标记 Job 成功。 * * @param tx 可选事务客户端。Projector 路径传入以共享事务;非 Projector 路径不传。 */ async markSucceeded(jobId: string, tx?: Prisma.TransactionClient): Promise { const client = tx ?? this.prisma; const now = new Date(); const result = await client.aiJob.updateMany({ where: { id: jobId, lifecycleStatus: 'running', }, data: { lifecycleStatus: 'succeeded', status: LIFECYCLE_TO_STATUS['succeeded'], // Shadow Write: completed finishedAt: now, // 清除运行时字段 lockedBy: null, lockedAt: null, lockUntil: null, errorCode: null, publicErrorMessage: null, internalErrorMessage: null, }, }); if (result.count === 0) { const current = await client.aiJob.findUnique({ where: { id: jobId }, select: { lifecycleStatus: true }, }); if (current && this.stateMachine.isTerminal(current.lifecycleStatus as LifecycleStatus)) { throw new JobAlreadyTerminalError(jobId, current.lifecycleStatus!); } throw new JobNotActiveError(jobId, current?.lifecycleStatus ?? 'unknown'); } return client.aiJob.findUniqueOrThrow({ where: { id: jobId } }); } // ═══════════════════════════════════════════════════════════ // 标记失败 // ═══════════════════════════════════════════════════════════ /** * 标记 Job 失败(终态)。 * * @param tx 可选事务客户端。Engine 非 Projector 错误路径不传。 */ async markFailed(jobId: string, error: JobErrorInfo, tx?: Prisma.TransactionClient): Promise { const client = tx ?? this.prisma; const now = new Date(); const result = await client.aiJob.updateMany({ where: { id: jobId, lifecycleStatus: 'running', }, data: { lifecycleStatus: 'failed', status: LIFECYCLE_TO_STATUS['failed'], // Shadow Write: failed finishedAt: now, lockedBy: null, lockedAt: null, lockUntil: null, errorCode: error.errorCode, publicErrorMessage: error.publicErrorMessage ?? null, internalErrorMessage: error.internalErrorMessage ?? null, }, }); if (result.count === 0) { const current = await client.aiJob.findUnique({ where: { id: jobId }, select: { lifecycleStatus: true }, }); if (current && this.stateMachine.isTerminal(current.lifecycleStatus as LifecycleStatus)) { throw new JobAlreadyTerminalError(jobId, current.lifecycleStatus!); } throw new JobNotActiveError(jobId, current?.lifecycleStatus ?? 'unknown'); } return client.aiJob.findUniqueOrThrow({ where: { id: jobId } }); } // ═══════════════════════════════════════════════════════════ // 标记取消 // ═══════════════════════════════════════════════════════════ /** * 标记 Job 取消。 * * queued 状态可直接取消;running 状态需先设置 cancelRequestedAt, * Engine/Heartbeat 检测后调用此方法。 * * @param tx 可选事务客户端。取消路径通常不传(独立操作)。 */ async markCancelled(jobId: string, tx?: Prisma.TransactionClient): Promise { const client = tx ?? this.prisma; const now = new Date(); // 尝试从 queued 或 running 取消 const result = await client.aiJob.updateMany({ where: { id: jobId, lifecycleStatus: { in: ['queued', 'running'] }, }, data: { lifecycleStatus: 'cancelled', // ── Shadow Write:cancelled 在旧 status 中不存在,映射为 failed + errorCode=CANCELLED ── status: LIFECYCLE_TO_STATUS['cancelled'], // 'failed' finishedAt: now, lockedBy: null, lockedAt: null, lockUntil: null, errorCode: 'CANCELLED', publicErrorMessage: 'Job was cancelled', internalErrorMessage: null, }, }); if (result.count === 0) { const current = await client.aiJob.findUnique({ where: { id: jobId }, select: { lifecycleStatus: true }, }); if (!current) { throw new JobLockConflictError(jobId); // Job 不存在,复用此错误 } if (this.stateMachine.isTerminal(current.lifecycleStatus as LifecycleStatus)) { throw new JobNotCancellableError(jobId, current.lifecycleStatus!); } throw new JobNotCancellableError(jobId, current.lifecycleStatus ?? 'unknown'); } return client.aiJob.findUniqueOrThrow({ where: { id: jobId } }); } // ═══════════════════════════════════════════════════════════ // 设置取消请求(不改变 lifecycleStatus) // ═══════════════════════════════════════════════════════════ /** * 对 running 状态的 Job 设置 cancelRequestedAt 信号。 * 不改变 lifecycleStatus — Engine/Heartbeat 检测后调用 markCancelled。 * * 对 queued 状态的 Job 直接取消。 */ async requestCancellation(jobId: string): Promise<{ status: 'cancelled' | 'cancel_requested' }> { const job = await this.prisma.aiJob.findUnique({ where: { id: jobId }, select: { lifecycleStatus: true }, }); if (!job) throw new JobLockConflictError(jobId); const status = this.stateMachine.parse(job.lifecycleStatus); if (status === 'queued') { // 直接取消(markCancelled 内部已有 CAS + 错误分类) await this.markCancelled(jobId); return { status: 'cancelled' }; } if (status === 'running') { // CAS 写入 cancelRequestedAt — 只有 lifecycleStatus 仍为 running 时才成功。 // 使用 updateMany 避免 read→write 之间的 TOCTOU 竞态窗口。 const result = await this.prisma.aiJob.updateMany({ where: { id: jobId, lifecycleStatus: 'running' }, data: { cancelRequestedAt: new Date() }, }); if (result.count === 0) { // 在 findUnique 和 updateMany 之间 Job 已进入终态或被其他 Worker 修改。 // 重新读取以返回精确错误。 const current = await this.prisma.aiJob.findUnique({ where: { id: jobId }, select: { lifecycleStatus: true }, }); const currentStatus = current?.lifecycleStatus ?? 'unknown'; if (currentStatus === 'cancelled') { return { status: 'cancelled' }; // 已被其他路径取消,幂等返回 } if (this.stateMachine.isTerminal( this.stateMachine.parse(currentStatus), )) { throw new JobNotCancellableError(jobId, currentStatus); } throw new JobNotCancellableError(jobId, currentStatus); } return { status: 'cancel_requested' }; } // 终态不可取消 throw new JobNotCancellableError(jobId, status); } // ═══════════════════════════════════════════════════════════ // 只读查询 // ═══════════════════════════════════════════════════════════ async findByLifecycleStatus( statuses: LifecycleStatus[], opts?: { queueName?: string; limit?: number }, ) { const where: Prisma.AiJobWhereInput = { lifecycleStatus: { in: statuses }, }; if (opts?.queueName) where.queueName = opts.queueName; return this.prisma.aiJob.findMany({ where, orderBy: [{ priority: 'asc' }, { createdAt: 'asc' }], take: opts?.limit ?? 50, }); } async findExpiredLocks(thresholdMs: number = 90_000) { const cutoff = new Date(Date.now() - thresholdMs); return this.prisma.aiJob.findMany({ where: { lifecycleStatus: 'running', lockUntil: { lt: cutoff }, }, orderBy: { createdAt: 'asc' }, }); } async findById(jobId: string) { return this.prisma.aiJob.findUnique({ where: { id: jobId } }); } }