From 167862a8a6a52d100c5d4e4a64ee6d10859b16ef Mon Sep 17 00:00:00 2001 From: Hampus Date: Wed, 9 Sep 2026 20:59:38 +0200 Subject: [PATCH] perf(api): harvest messages a page at a time (#2633) --- .../api/worker/tasks/HarvestUserData.test.ts | 68 +++++++ .../src/api/worker/tasks/HarvestUserData.ts | 175 +++++++++--------- 2 files changed, 159 insertions(+), 84 deletions(-) create mode 100644 fluxer_api/src/api/worker/tasks/HarvestUserData.test.ts diff --git a/fluxer_api/src/api/worker/tasks/HarvestUserData.test.ts b/fluxer_api/src/api/worker/tasks/HarvestUserData.test.ts new file mode 100644 index 000000000..a3e6f8c14 --- /dev/null +++ b/fluxer_api/src/api/worker/tasks/HarvestUserData.test.ts @@ -0,0 +1,68 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +import {describe, expect, it} from 'vitest'; +import {type ChannelID, createChannelID, createMessageID, createUserID, type MessageID} from '../../BrandedTypes'; +import {harvestMessages} from './HarvestUserData'; + +const AUTHOR = createUserID(1000000000000000000n); + +function makeRepository(refs: Array<{channelId: ChannelID; messageId: MessageID}>, missing = new Set()) { + let pages = 0; + return { + pages: () => pages, + listMessagesByAuthor: async (_userId: typeof AUTHOR, limit: number, lastMessageId?: MessageID) => { + pages++; + const start = lastMessageId ? refs.findIndex((r) => r.messageId === lastMessageId) + 1 : 0; + return refs.slice(start, start + limit); + }, + getMessage: async (_channelId: ChannelID, messageId: MessageID) => + missing.has(messageId.toString()) ? null : {content: `body ${messageId.toString()}`, attachments: undefined}, + }; +} + +function refsAcross(channelCount: number, perChannel: number) { + const out: Array<{channelId: ChannelID; messageId: MessageID}> = []; + let id = 1500000000000000000n; + for (let c = 0; c < channelCount; c++) { + const channelId = createChannelID(2000000000000000000n + BigInt(c)); + for (let m = 0; m < perChannel; m++) { + out.push({channelId, messageId: createMessageID(id)}); + id += 1n; + } + } + return out; +} + +describe('harvestMessages', () => { + it('reads past a single page instead of stopping at one query', async () => { + const refs = refsAcross(1, 2500); + const repo = makeRepository(refs); + const result = await harvestMessages(repo, AUTHOR, Date.now(), null); + expect(result.totalMessages).toBe(2500); + expect(repo.pages()).toBeGreaterThan(1); + }); + + it('groups every message under its own channel', async () => { + const repo = makeRepository(refsAcross(3, 4)); + const result = await harvestMessages(repo, AUTHOR, Date.now(), null); + expect(result.channelMessagesMap.size).toBe(3); + for (const messages of result.channelMessagesMap.values()) { + expect(messages).toHaveLength(4); + } + expect(result.totalMessages).toBe(12); + }); + + it('leaves out a message the repository cannot return', async () => { + const refs = refsAcross(1, 5); + const repo = makeRepository(refs, new Set([refs[2].messageId.toString()])); + const result = await harvestMessages(repo, AUTHOR, Date.now(), null); + expect(result.totalMessages).toBe(4); + }); + + it('returns an empty map for an account with no messages', async () => { + const repo = makeRepository([]); + const result = await harvestMessages(repo, AUTHOR, Date.now(), null); + expect(result.totalMessages).toBe(0); + expect(result.channelMessagesMap.size).toBe(0); + }); +}); diff --git a/fluxer_api/src/api/worker/tasks/HarvestUserData.ts b/fluxer_api/src/api/worker/tasks/HarvestUserData.ts index 0ae9c7fd9..3b4c87784 100644 --- a/fluxer_api/src/api/worker/tasks/HarvestUserData.ts +++ b/fluxer_api/src/api/worker/tasks/HarvestUserData.ts @@ -230,7 +230,7 @@ interface HarvestMessagesFilterArgs { findChannel: (channelId: ChannelID) => Promise; } -async function harvestMessages( +export async function harvestMessages( channelRepository: { listMessagesByAuthor: ( userId: UserID, @@ -255,36 +255,78 @@ async function harvestMessages( filterArgs: HarvestMessagesFilterArgs | null, ): Promise { const channelMessagesMap = new Map>(); + const channelEligibility = filterArgs ? new Map() : null; Logger.debug('Fetching all user messages'); const startFetchTime = Date.now(); - const messageRefs: Array<{channelId: ChannelID; messageId: MessageID}> = []; let lastMessageId: MessageID | undefined; + let scannedMessages = 0; + let totalMessages = 0; + + const readMessage = async ({ + channelId, + messageId, + }: { + channelId: ChannelID; + messageId: MessageID; + }): Promise => { + try { + const message = await channelRepository.getMessage(channelId, messageId); + if (!message) { + Logger.warn( + {channelId: channelId.toString(), messageId: messageId.toString()}, + 'Message not found during harvest', + ); + return null; + } + const timestamp = snowflakeToDate(messageId); + const attachments: Array = []; + if (message.attachments) { + for (const attachment of message.attachments) { + attachments.push({ + attachment_id: attachment.id.toString(), + filename: attachment.filename, + size: attachment.size.toString(), + content_type: attachment.contentType, + content_hash: null, + archive_path: null, + cdn_url: makeAttachmentCdnUrl(channelId, attachment.id, attachment.filename), + width: attachment.width, + height: attachment.height, + }); + } + } + return { + channelId: channelId.toString(), + messageData: { + id: messageId.toString(), + timestamp: timestamp.toISOString(), + content: message.content ?? '', + attachments, + }, + }; + } catch (error) { + Logger.error( + {error, channelId: channelId.toString(), messageId: messageId.toString()}, + 'Failed to process message during harvest', + ); + return null; + } + }; + while (true) { const page = await channelRepository.listMessagesByAuthor(userId, HARVEST_MESSAGE_CHUNK_SIZE, lastMessageId); if (page.length === 0) { break; } - messageRefs.push(...page); + scannedMessages += page.length; lastMessageId = page[page.length - 1].messageId; - if (page.length < HARVEST_MESSAGE_CHUNK_SIZE) { - break; - } - } - Logger.debug( - { - totalMessages: messageRefs.length, - fetchElapsed: Date.now() - startFetchTime, - totalElapsed: Date.now() - startTime, - }, - 'All messages retrieved', - ); - if (messageRefs.length === 0) { - return {channelMessagesMap, totalMessages: 0}; - } - const channelEligibility = filterArgs ? new Map() : null; - const filteredRefs: Array<{channelId: ChannelID; messageId: MessageID}> = []; - if (filterArgs) { - for (const ref of messageRefs) { + + const pageRefs: Array<{channelId: ChannelID; messageId: MessageID}> = []; + for (const ref of page) { + if (!filterArgs) { + pageRefs.push(ref); + continue; + } const ts = snowflakeToDate(ref.messageId).getTime(); if (!isTimestampInWindow(ts, filterArgs.filter)) { continue; @@ -297,76 +339,41 @@ async function harvestMessages( channelEligibility!.set(channelIdStr, eligible); } if (eligible) { - filteredRefs.push(ref); + pageRefs.push(ref); } } - } else { - filteredRefs.push(...messageRefs); - } - if (filteredRefs.length === 0) { - return {channelMessagesMap, totalMessages: 0}; - } - const messages: Array = []; - for (let i = 0; i < filteredRefs.length; i += CONCURRENT_MESSAGE_LIMIT) { - const batch = filteredRefs.slice(i, i + CONCURRENT_MESSAGE_LIMIT); - const batchPromises = batch.map(async ({channelId, messageId}): Promise => { - try { - const message = await channelRepository.getMessage(channelId, messageId); - if (!message) { - Logger.warn( - {channelId: channelId.toString(), messageId: messageId.toString()}, - 'Message not found during harvest', - ); - return null; + + for (let i = 0; i < pageRefs.length; i += CONCURRENT_MESSAGE_LIMIT) { + const batchResults = await Promise.all(pageRefs.slice(i, i + CONCURRENT_MESSAGE_LIMIT).map(readMessage)); + for (const result of batchResults) { + if (result === null) { + continue; } - const timestamp = snowflakeToDate(messageId); - const attachments: Array = []; - if (message.attachments) { - for (const attachment of message.attachments) { - attachments.push({ - attachment_id: attachment.id.toString(), - filename: attachment.filename, - size: attachment.size.toString(), - content_type: attachment.contentType, - content_hash: null, - archive_path: null, - cdn_url: makeAttachmentCdnUrl(channelId, attachment.id, attachment.filename), - width: attachment.width, - height: attachment.height, - }); - } + let bucket = channelMessagesMap.get(result.channelId); + if (!bucket) { + bucket = []; + channelMessagesMap.set(result.channelId, bucket); } - return { - channelId: channelId.toString(), - messageData: { - id: messageId.toString(), - timestamp: timestamp.toISOString(), - content: message.content ?? '', - attachments, - }, - }; - } catch (error) { - Logger.error( - {error, channelId: channelId.toString(), messageId: messageId.toString()}, - 'Failed to process message during harvest', - ); - return null; - } - }); - const batchResults = await Promise.all(batchPromises); - for (const result of batchResults) { - if (result !== null) { - messages.push(result); + bucket.push(result.messageData); + totalMessages++; } } - } - for (const {channelId, messageData} of messages) { - if (!channelMessagesMap.has(channelId)) { - channelMessagesMap.set(channelId, []); + + if (page.length < HARVEST_MESSAGE_CHUNK_SIZE) { + break; } - channelMessagesMap.get(channelId)!.push(messageData); } - return {channelMessagesMap, totalMessages: messages.length}; + + Logger.debug( + { + scannedMessages, + totalMessages, + fetchElapsed: Date.now() - startFetchTime, + totalElapsed: Date.now() - startTime, + }, + 'All messages retrieved', + ); + return {channelMessagesMap, totalMessages}; } function buildUserDataJson(params: UserDataJsonParams) {