mirror of
https://github.com/fluxerapp/fluxer
synced 2026-10-07 19:22:14 +09:00
fix(worker): rebuild the deletion queue and cancel system dms (#2510)
This commit is contained in:
Hampus
@@ -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');
|
||||
});
|
||||
});
|
||||
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user