diff --git a/fluxer_api/src/api/infrastructure/KVBulkMessageDeletionQueueService.test.ts b/fluxer_api/src/api/infrastructure/KVBulkMessageDeletionQueueService.test.ts new file mode 100644 index 000000000..ae4540c35 --- /dev/null +++ b/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(); + readonly sortedSets = new Map>(); + + async get(key: string): Promise { + return this.strings.get(key) ?? null; + } + + async set(key: string, value: string): Promise { + this.strings.set(key, value); + return 'OK'; + } + + async del(...keys: Array): Promise { + 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 { + return this.strings.has(key) || this.sortedSets.has(key) ? 1 : 0; + } + + async zcard(key: string): Promise { + return this.sortedSets.get(key)?.size ?? 0; + } + + async zrangebyscore( + key: string, + _min: string | number, + max: string | number, + _limitToken?: string, + offset?: number, + count?: number, + ): Promise> { + const members = [...(this.sortedSets.get(key) ?? new Map()).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 { + let set = this.sortedSets.get(queueKey); + if (!set) { + set = new Map(); + this.sortedSets.set(queueKey, set); + } + set.set(value, score); + this.strings.set(secondaryKey, value); + } + + async removeBulkDeletion(queueKey: string, secondaryKey: string): Promise { + 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 { + if (this.strings.has(key)) { + return false; + } + this.strings.set(key, token); + return true; + } + + async releaseLock(key: string, token: string): Promise { + 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) { + const kvClient = new FakeKVProvider(); + const scanCalls: Array = []; + const userRepository = { + async scanAllUsersPage(limit: number): Promise<{users: Array; pageState: null}> { + scanCalls.push(limit); + return {users, pageState: null}; + }, + async findUnique(userId: UserID): Promise { + 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 { + 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'); + }); +}); diff --git a/fluxer_api/src/api/infrastructure/KVBulkMessageDeletionQueueService.ts b/fluxer_api/src/api/infrastructure/KVBulkMessageDeletionQueueService.ts index 54fb21922..be5184bb8 100644 --- a/fluxer_api/src/api/infrastructure/KVBulkMessageDeletionQueueService.ts +++ b/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 { + 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 { + 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 { try { const entry: QueuedBulkMessageDeletion = { @@ -81,6 +151,34 @@ export class KVBulkMessageDeletionQueueService { } } + async acquireRebuildLock(): Promise { + 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 { + 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 { try { return await this.kvClient.zcard(QUEUE_KEY); diff --git a/fluxer_api/src/api/middleware/ServiceSingletons.ts b/fluxer_api/src/api/middleware/ServiceSingletons.ts index 465e46b6b..56036f04b 100644 --- a/fluxer_api/src/api/middleware/ServiceSingletons.ts +++ b/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; diff --git a/fluxer_api/src/api/worker/tasks/ProcessPendingBulkMessageDeletions.ts b/fluxer_api/src/api/worker/tasks/ProcessPendingBulkMessageDeletions.ts index d9b4cac3e..a902fdd2c 100644 --- a/fluxer_api/src/api/worker/tasks/ProcessPendingBulkMessageDeletions.ts +++ b/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'); diff --git a/fluxer_api/src/api/worker/tasks/SendSystemDm.ts b/fluxer_api/src/api/worker/tasks/SendSystemDm.ts index 32de77c8e..2e4f987c8 100644 --- a/fluxer_api/src/api/worker/tasks/SendSystemDm.ts +++ b/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 { 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 { diff --git a/fluxer_api/src/api/worker/tests/SendSystemDmCancel.test.ts b/fluxer_api/src/api/worker/tests/SendSystemDmCancel.test.ts new file mode 100644 index 000000000..3d680bfa5 --- /dev/null +++ b/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 { + 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 = []; + const sentUserIds: Array = []; + 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]); + }); +});