fix(worker): rebuild the deletion queue and cancel system dms (#2510)

fluxer_api/src/api/infrastructure/KVBulkMessageDeletionQueueService.test.ts

@@ -0,0 +1,193 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+import type {IKVProvider} from '@pkgs/kv_client/src/IKVProvider';
+import type {WorkerTaskHelpers} from '@pkgs/worker/src/contracts/WorkerTask';
+import {afterEach, describe, expect, test} from 'vitest';
+import {createUserID, type UserID} from '../BrandedTypes';
+import {EMPTY_USER_ROW} from '../database/types/UserTypes';
+import {User} from '../models/User';
+import {NoopLogger} from '../test/mocks/NoopLogger';
+import type {UserRepository} from '../user/repositories/UserRepository';
+import processPendingBulkMessageDeletions from '../worker/tasks/ProcessPendingBulkMessageDeletions';
+import {clearWorkerDependencies, setWorkerDependenciesForTest} from '../worker/WorkerContext';
+import type {WorkerDependencies} from '../worker/WorkerDependencies';
+import {KVBulkMessageDeletionQueueService} from './KVBulkMessageDeletionQueueService';
+
+class FakeKVProvider {
+	readonly strings = new Map<string, string>();
+	readonly sortedSets = new Map<string, Map<string, number>>();
+
+	async get(key: string): Promise<string | null> {
+		return this.strings.get(key) ?? null;
+	}
+
+	async set(key: string, value: string): Promise<string | null> {
+		this.strings.set(key, value);
+		return 'OK';
+	}
+
+	async del(...keys: Array<string>): Promise<number> {
+		let removed = 0;
+		for (const key of keys) {
+			if (this.strings.delete(key)) {
+				removed++;
+			}
+			if (this.sortedSets.delete(key)) {
+				removed++;
+			}
+		}
+		return removed;
+	}
+
+	async exists(key: string): Promise<number> {
+		return this.strings.has(key) || this.sortedSets.has(key) ? 1 : 0;
+	}
+
+	async zcard(key: string): Promise<number> {
+		return this.sortedSets.get(key)?.size ?? 0;
+	}
+
+	async zrangebyscore(
+		key: string,
+		_min: string | number,
+		max: string | number,
+		_limitToken?: string,
+		offset?: number,
+		count?: number,
+	): Promise<Array<string>> {
+		const members = [...(this.sortedSets.get(key) ?? new Map<string, number>()).entries()]
+			.filter(([, score]) => score <= Number(max))
+			.sort((a, b) => a[1] - b[1])
+			.map(([member]) => member);
+		return members.slice(offset ?? 0, (offset ?? 0) + (count ?? members.length));
+	}
+
+	async scheduleBulkDeletion(queueKey: string, secondaryKey: string, score: number, value: string): Promise<void> {
+		let set = this.sortedSets.get(queueKey);
+		if (!set) {
+			set = new Map<string, number>();
+			this.sortedSets.set(queueKey, set);
+		}
+		set.set(value, score);
+		this.strings.set(secondaryKey, value);
+	}
+
+	async removeBulkDeletion(queueKey: string, secondaryKey: string): Promise<boolean> {
+		const value = this.strings.get(secondaryKey);
+		if (!value) {
+			return false;
+		}
+		this.sortedSets.get(queueKey)?.delete(value);
+		this.strings.delete(secondaryKey);
+		return true;
+	}
+
+	async acquireLock(key: string, token: string, _ttlSeconds: number): Promise<boolean> {
+		if (this.strings.has(key)) {
+			return false;
+		}
+		this.strings.set(key, token);
+		return true;
+	}
+
+	async releaseLock(key: string, token: string): Promise<boolean> {
+		if (this.strings.get(key) !== token) {
+			return false;
+		}
+		this.strings.delete(key);
+		return true;
+	}
+}
+
+function createUser(userId: UserID, pendingBulkMessageDeletionAt: Date | null): User {
+	return new User({
+		...EMPTY_USER_ROW,
+		user_id: userId,
+		username: `user${userId}`,
+		pending_bulk_message_deletion_at: pendingBulkMessageDeletionAt,
+	});
+}
+
+function createHelpers(): WorkerTaskHelpers {
+	return {
+		logger: new NoopLogger(),
+		jobId: 1n,
+		addJob: async () => 0n,
+		reportProgress: async () => {},
+		shouldCancel: async () => false,
+		setContextLink: async () => {},
+	};
+}
+
+function createHarness(users: Array<User>) {
+	const kvClient = new FakeKVProvider();
+	const scanCalls: Array<number> = [];
+	const userRepository = {
+		async scanAllUsersPage(limit: number): Promise<{users: Array<User>; pageState: null}> {
+			scanCalls.push(limit);
+			return {users, pageState: null};
+		},
+		async findUnique(userId: UserID): Promise<User | null> {
+			return users.find((user) => user.id === userId) ?? null;
+		},
+	} as unknown as UserRepository;
+	const jobs: Array<{name: string; payload: unknown}> = [];
+	const workerService = {
+		async addJob(name: string, payload: unknown): Promise<bigint> {
+			jobs.push({name, payload});
+			return 0n;
+		},
+	} as unknown as WorkerDependencies['workerService'];
+	const bulkMessageDeletionQueueService = new KVBulkMessageDeletionQueueService(
+		kvClient as unknown as IKVProvider,
+		userRepository,
+	);
+	setWorkerDependenciesForTest({bulkMessageDeletionQueueService, userRepository, workerService});
+	return {kvClient, scanCalls, jobs};
+}
+
+describe('processPendingBulkMessageDeletions', () => {
+	afterEach(() => {
+		clearWorkerDependencies();
+	});
+
+	test('rebuilds the queue from the users table when the state version is missing', async () => {
+		const userId = createUserID(1234n);
+		const scheduledAt = new Date(Date.now() - 60_000);
+		const harness = createHarness([createUser(userId, scheduledAt)]);
+
+		await processPendingBulkMessageDeletions({}, createHelpers());
+
+		expect(harness.scanCalls).toEqual([1000]);
+		expect(harness.jobs).toEqual([
+			{
+				name: 'bulkDeleteUserMessages',
+				payload: {userId: userId.toString(), scheduledAt: scheduledAt.getTime()},
+			},
+		]);
+		expect(harness.kvClient.strings.has('bulk_message_deletion_queue:state_version')).toBe(true);
+	});
+
+	test('does not scan the users table when the state version is fresh', async () => {
+		const userId = createUserID(5678n);
+		const harness = createHarness([createUser(userId, new Date(Date.now() - 60_000))]);
+		harness.kvClient.strings.set('bulk_message_deletion_queue:state_version', Date.now().toString());
+
+		await processPendingBulkMessageDeletions({}, createHelpers());
+
+		expect(harness.scanCalls).toEqual([]);
+		expect(harness.jobs).toEqual([]);
+	});
+
+	test('skips the run when another worker holds the rebuild lock', async () => {
+		const userId = createUserID(9012n);
+		const harness = createHarness([createUser(userId, new Date(Date.now() - 60_000))]);
+		harness.kvClient.strings.set('bulk_message_deletion_queue:rebuild_lock', 'other-worker');
+
+		await processPendingBulkMessageDeletions({}, createHelpers());
+
+		expect(harness.scanCalls).toEqual([]);
+		expect(harness.jobs).toEqual([]);
+		expect(harness.kvClient.strings.get('bulk_message_deletion_queue:rebuild_lock')).toBe('other-worker');
+	});
+});

fluxer_api/src/api/infrastructure/KVBulkMessageDeletionQueueService.ts

@@ -1,8 +1,11 @@
 // SPDX-License-Identifier: AGPL-3.0-or-later
 
+import {generateLockToken} from '@pkgs/cache/src/CacheLockValidation';
 import type {IKVProvider} from '@pkgs/kv_client/src/IKVProvider';
+import {ms, seconds} from 'itty-time';
 import type {UserID} from '../BrandedTypes';
 import {Logger} from '../Logger';
+import type {UserRepository} from '../user/repositories/UserRepository';
 
 interface QueuedBulkMessageDeletion {
 	userId: bigint;
@@ -11,9 +14,15 @@ interface QueuedBulkMessageDeletion {
 
 const QUEUE_KEY = 'bulk_message_deletion_queue';
 const SECONDARY_KEY_PREFIX = 'bulk_message_deletion_queue:';
+const STATE_VERSION_KEY = 'bulk_message_deletion_queue:state_version';
+const REBUILD_LOCK_KEY = 'bulk_message_deletion_queue:rebuild_lock';
+const REBUILD_LOCK_TTL = seconds('5 minutes');
 
 export class KVBulkMessageDeletionQueueService {
-	constructor(private readonly kvClient: IKVProvider) {}
+	constructor(
+		private readonly kvClient: IKVProvider,
+		private readonly userRepository: UserRepository,
+	) {}
 
 	private getSecondaryKey(userId: UserID): string {
 		return `${SECONDARY_KEY_PREFIX}${userId}`;
@@ -31,6 +40,67 @@ export class KVBulkMessageDeletionQueueService {
 		};
 	}
 
+	async needsRebuild(): Promise<boolean> {
+		try {
+			const versionExists = await this.kvClient.exists(STATE_VERSION_KEY);
+			if (!versionExists) {
+				Logger.debug('Bulk message deletion queue needs rebuild: no state version');
+				return true;
+			}
+			const stateVersionStr = await this.kvClient.get(STATE_VERSION_KEY);
+			if (stateVersionStr) {
+				const stateVersion = Number.parseInt(stateVersionStr, 10);
+				const ageMs = Date.now() - stateVersion;
+				if (ageMs > ms('1 day')) {
+					Logger.debug({ageMs, maxAgeMs: ms('1 day')}, 'Bulk message deletion queue needs rebuild: state too old');
+					return true;
+				}
+			}
+			return false;
+		} catch (error) {
+			Logger.error({error}, 'Failed to check if bulk message deletion queue needs rebuild');
+			throw error;
+		}
+	}
+
+	async rebuildState(): Promise<void> {
+		Logger.info('Starting bulk message deletion queue rebuild from primary database');
+		try {
+			await this.kvClient.del(QUEUE_KEY);
+			await this.kvClient.del(STATE_VERSION_KEY);
+			let pageState: string | null = null;
+			let totalProcessed = 0;
+			let totalQueued = 0;
+			const batchSize = 1000;
+			while (true) {
+				const page = await this.userRepository.scanAllUsersPage(batchSize, pageState);
+				const users = page.users;
+				if (users.length === 0) {
+					break;
+				}
+				for (const user of users) {
+					if (user.pendingBulkMessageDeletionAt) {
+						await this.scheduleDeletion(user.id, user.pendingBulkMessageDeletionAt);
+						totalQueued++;
+					}
+				}
+				totalProcessed += users.length;
+				pageState = page.pageState;
+				if (totalProcessed % 10000 === 0) {
+					Logger.debug({totalProcessed, totalQueued}, 'Bulk message deletion queue rebuild progress');
+				}
+				if (!pageState) {
+					break;
+				}
+			}
+			await this.kvClient.set(STATE_VERSION_KEY, Date.now().toString());
+			Logger.info({totalProcessed, totalQueued}, 'Bulk message deletion queue rebuild completed');
+		} catch (error) {
+			Logger.error({error}, 'Failed to rebuild bulk message deletion queue state');
+			throw error;
+		}
+	}
+
 	async scheduleDeletion(userId: UserID, scheduledAt: Date): Promise<void> {
 		try {
 			const entry: QueuedBulkMessageDeletion = {
@@ -81,6 +151,34 @@ export class KVBulkMessageDeletionQueueService {
 		}
 	}
 
+	async acquireRebuildLock(): Promise<string | null> {
+		try {
+			const token = generateLockToken();
+			const acquired = await this.kvClient.acquireLock(REBUILD_LOCK_KEY, token, REBUILD_LOCK_TTL);
+			if (acquired) {
+				Logger.debug({token}, 'Acquired bulk message deletion rebuild lock');
+				return token;
+			}
+			return null;
+		} catch (error) {
+			Logger.error({error}, 'Failed to acquire bulk message deletion rebuild lock');
+			throw error;
+		}
+	}
+
+	async releaseRebuildLock(token: string): Promise<boolean> {
+		try {
+			const released = await this.kvClient.releaseLock(REBUILD_LOCK_KEY, token);
+			if (released) {
+				Logger.debug({token}, 'Released bulk message deletion rebuild lock');
+			}
+			return released;
+		} catch (error) {
+			Logger.error({error, token}, 'Failed to release bulk message deletion rebuild lock');
+			throw error;
+		}
+	}
+
 	async getQueueSize(): Promise<number> {
 		try {
 			return await this.kvClient.zcard(QUEUE_KEY);

fluxer_api/src/api/middleware/ServiceSingletons.ts

@@ -218,7 +218,7 @@ let bulkMessageDeletionQueue: KVBulkMessageDeletionQueueService | null = null;
 export function getKVBulkMessageDeletionQueue(): KVBulkMessageDeletionQueueService {
 	const kvClient = getKVClient();
 	if (!bulkMessageDeletionQueue || bulkMessageDeletionQueueClient !== kvClient) {
-		bulkMessageDeletionQueue = new KVBulkMessageDeletionQueueService(kvClient);
+		bulkMessageDeletionQueue = new KVBulkMessageDeletionQueueService(kvClient, getUserRepository());
 		bulkMessageDeletionQueueClient = kvClient;
 	}
 	return bulkMessageDeletionQueue;

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

@@ -8,6 +8,19 @@ import {getWorkerDependencies} from '../WorkerContext';
 const processPendingBulkMessageDeletions: WorkerTaskHandler = async (_payload, helpers) => {
 	helpers.logger.debug('Processing pending bulk message deletions');
 	const {bulkMessageDeletionQueueService, userRepository, workerService} = getWorkerDependencies();
+	if (await bulkMessageDeletionQueueService.needsRebuild()) {
+		Logger.info('Bulk message deletion queue needs rebuild, acquiring lock');
+		const lockToken = await bulkMessageDeletionQueueService.acquireRebuildLock();
+		if (!lockToken) {
+			Logger.info('Another worker is rebuilding the bulk message deletion queue, skipping this run');
+			return;
+		}
+		try {
+			await bulkMessageDeletionQueueService.rebuildState();
+		} finally {
+			await bulkMessageDeletionQueueService.releaseRebuildLock(lockToken);
+		}
+	}
 	const nowMs = Date.now();
 	const pendingDeletions = await bulkMessageDeletionQueueService.getReadyDeletions(nowMs, 100);
 	Logger.debug({count: pendingDeletions.length}, 'Pending bulk message deletions found');

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

@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: AGPL-3.0-or-later
 
-import type {WorkerTaskHelpers} from '@pkgs/worker/src/contracts/WorkerTask';
+import {JobCancelledError, type WorkerTaskHelpers} from '@pkgs/worker/src/contracts/WorkerTask';
 import {z} from 'zod';
 import {createUserID, type UserID} from '../../BrandedTypes';
 import {createRequestCache} from '../../middleware/RequestCacheMiddleware';
@@ -16,10 +16,7 @@ const PayloadSchema = z.object({
 export async function sendSystemDm(payload: unknown, helpers: WorkerTaskHelpers): Promise<void> {
 	const {content, user_ids} = PayloadSchema.parse(payload);
 	const deps = getWorkerDependencies();
-	const systemUser = await deps.userRepository.findUnique(SYSTEM_USER_ID);
-	if (!systemUser) {
-		throw new Error('System user (id=0) not found');
-	}
+	const systemUser = await deps.userRepository.findUniqueAssert(SYSTEM_USER_ID);
 	const userChannelService = new UserChannelService(
 		deps.userRepository,
 		deps.channelService,
@@ -38,7 +35,8 @@ export async function sendSystemDm(payload: unknown, helpers: WorkerTaskHelpers)
 				{sent, failed, remaining: user_ids.length - sent - failed},
 				'System DM job cancelled mid-flight',
 			);
-			break;
+			requestCache.clear();
+			throw new JobCancelledError();
 		}
 		const recipientId = createUserID(BigInt(raw));
 		try {

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

@@ -0,0 +1,167 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+import type {JsMsg} from 'nats';
+import {afterEach, beforeAll, describe, expect, it, vi} from 'vitest';
+import type {ChannelService} from '../../channel/services/ChannelService';
+import type {IJobLedgerRepository} from '../../jobs/IJobLedgerRepository';
+import {setInjectedWorkerService} from '../../middleware/ServiceRegistry';
+import {NoopWorkerService} from '../../test/NoopWorkerService';
+import type {UserRepository} from '../../user/repositories/UserRepository';
+import {sendSystemDm} from '../tasks/SendSystemDm';
+import {clearWorkerDependencies, setWorkerDependenciesForTest} from '../WorkerContext';
+import {WorkerRunner} from '../WorkerRunner';
+
+const TASK_TYPE = 'sendSystemDm';
+const LEDGER_JOB_ID = 88n;
+
+const queueStub = {
+	getConnectionManager: () => {
+		throw new Error('WorkerRunner tests never consume messages');
+	},
+	getStreamName: () => 'JOBS',
+	publishToDlq: vi.fn(),
+};
+
+class TestWorkerRunner extends WorkerRunner {
+	async runJob(taskType: string, msg: JsMsg): Promise<boolean> {
+		return await this.processJob(taskType, msg);
+	}
+}
+
+function createLedgerStub(cancelAfter: number) {
+	const cancelChecks = {count: 0};
+	const markCancelled = vi.fn(async () => {});
+	const markSucceeded = vi.fn(async () => {});
+	const ledger = {
+		markRunning: async () => {},
+		markSucceeded,
+		markCancelled,
+		markDeadletter: async () => {},
+		reportProgress: async () => {},
+		setContextLink: async () => {},
+		isCancelRequested: async () => {
+			cancelChecks.count += 1;
+			return cancelChecks.count > cancelAfter;
+		},
+	} as unknown as IJobLedgerRepository;
+	return {ledger, markCancelled, markSucceeded};
+}
+
+function createWorkerDependencies() {
+	const sentChannelIds: Array<bigint> = [];
+	const sentUserIds: Array<bigint> = [];
+	const systemUser = {id: 0n, username: 'Fluxer', bot: true, system: true};
+	const userRepository = {
+		findUnique: async () => systemUser,
+		findUniqueAssert: async () => systemUser,
+		findExistingDmState: async () => ({id: 500n}),
+		isDmChannelOpen: async () => true,
+	} as unknown as UserRepository;
+	const channelService = {
+		messages: {
+			send: {
+				sendMessage: async ({channelId, user}: {channelId: bigint; user: {id: bigint}}) => {
+					sentChannelIds.push(channelId);
+					sentUserIds.push(user.id);
+				},
+			},
+		},
+	} as unknown as ChannelService;
+	setWorkerDependenciesForTest({userRepository, channelService});
+	return {sentChannelIds, sentUserIds};
+}
+
+function createJobMessage() {
+	const envelope = {
+		payload: {
+			content: 'scheduled maintenance tonight',
+			user_ids: ['11', '12', '13'],
+			__jobId: LEDGER_JOB_ID.toString(),
+		},
+		max_attempts: 5,
+		priority: 0,
+		created_at: new Date().toISOString(),
+	};
+	return {
+		seq: 1,
+		subject: `jobs.${TASK_TYPE}`,
+		redelivered: false,
+		data: new TextEncoder().encode(JSON.stringify(envelope)),
+		info: {deliveryCount: 1},
+		ack: vi.fn(),
+		nak: vi.fn(),
+		term: vi.fn(),
+		working: vi.fn(),
+	};
+}
+
+describe('System DM cancellation', () => {
+	beforeAll(() => {
+		setInjectedWorkerService(new NoopWorkerService());
+	});
+
+	afterEach(() => {
+		clearWorkerDependencies();
+		queueStub.publishToDlq.mockClear();
+	});
+
+	it('settles the ledger row as cancelled after stopping mid-broadcast', async () => {
+		const deps = createWorkerDependencies();
+		const {ledger, markCancelled, markSucceeded} = createLedgerStub(1);
+		const runner = new TestWorkerRunner({
+			tasks: {[TASK_TYPE]: sendSystemDm},
+			queue: queueStub,
+			consumerName: 'workers_batch',
+			laneName: 'batch',
+			ledger,
+			concurrency: 1,
+		});
+		const msg = createJobMessage();
+
+		await expect(runner.runJob(TASK_TYPE, msg as unknown as JsMsg)).resolves.toBe(false);
+
+		expect(deps.sentChannelIds).toHaveLength(1);
+		expect(markCancelled).toHaveBeenCalledWith(LEDGER_JOB_ID);
+		expect(markSucceeded).not.toHaveBeenCalled();
+		expect(msg.ack).toHaveBeenCalledTimes(1);
+		expect(msg.nak).not.toHaveBeenCalled();
+	});
+
+	it('settles the ledger row as succeeded when no cancel is requested', async () => {
+		const deps = createWorkerDependencies();
+		const {ledger, markCancelled, markSucceeded} = createLedgerStub(Number.POSITIVE_INFINITY);
+		const runner = new TestWorkerRunner({
+			tasks: {[TASK_TYPE]: sendSystemDm},
+			queue: queueStub,
+			consumerName: 'workers_batch',
+			laneName: 'batch',
+			ledger,
+			concurrency: 1,
+		});
+		const msg = createJobMessage();
+
+		await expect(runner.runJob(TASK_TYPE, msg as unknown as JsMsg)).resolves.toBe(true);
+
+		expect(deps.sentChannelIds).toHaveLength(3);
+		expect(markSucceeded).toHaveBeenCalledTimes(1);
+		expect(markCancelled).not.toHaveBeenCalled();
+	});
+
+	it('sends every message as the synthetic system account', async () => {
+		const deps = createWorkerDependencies();
+		const {ledger} = createLedgerStub(Number.POSITIVE_INFINITY);
+		const runner = new TestWorkerRunner({
+			tasks: {[TASK_TYPE]: sendSystemDm},
+			queue: queueStub,
+			consumerName: 'workers_batch',
+			laneName: 'batch',
+			ledger,
+			concurrency: 1,
+		});
+		const msg = createJobMessage();
+
+		await expect(runner.runJob(TASK_TYPE, msg as unknown as JsMsg)).resolves.toBe(true);
+
+		expect(deps.sentUserIds).toEqual([0n, 0n, 0n]);
+	});
+});