fix(kv): claim parked jobs by member instead of secondary key (#2232)

fluxer_api/src/api/infrastructure/KVScheduledJobQueueService.ts

@@ -13,6 +13,11 @@ export interface ParkedScheduledJob {
 	ledgerJobId: string | null;
 }
 
+export interface ReadyScheduledJob {
+	job: ParkedScheduledJob;
+	member: string;
+}
+
 const QUEUE_KEY = 'scheduled_job_queue';
 const SECONDARY_KEY_PREFIX = 'scheduled_job_queue:';
 
@@ -75,22 +80,22 @@ export class KVScheduledJobQueueService {
 		}
 	}
 
-	async claimJob(jobIdentity: string): Promise<boolean> {
+	async claimJob(jobIdentity: string, member = ''): Promise<boolean> {
 		try {
-			return await this.kvClient.removeBulkDeletion(QUEUE_KEY, this.getSecondaryKey(jobIdentity));
+			return await this.kvClient.removeBulkDeletion(QUEUE_KEY, this.getSecondaryKey(jobIdentity), member);
 		} catch (error) {
 			Logger.error({error, jobIdentity}, 'Failed to claim parked scheduled job');
 			throw error;
 		}
 	}
 
-	async getReadyJobs(nowMs: number, limit: number): Promise<Array<ParkedScheduledJob>> {
+	async getReadyJobs(nowMs: number, limit: number): Promise<Array<ReadyScheduledJob>> {
 		try {
 			const results = await this.kvClient.zrangebyscore(QUEUE_KEY, '-inf', nowMs, 'LIMIT', 0, limit);
-			const jobs: Array<ParkedScheduledJob> = [];
+			const jobs: Array<ReadyScheduledJob> = [];
 			for (const result of results) {
 				try {
-					jobs.push(this.deserializeQueueItem(result));
+					jobs.push({job: this.deserializeQueueItem(result), member: result});
 				} catch (error) {
 					Logger.error({error, result}, 'Failed to parse parked scheduled job entry');
 				}

fluxer_api/src/api/test/mocks/MockKVProvider.ts

@@ -667,9 +667,13 @@ export class MockKVProvider implements IKVProvider {
 		this.expiries.delete(secondaryKey);
 	}
 
-	async removeBulkDeletion(queueKey: string, secondaryKey: string): Promise<boolean> {
-		this.removeBulkDeletionSpy(queueKey, secondaryKey);
+	async removeBulkDeletion(queueKey: string, secondaryKey: string, member = ''): Promise<boolean> {
+		this.removeBulkDeletionSpy(queueKey, secondaryKey, member);
 		this.evictIfExpired(secondaryKey);
+		if (member !== '' && (await this.zrem(queueKey, member)) === 1) {
+			this.deleteKey(secondaryKey);
+			return true;
+		}
 		const value = this.stringStore.get(secondaryKey);
 		if (value === undefined) {
 			return false;

fluxer_api/src/api/worker/tasks/ProcessScheduledJobQueue.ts

@@ -28,9 +28,9 @@ const processScheduledJobQueue: WorkerTaskHandler = async (_payload, helpers) =>
 	let claimedElsewhereCount = 0;
 	let droppedCount = 0;
 	let failedCount = 0;
-	for (const job of readyJobs) {
+	for (const {job, member} of readyJobs) {
 		try {
-			const claimed = await scheduledJobQueueService.claimJob(job.jobIdentity);
+			const claimed = await scheduledJobQueueService.claimJob(job.jobIdentity, member);
 			if (!claimed) {
 				claimedElsewhereCount += 1;
 				continue;

fluxer_api/src/api/worker/tests/ScheduledJobQueue.test.ts

@@ -167,13 +167,14 @@ describe('Scheduled job due queue', () => {
 		expect(await scheduledJobQueue.getReadyJobs(Date.now(), 10)).toHaveLength(0);
 		const due = await scheduledJobQueue.getReadyJobs(runAt.getTime() - LANE_ACK_WAIT_MS, 10);
 		expect(due).toHaveLength(1);
-		expect(due[0]).toEqual({
+		expect(due[0]!.job).toEqual({
 			jobIdentity: '4242',
 			taskType: TASK_TYPE,
 			payload: {scheduledMessageId: '77'},
 			runAtMs: runAt.getTime(),
 			ledgerJobId: '4242',
 		});
+		expect(due[0]!.member).toContain('"jobIdentity":"4242"');
 	});
 
 	it('releases a parked job exactly once when its due time arrives', async () => {
@@ -277,3 +278,66 @@ describe('Scheduled job due queue', () => {
 		expect(await scheduledJobQueue.getQueueSize()).toBe(0);
 	});
 });
+
+describe('Parked job claim survives secondary key loss', () => {
+	it('claims by member when the secondary key was evicted', async () => {
+		const kv = new MockKVProvider();
+		const service = new KVScheduledJobQueueService(kv);
+		const runAtMs = Date.now() + 86_400_000;
+		const job = {
+			jobIdentity: 'job-evicted',
+			taskType: TASK_TYPE,
+			payload: {message_id: '1'} as WorkerJobPayload,
+			runAtMs,
+			ledgerJobId: null,
+		};
+
+		await service.parkJob(job, new Date(Date.now() - 1000));
+		expect(await service.getQueueSize()).toBe(1);
+
+		await kv.del('scheduled_job_queue:job-evicted');
+
+		const ready = await service.getReadyJobs(Date.now(), 10);
+		expect(ready).toHaveLength(1);
+
+		expect(await service.claimJob(ready[0]!.job.jobIdentity, ready[0]!.member)).toBe(true);
+		expect(await service.getQueueSize()).toBe(0);
+	});
+
+	it('leaves the entry wedged when claiming without the member', async () => {
+		const kv = new MockKVProvider();
+		const service = new KVScheduledJobQueueService(kv);
+		const job = {
+			jobIdentity: 'job-wedged',
+			taskType: TASK_TYPE,
+			payload: {message_id: '2'} as WorkerJobPayload,
+			runAtMs: Date.now() + 86_400_000,
+			ledgerJobId: null,
+		};
+
+		await service.parkJob(job, new Date(Date.now() - 1000));
+		await kv.del('scheduled_job_queue:job-wedged');
+
+		expect(await service.claimJob('job-wedged')).toBe(false);
+		expect(await service.getQueueSize()).toBe(1);
+	});
+
+	it('still reports a second claim of the same member as already claimed', async () => {
+		const kv = new MockKVProvider();
+		const service = new KVScheduledJobQueueService(kv);
+		const job = {
+			jobIdentity: 'job-twice',
+			taskType: TASK_TYPE,
+			payload: {message_id: '3'} as WorkerJobPayload,
+			runAtMs: Date.now() + 86_400_000,
+			ledgerJobId: null,
+		};
+
+		await service.parkJob(job, new Date(Date.now() - 1000));
+		const ready = await service.getReadyJobs(Date.now(), 10);
+		const member = ready[0]!.member;
+
+		expect(await service.claimJob('job-twice', member)).toBe(true);
+		expect(await service.claimJob('job-twice', member)).toBe(false);
+	});
+});