fix: requestCancellation TOCTOU 竞态 — update 改为 updateMany CAS
- running 路径:prisma.aiJob.update() → updateMany() CAS - CAS 失败时重读当前状态:已取消则幂等返回,已终态则抛出 JobNotCancellableError - 新增 2 个 TOCTOU 场景测试(已变 cancelled → 幂等,已变 succeeded → 报错) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
parent
b08d27a449
commit
38d1c4f414
@ -369,22 +369,42 @@ describe('AiJobLifecycleRepository', () => {
|
|||||||
expect(result.status).toBe('cancelled');
|
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.findUnique as jest.Mock).mockResolvedValue({ lifecycleStatus: 'running' });
|
||||||
(prisma.aiJob.update as jest.Mock).mockResolvedValue({
|
// updateMany (CAS) succeeds
|
||||||
id: 'job-001',
|
(prisma.aiJob.updateMany as jest.Mock).mockResolvedValue({ count: 1 });
|
||||||
lifecycleStatus: 'running',
|
|
||||||
cancelRequestedAt: new Date(),
|
|
||||||
});
|
|
||||||
|
|
||||||
const result = await repo.requestCancellation('job-001');
|
const result = await repo.requestCancellation('job-001');
|
||||||
expect(result.status).toBe('cancel_requested');
|
expect(result.status).toBe('cancel_requested');
|
||||||
expect(prisma.aiJob.update).toHaveBeenCalledWith(
|
expect(prisma.aiJob.updateMany).toHaveBeenCalledWith(
|
||||||
expect.objectContaining({
|
expect.objectContaining({
|
||||||
where: { id: 'job-001', lifecycleStatus: 'running' },
|
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', () => {
|
describe('findByLifecycleStatus', () => {
|
||||||
|
|||||||
@ -375,17 +375,38 @@ export class AiJobLifecycleRepository {
|
|||||||
const status = this.stateMachine.parse(job.lifecycleStatus);
|
const status = this.stateMachine.parse(job.lifecycleStatus);
|
||||||
|
|
||||||
if (status === 'queued') {
|
if (status === 'queued') {
|
||||||
// 直接取消
|
// 直接取消(markCancelled 内部已有 CAS + 错误分类)
|
||||||
await this.markCancelled(jobId);
|
await this.markCancelled(jobId);
|
||||||
return { status: 'cancelled' };
|
return { status: 'cancelled' };
|
||||||
}
|
}
|
||||||
|
|
||||||
if (status === 'running') {
|
if (status === 'running') {
|
||||||
// 设置取消信号
|
// CAS 写入 cancelRequestedAt — 只有 lifecycleStatus 仍为 running 时才成功。
|
||||||
await this.prisma.aiJob.update({
|
// 使用 updateMany 避免 read→write 之间的 TOCTOU 竞态窗口。
|
||||||
|
const result = await this.prisma.aiJob.updateMany({
|
||||||
where: { id: jobId, lifecycleStatus: 'running' },
|
where: { id: jobId, lifecycleStatus: 'running' },
|
||||||
data: { cancelRequestedAt: new Date() },
|
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' };
|
return { status: 'cancel_requested' };
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user