diff --git a/src/services/scheduledExports.test.ts b/src/services/scheduledExports.test.ts index 815d9f8..dcf1a9d 100644 --- a/src/services/scheduledExports.test.ts +++ b/src/services/scheduledExports.test.ts @@ -121,3 +121,49 @@ test('worker starts, processes due schedules, and idles cleanly', async () => { await pool.end(); } }); + +test('worker handles cancellation cooperatively and preserves idempotency on restart', async () => { + const { repository, pool } = createUsageRepository(); + try { + await repository.create({ userId: 'u1', apiId: 'api-1', endpointId: 'ep-1', apiKeyId: 'key-1', developerId: 'dev-1', amount: 15n, requestId: 'req-1', createdAt: new Date('2026-06-01T00:00:00.000Z') }); + + const store = new InMemoryScheduleStore(); + const objectStorage = new HmacObjectStorageClient(); + + let uploadsStarted = 0; + const originalUpload = objectStorage.uploadObject.bind(objectStorage); + objectStorage.uploadObject = async (input) => { + uploadsStarted++; + await new Promise(resolve => setTimeout(resolve, 50)); + if (input.signal?.aborted) { + throw new Error('AbortError'); + } + return originalUpload(input); + }; + + const service = new ScheduledExportsService({ findByApiId: async () => listAllEvents(pool) }, store, objectStorage); + const schedule = await service.createSchedule({ developerId: 'dev-1', name: 'Minute export', cron: '* * * * *', s3Bucket: 'exports', s3Region: 'us-east-1', s3Endpoint: 'https://s3.example.com', s3AccessKeyId: 'akid', s3SecretAccessKey: 'secret', enabled: true }); + + await store.update(schedule.id, { nextRunAt: new Date(0) }); + const expectedStamp = new Date(0).toISOString().replace(/[:.]/g, '-'); + + const worker = createScheduledExportsWorker(service, { intervalMs: 25 }); + worker.start(); + + await new Promise(resolve => setTimeout(resolve, 10)); + worker.stop(); + await worker.awaitIdle(); + + assert.equal(objectStorage.uploads.length, 0); + + worker.start(); + await new Promise(resolve => setTimeout(resolve, 150)); + worker.stop(); + await worker.awaitIdle(); + + assert.equal(objectStorage.uploads.length >= 2, true); + assert.match(objectStorage.uploads[0].key, new RegExp(`-${expectedStamp}.csv`)); + } finally { + await pool.end(); + } +}); diff --git a/src/services/scheduledExports.ts b/src/services/scheduledExports.ts index 620746a..00c8e1d 100644 --- a/src/services/scheduledExports.ts +++ b/src/services/scheduledExports.ts @@ -39,6 +39,7 @@ export interface ObjectStorageClient { secretAccessKey: string; region: string; endpoint: string; + signal?: AbortSignal; }): Promise; createSignedDownloadUrl(input: { bucket: string; @@ -187,7 +188,11 @@ export class HmacObjectStorageClient implements ObjectStorageClient { secretAccessKey: string; region: string; endpoint: string; + signal?: AbortSignal; }): Promise { + if (input.signal?.aborted) { + throw input.signal.reason; + } void input.accessKeyId; void input.secretAccessKey; void input.region; @@ -265,13 +270,16 @@ export class ScheduledExportsService { return updated ? this.redactSecret(updated) : undefined; } - async runDueSchedules(now: Date = new Date()): Promise { + async runDueSchedules(now: Date = new Date(), signal?: AbortSignal): Promise { const schedules = await this.scheduleStore.list(); const dueSchedules = schedules.filter((schedule) => schedule.enabled && schedule.nextRunAt <= now); const results: ExportRunResult[] = []; for (const schedule of dueSchedules) { - results.push(await this.runSchedule(schedule, now)); + if (signal?.aborted) { + throw signal.reason; + } + results.push(await this.runSchedule(schedule, now, signal)); await this.scheduleStore.update(schedule.id, { lastRunAt: now, nextRunAt: computeNextRunAt(schedule.cron, now), @@ -281,7 +289,10 @@ export class ScheduledExportsService { return results; } - async runSchedule(schedule: ExportSchedule, now: Date = new Date()): Promise { + async runSchedule(schedule: ExportSchedule, now: Date = new Date(), signal?: AbortSignal): Promise { + if (signal?.aborted) { + throw signal.reason; + } const allEvents = await this.usageEventsRepository.findByApiId('', undefined, now, undefined, 0); const scopedEvents = allEvents.filter((event: BillingUsageEvent) => { if (event.developerId !== schedule.developerId) return false; @@ -290,7 +301,7 @@ export class ScheduledExportsService { }); const prefix = schedule.s3PathPrefix ? `${schedule.s3PathPrefix.replace(/\/$/, '')}/` : ''; - const stamp = now.toISOString().replace(/[:.]/g, '-'); + const stamp = schedule.nextRunAt.toISOString().replace(/[:.]/g, '-'); const csvKey = `${prefix}usage-events-${schedule.id}-${stamp}.csv`; const jsonKey = `${prefix}usage-events-${schedule.id}-${stamp}.json`; @@ -304,6 +315,7 @@ export class ScheduledExportsService { secretAccessKey: schedule.s3SecretAccessKey, region: schedule.s3Region, endpoint: schedule.s3Endpoint, + signal, }), this.objectStorageClient.uploadObject({ bucket: schedule.s3Bucket, @@ -314,6 +326,7 @@ export class ScheduledExportsService { secretAccessKey: schedule.s3SecretAccessKey, region: schedule.s3Region, endpoint: schedule.s3Endpoint, + signal, }), ]); @@ -369,16 +382,23 @@ export function createScheduledExportsWorker( const log = options.logger ?? logger; let timer: NodeJS.Timeout | null = null; let running: Promise | null = null; + let abortController: AbortController | null = null; const tick = async (): Promise => { if (running) return; - running = service.runDueSchedules(); + abortController = new AbortController(); + running = service.runDueSchedules(new Date(), abortController.signal); try { await running; - } catch (error) { - log.error('scheduled export worker failed', error); + } catch (error: any) { + if (error?.name === 'AbortError') { + log.info('scheduled export worker canceled'); + } else { + log.error('scheduled export worker failed', error); + } } finally { running = null; + abortController = null; } }; @@ -392,6 +412,9 @@ export function createScheduledExportsWorker( if (!timer) return; clearInterval(timer); timer = null; + if (abortController) { + abortController.abort(new Error('AbortError')); + } }, async awaitIdle() { if (running) await running.catch(() => undefined);