fix(worker): count each channel once in a bulk reindex (#2286)

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

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: AGPL-3.0-or-later
 
 import type {WorkerTaskHandler} from '@pkgs/worker/src/contracts/WorkerTask';
+import {seconds} from 'itty-time';
 import {z} from 'zod';
 import type {MessageID, UserID} from '../../BrandedTypes';
 import {createChannelID} from '../../BrandedTypes';
@@ -15,6 +16,7 @@ const PayloadSchema = z.object({
 	channelCount: z.number().optional(),
 });
 const BULK_BATCH_SIZE = 5000;
+const COMPLETION_KEY_TTL_SECONDS = seconds('1 day');
 function getMessageIndexServices(): Array<IMessageSearchService> {
 	const services: Array<IMessageSearchService> = [];
 	const defaultService = getMessageSearchService();
@@ -64,7 +66,9 @@ const indexChannelMessages: WorkerTaskHandler = async (payload) => {
 			}
 		}
 		if (validated.completionKey && validated.channelCount) {
-			const completed = await kvClient.incr(validated.completionKey);
+			await kvClient.sadd(validated.completionKey, validated.channelId);
+			await kvClient.expire(validated.completionKey, COMPLETION_KEY_TTL_SECONDS);
+			const completed = await kvClient.scard(validated.completionKey);
 			if (completed >= validated.channelCount) {
 				try {
 					await Promise.all(searchServices.map((s) => s.refreshIndex()));

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

@@ -0,0 +1,67 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+import type {WorkerTaskHelpers} from '@pkgs/worker/src/contracts/WorkerTask';
+import {afterEach, describe, expect, it, vi} from 'vitest';
+import type {ChannelRepository} from '../../channel/ChannelRepository';
+import {setInjectedSearchProvider} from '../../SearchFactory';
+import type {IMessageSearchService} from '../../search/IMessageSearchService';
+import type {ISearchProvider} from '../../search/ISearchProvider';
+import {MockKVProvider} from '../../test/mocks/MockKVProvider';
+import {NoopLogger} from '../../test/mocks/NoopLogger';
+import type {UserRepository} from '../../user/repositories/UserRepository';
+import indexChannelMessages from '../tasks/IndexChannelMessages';
+import {clearWorkerDependencies, setWorkerDependenciesForTest} from '../WorkerContext';
+
+const COMPLETION_KEY = 'bulk_reindex:7000:remaining';
+const HELPERS = {logger: new NoopLogger()} as unknown as WorkerTaskHelpers;
+
+function createHarness(): {kvClient: MockKVProvider; refreshIndex: ReturnType<typeof vi.fn>} {
+	const kvClient = new MockKVProvider();
+	const refreshIndex = vi.fn().mockResolvedValue(undefined);
+	const messageSearchService = {
+		bulkIndexMessages: async () => {},
+		refreshIndex,
+	} as unknown as IMessageSearchService;
+	setInjectedSearchProvider({
+		getMessageSearchService: () => messageSearchService,
+	} as unknown as ISearchProvider);
+	setWorkerDependenciesForTest({
+		kvClient,
+		channelRepository: {
+			listMessages: async () => [],
+			findUnique: async () => null,
+		} as unknown as ChannelRepository,
+		userRepository: {listUsers: async () => []} as unknown as UserRepository,
+	});
+	return {kvClient, refreshIndex};
+}
+
+describe('indexChannelMessages', () => {
+	afterEach(() => {
+		clearWorkerDependencies();
+		setInjectedSearchProvider(undefined);
+	});
+
+	it('counts a redelivered channel once towards bulk reindex completion', async () => {
+		const {kvClient, refreshIndex} = createHarness();
+		const payload = {channelId: '5001', completionKey: COMPLETION_KEY, channelCount: 2};
+
+		await indexChannelMessages(payload, HELPERS);
+		await indexChannelMessages(payload, HELPERS);
+
+		expect(refreshIndex).not.toHaveBeenCalled();
+
+		await indexChannelMessages({...payload, channelId: '5002'}, HELPERS);
+
+		expect(refreshIndex).toHaveBeenCalledTimes(1);
+		expect(await kvClient.exists(COMPLETION_KEY)).toBe(0);
+	});
+
+	it('expires the completion key when the reindex never finishes', async () => {
+		const {kvClient} = createHarness();
+
+		await indexChannelMessages({channelId: '5001', completionKey: COMPLETION_KEY, channelCount: 2}, HELPERS);
+
+		expect(await kvClient.ttl(COMPLETION_KEY)).toBeGreaterThan(0);
+	});
+});