feat(core): integrate RetryManager into SegmentDestination upload pipeline

packages/core/src/__tests__/api.test.ts

@@ -25,7 +25,11 @@ describe('#sendEvents', () => {
       .mockReturnValue('2001-01-01T00:00:00.000Z');
   });
 
-  async function sendAnEventPer(writeKey: string, toUrl: string) {
+  async function sendAnEventPer(
+    writeKey: string,
+    toUrl: string,
+    retryCount?: number
+  ) {
     const mockResponse = Promise.resolve('MANOS');
     // eslint-disable-next-line @typescript-eslint/ban-ts-comment
     // @ts-ignore
@@ -60,9 +64,19 @@ describe('#sendEvents', () => {
       writeKey: writeKey,
       url: toUrl,
       events: [event],
+      retryCount,
     });
 
-    expect(fetch).toHaveBeenCalledWith(toUrl, {
+    return event;
+  }
+
+  it('sends an event', async () => {
+    const toSegmentBatchApi = 'https://api.segment.io/v1.b';
+    const writeKey = 'SEGMENT_KEY';
+
+    const event = await sendAnEventPer(writeKey, toSegmentBatchApi);
+
+    expect(fetch).toHaveBeenCalledWith(toSegmentBatchApi, {
       method: 'POST',
       body: JSON.stringify({
         batch: [event],
@@ -71,21 +85,68 @@ describe('#sendEvents', () => {
       }),
       headers: {
         'Content-Type': 'application/json; charset=utf-8',
+        'X-Retry-Count': '0',
       },
+      keepalive: true,
     });
-  }
-
-  it('sends an event', async () => {
-    const toSegmentBatchApi = 'https://api.segment.io/v1.b';
-    const writeKey = 'SEGMENT_KEY';
-
-    await sendAnEventPer(writeKey, toSegmentBatchApi);
   });
 
   it('sends an event to proxy', async () => {
     const toProxyUrl = 'https://myprox.io/b';
     const writeKey = 'SEGMENT_KEY';
 
-    await sendAnEventPer(writeKey, toProxyUrl);
+    const event = await sendAnEventPer(writeKey, toProxyUrl);
+
+    expect(fetch).toHaveBeenCalledWith(toProxyUrl, {
+      method: 'POST',
+      body: JSON.stringify({
+        batch: [event],
+        sentAt: '2001-01-01T00:00:00.000Z',
+        writeKey: 'SEGMENT_KEY',
+      }),
+      headers: {
+        'Content-Type': 'application/json; charset=utf-8',
+        'X-Retry-Count': '0',
+      },
+      keepalive: true,
+    });
+  });
+
+  it('sends X-Retry-Count header with default value 0', async () => {
+    const url = 'https://api.segment.io/v1.b';
+    await sendAnEventPer('KEY', url);
+
+    expect(fetch).toHaveBeenCalledWith(
+      url,
+      expect.objectContaining({
+        headers: expect.objectContaining({
+          'X-Retry-Count': '0',
+        }),
+      })
+    );
+  });
+
+  it('sends X-Retry-Count header with provided retry count', async () => {
+    const url = 'https://api.segment.io/v1.b';
+    await sendAnEventPer('KEY', url, 5);
+
+    expect(fetch).toHaveBeenCalledWith(
+      url,
+      expect.objectContaining({
+        headers: expect.objectContaining({
+          'X-Retry-Count': '5',
+        }),
+      })
+    );
+  });
+
+  it('sends X-Retry-Count as string format', async () => {
+    const url = 'https://api.segment.io/v1.b';
+    await sendAnEventPer('KEY', url, 42);
+
+    const callArgs = (fetch as jest.Mock).mock.calls[0];
+    const headers = callArgs[1].headers;
+    expect(typeof headers['X-Retry-Count']).toBe('string');
+    expect(headers['X-Retry-Count']).toBe('42');
   });
 });

packages/core/src/api.ts

@@ -4,10 +4,12 @@ export const uploadEvents = async ({
   writeKey,
   url,
   events,
+  retryCount = 0,
 }: {
   writeKey: string;
   url: string;
   events: SegmentEvent[];
+  retryCount?: number;
 }) => {
   return await fetch(url, {
     method: 'POST',
@@ -19,6 +21,7 @@ export const uploadEvents = async ({
     }),
     headers: {
       'Content-Type': 'application/json; charset=utf-8',
+      'X-Retry-Count': retryCount.toString(),
     },
   });
 };

packages/core/src/plugins/QueueFlushingPlugin.ts

@@ -15,12 +15,12 @@ export class QueueFlushingPlugin extends UtilityPlugin {
   type = PluginType.after;
 
   private storeKey: string;
-  private isPendingUpload = false;
   private queueStore: Store<{ events: SegmentEvent[] }> | undefined;
   private onFlush: (events: SegmentEvent[]) => Promise<void>;
   private isRestoredResolve: () => void;
   private isRestored: Promise<void>;
   private timeoutWarned = false;
+  private flushPromise?: Promise<void>;
 
   /**
    * @param onFlush callback to execute when the queue is flushed (either by reaching the limit or manually) e.g. code to upload events to your destination
@@ -63,16 +63,36 @@ export class QueueFlushingPlugin extends UtilityPlugin {
 
   async execute(event: SegmentEvent): Promise<SegmentEvent | undefined> {
     await this.queueStore?.dispatch((state) => {
-      const events = [...state.events, event];
+      const stampedEvent = { ...event, _queuedAt: Date.now() };
+      const events = [...state.events, stampedEvent];
       return { events };
     });
     return event;
   }
 
   /**
-   * Calls the onFlush callback with the events in the queue
+   * Calls the onFlush callback with the events in the queue.
+   * Ensures only one flush operation runs at a time.
    */
   async flush() {
+    // Safety: prevent concurrent flush operations
+    if (this.flushPromise) {
+      this.analytics?.logger.info(
+        'Flush already in progress, waiting for completion'
+      );
+      await this.flushPromise;
+      return;
+    }
+
+    this.flushPromise = this._doFlush();
+    try {
+      await this.flushPromise;
+    } finally {
+      this.flushPromise = undefined;
+    }
+  }
+
+  private async _doFlush(): Promise<void> {
     // Wait for the queue to be restored
     try {
       await this.isRestored;
@@ -103,14 +123,7 @@ export class QueueFlushingPlugin extends UtilityPlugin {
     }
 
     const events = (await this.queueStore?.getState(true))?.events ?? [];
-    if (!this.isPendingUpload) {
-      try {
-        this.isPendingUpload = true;
-        await this.onFlush(events);
-      } finally {
-        this.isPendingUpload = false;
-      }
-    }
+    await this.onFlush(events);
   }
 
   /**
@@ -130,6 +143,26 @@ export class QueueFlushingPlugin extends UtilityPlugin {
       return { events: filteredEvents };
     });
   }
+
+  /**
+   * Removes events from the queue by their messageId
+   * @param messageIds array of messageId strings to remove
+   */
+  async dequeueByMessageIds(messageIds: string[]): Promise<void> {
+    await this.queueStore?.dispatch((state) => {
+      if (messageIds.length === 0 || state.events.length === 0) {
+        return state;
+      }
+
+      const idsToRemove = new Set(messageIds);
+      const filteredEvents = state.events.filter(
+        (e) => e.messageId == null || !idsToRemove.has(e.messageId)
+      );
+
+      return { events: filteredEvents };
+    });
+  }
+
   /**
    * Clear all events from the queue
    */