From 7f3e72d982e33cdc3a59198f4a4989178d5bc4df Mon Sep 17 00:00:00 2001 From: Evan Hahn Date: Tue, 13 Oct 2026 10:52:21 -0500 Subject: [PATCH] Started waiting for active Tinybird batch uploads before shutdown closes https://linear.app/ghost/issue/NY-1575 If the Tinybird sync worker stopped while a batch was uploading, we might NACK those pub/sub messages which we didn't need to. This fixes that. --- .../TinybirdSyncBatchWorker.ts | 17 ++++ .../tinybird-sync-batch-worker.test.ts | 84 +++++++++++++++++++ 2 files changed, 101 insertions(+) diff --git a/src/services/tinybird-sync-worker/TinybirdSyncBatchWorker.ts b/src/services/tinybird-sync-worker/TinybirdSyncBatchWorker.ts index cebbead0..820c13e6 100644 --- a/src/services/tinybird-sync-worker/TinybirdSyncBatchWorker.ts +++ b/src/services/tinybird-sync-worker/TinybirdSyncBatchWorker.ts @@ -37,6 +37,7 @@ class TinybirdSyncBatchWorker { #flushInterval: number; #flushTimer: NodeJS.Timeout | null; #isShuttingDown: boolean; + #activeFlushes = new Set>(); constructor(subscriptionName: string, tinybirdClients: TinybirdSyncClients, config: TinybirdSyncBatchWorkerConfig) { logger.info({event: 'TinybirdSyncBatchWorkerCreating', subscriptionName}); @@ -78,11 +79,17 @@ class TinybirdSyncBatchWorker { this.#flushTimer = null; } + await Promise.allSettled(this.#activeFlushes); await this.#subscriber.close(); await this.#flushAllBatches(); } async #handleMessage(message: Message): Promise { + if (this.#isShuttingDown) { + message.nack(); + return; + } + const tinybirdSyncEvent = parseMessage(message); if (!tinybirdSyncEvent) { message.nack(); @@ -106,6 +113,16 @@ class TinybirdSyncBatchWorker { } async #flushBatch(type: TinybirdSyncEvent['type']): Promise { + const flush = this.#sendBatch(type); + this.#activeFlushes.add(flush); + try { + await flush; + } finally { + this.#activeFlushes.delete(flush); + } + } + + async #sendBatch(type: TinybirdSyncEvent['type']): Promise { const batch = this.#batches[type]; if (batch.length === 0) { return; diff --git a/test/unit/services/tinybird-sync-worker/tinybird-sync-batch-worker.test.ts b/test/unit/services/tinybird-sync-worker/tinybird-sync-batch-worker.test.ts index b9f7b582..72a4872c 100644 --- a/test/unit/services/tinybird-sync-worker/tinybird-sync-batch-worker.test.ts +++ b/test/unit/services/tinybird-sync-worker/tinybird-sync-batch-worker.test.ts @@ -104,13 +104,16 @@ describe('TinybirdSyncBatchWorker', () => { ]); await handleMessage(steps[1]); + expect(stepClient.postEventBatch).toHaveBeenCalledOnce(); [...runs, ...steps].forEach(value => expect(value.ack).toHaveBeenCalledOnce()); }); it('nacks invalid messages without sending them to Tinybird', async () => { const value = message('not-json'); + await handleMessage(value); + expect(value.nack).toHaveBeenCalledOnce(); expect(value.ack).not.toHaveBeenCalled(); expect(runClient.postEventBatch).not.toHaveBeenCalled(); @@ -121,9 +124,11 @@ describe('TinybirdSyncBatchWorker', () => { runClient.postEventBatch.mockRejectedValue(new Error('Tinybird unavailable')); const runs = [message(runEvent()), message(runEvent('6a99cd8cb5ac7c0052553385'))]; const steps = [message(stepEvent()), message(stepEvent('6a99cd8cb5ac7c0052553386'))]; + for (const value of [...runs, ...steps]) { await handleMessage(value); } + runs.forEach(value => expect(value.nack).toHaveBeenCalledOnce()); steps.forEach(value => expect(value.ack).toHaveBeenCalledOnce()); }); @@ -131,9 +136,11 @@ describe('TinybirdSyncBatchWorker', () => { it('flushes partial batches when stopped', async () => { const run = message(runEvent()); const step = message(stepEvent()); + await handleMessage(run); await handleMessage(step); await worker.stop(); + expect(runClient.postEventBatch).toHaveBeenCalledOnce(); expect(stepClient.postEventBatch).toHaveBeenCalledOnce(); expect(run.ack).toHaveBeenCalledOnce(); @@ -147,15 +154,92 @@ describe('TinybirdSyncBatchWorker', () => { expect(subscriberMocks.close).toHaveBeenCalledOnce(); }); + it('waits for a successful batch upload before closing', async () => { + const upload = Promise.withResolvers(); + runClient.postEventBatch.mockReturnValueOnce(upload.promise); + const run = message(runEvent()); + await handleMessage(message(runEvent())); + + const handling = handleMessage(run); + const stopping = worker.stop(); + const closedDuringUpload = subscriberMocks.close.mock.calls.length > 0; + upload.resolve(); + await handling; + await stopping; + + expect(closedDuringUpload).toBe(false); + expect(run.ack).toHaveBeenCalledOnce(); + expect(run.ack).toHaveBeenCalledBefore(subscriberMocks.close); + expect(subscriberMocks.close).toHaveBeenCalledOnce(); + }); + + it('waits for a failed batch upload before closing', async () => { + const upload = Promise.withResolvers(); + runClient.postEventBatch.mockReturnValueOnce(upload.promise); + const run = message(runEvent()); + await handleMessage(message(runEvent())); + + const handling = handleMessage(run); + const stopping = worker.stop(); + const closedDuringUpload = subscriberMocks.close.mock.calls.length > 0; + upload.reject(new Error('Tinybird unavailable')); + await handling; + await stopping; + + expect(closedDuringUpload).toBe(false); + expect(run.nack).toHaveBeenCalledOnce(); + expect(run.nack).toHaveBeenCalledBefore(subscriberMocks.close); + expect(subscriberMocks.close).toHaveBeenCalledOnce(); + }); + + it('nacks messages received after shutdown starts', async () => { + const stopping = worker.stop(); + const run = message(runEvent()); + + await handleMessage(run); + + await stopping; + expect(run.nack).toHaveBeenCalledOnce(); + expect(run.ack).not.toHaveBeenCalled(); + expect(runClient.postEventBatch).not.toHaveBeenCalled(); + }); + + it('waits for an active timer flush before closing', async () => { + await worker.stop(); + vi.useFakeTimers(); + start(2, 1000); + subscriberMocks.close.mockClear(); + const upload = Promise.withResolvers(); + runClient.postEventBatch.mockReturnValueOnce(upload.promise); + const run = message(runEvent()); + await handleMessage(run); + + vi.advanceTimersByTime(1000); + + expect(runClient.postEventBatch).toHaveBeenCalledOnce(); + + const stopping = worker.stop(); + const closedDuringUpload = subscriberMocks.close.mock.calls.length > 0; + upload.resolve(); + await stopping; + + expect(closedDuringUpload).toBe(false); + expect(run.ack).toHaveBeenCalledBefore(subscriberMocks.close); + expect(subscriberMocks.close).toHaveBeenCalledOnce(); + expect(vi.getTimerCount()).toBe(0); + }); + it('flushes partial batches on configured interval', async () => { await worker.stop(); vi.useFakeTimers(); start(2, 1000); const run = message(runEvent()); const step = message(stepEvent()); + await handleMessage(run); await handleMessage(step); await vi.advanceTimersByTimeAsync(1000); + expect(runClient.postEventBatch).toHaveBeenCalledOnce(); expect(stepClient.postEventBatch).toHaveBeenCalledOnce(); });