diff --git a/fluxer_api/src/api/worker/tasks/IndexChannelMessages.ts b/fluxer_api/src/api/worker/tasks/IndexChannelMessages.ts index 177ce1ef1..335efb975 100644 --- a/fluxer_api/src/api/worker/tasks/IndexChannelMessages.ts +++ b/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 { const services: Array = []; 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())); diff --git a/fluxer_api/src/api/worker/tests/IndexChannelMessages.test.ts b/fluxer_api/src/api/worker/tests/IndexChannelMessages.test.ts new file mode 100644 index 000000000..cd88e5987 --- /dev/null +++ b/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} { + 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); + }); +});