diff --git a/src/modules/ai-job/ai-job-lifecycle.repository.spec.ts b/src/modules/ai-job/ai-job-lifecycle.repository.spec.ts index 6473c63..f275d07 100644 --- a/src/modules/ai-job/ai-job-lifecycle.repository.spec.ts +++ b/src/modules/ai-job/ai-job-lifecycle.repository.spec.ts @@ -369,22 +369,42 @@ describe('AiJobLifecycleRepository', () => { expect(result.status).toBe('cancelled'); }); - it('running Job 设置 cancelRequestedAt 信号', async () => { + it('running Job CAS 写入 cancelRequestedAt 信号', async () => { (prisma.aiJob.findUnique as jest.Mock).mockResolvedValue({ lifecycleStatus: 'running' }); - (prisma.aiJob.update as jest.Mock).mockResolvedValue({ - id: 'job-001', - lifecycleStatus: 'running', - cancelRequestedAt: new Date(), - }); + // updateMany (CAS) succeeds + (prisma.aiJob.updateMany as jest.Mock).mockResolvedValue({ count: 1 }); const result = await repo.requestCancellation('job-001'); expect(result.status).toBe('cancel_requested'); - expect(prisma.aiJob.update).toHaveBeenCalledWith( + expect(prisma.aiJob.updateMany).toHaveBeenCalledWith( expect.objectContaining({ where: { id: 'job-001', lifecycleStatus: 'running' }, + data: { cancelRequestedAt: expect.any(Date) }, }), ); }); + + it('running Job TOCTOU:CAS 失败且已变为 cancelled → 幂等返回', async () => { + (prisma.aiJob.findUnique as jest.Mock).mockResolvedValue({ lifecycleStatus: 'running' }); + // CAS fails (job already cancelled between read and write) + (prisma.aiJob.updateMany as jest.Mock).mockResolvedValue({ count: 0 }); + // Re-read: already cancelled + (prisma.aiJob.findUnique as jest.Mock).mockResolvedValueOnce({ lifecycleStatus: 'running' }); + // The second findUnique call (after CAS failure) + (prisma.aiJob.findUnique as jest.Mock).mockResolvedValueOnce({ lifecycleStatus: 'cancelled' }); + + const result = await repo.requestCancellation('job-001'); + expect(result.status).toBe('cancelled'); // 幂等返回 + }); + + it('running Job TOCTOU:CAS 失败且已变为 succeeded → 抛出错误', async () => { + (prisma.aiJob.findUnique as jest.Mock).mockResolvedValue({ lifecycleStatus: 'running' }); + (prisma.aiJob.updateMany as jest.Mock).mockResolvedValue({ count: 0 }); + // Re-read: already succeeded (terminal) + (prisma.aiJob.findUnique as jest.Mock).mockResolvedValueOnce({ lifecycleStatus: 'succeeded' }); + + await expect(repo.requestCancellation('job-001')).rejects.toThrow(JobNotCancellableError); + }); }); describe('findByLifecycleStatus', () => { diff --git a/src/modules/ai-job/ai-job-lifecycle.repository.ts b/src/modules/ai-job/ai-job-lifecycle.repository.ts index 15fcaf8..97687ce 100644 --- a/src/modules/ai-job/ai-job-lifecycle.repository.ts +++ b/src/modules/ai-job/ai-job-lifecycle.repository.ts @@ -375,17 +375,38 @@ export class AiJobLifecycleRepository { const status = this.stateMachine.parse(job.lifecycleStatus); if (status === 'queued') { - // 直接取消 + // 直接取消(markCancelled 内部已有 CAS + 错误分类) await this.markCancelled(jobId); return { status: 'cancelled' }; } if (status === 'running') { - // 设置取消信号 - await this.prisma.aiJob.update({ + // 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' }; }