fix(session-ingest): delay queue retries safely

Author: johnnyeric

services/session-ingest/src/dos/SessionIngestDO.test.ts

@@ -0,0 +1,16 @@
+import { describe, expect, it, vi } from 'vitest';
+
+vi.mock('cloudflare:workers', () => ({
+  DurableObject: class DurableObject {
+    constructor(_state: unknown, _env: unknown) {}
+  },
+}));
+
+import { ingestOrderCursor } from './SessionIngestDO';
+
+describe('SessionIngestDO ingest ordering', () => {
+  it('uses ingested_at with id only as a tie-breaker for cursor progression', () => {
+    expect(ingestOrderCursor({ ingested_at: 100, id: 7 })).toEqual({ ingestedAt: 100, id: 7 });
+    expect(ingestOrderCursor({ ingested_at: null, id: 3 })).toEqual({ ingestedAt: null, id: 3 });
+  });
+});

services/session-ingest/src/dos/SessionIngestDO.ts

@@ -1,5 +1,5 @@
 import { DurableObject } from 'cloudflare:workers';
-import { eq, ne, gt, gte, lt, and, inArray, isNotNull } from 'drizzle-orm';
+import { eq, ne, gt, gte, lt, and, or, inArray, isNull, isNotNull } from 'drizzle-orm';
 import { drizzle, type DrizzleSqliteDODatabase } from 'drizzle-orm/durable-sqlite';
 import { migrate } from 'drizzle-orm/durable-sqlite/migrator';
 
@@ -86,6 +86,29 @@ const INGEST_META_EXTRACTORS: Array<{
 
 type Changes = Array<{ name: ExtractableMetaKey; value: string | null }>;
 
+export type IngestOrderCursor = { ingestedAt: number | null; id: number };
+
+export function afterIngestOrderCursor(cursor: IngestOrderCursor) {
+  if (cursor.ingestedAt === null) {
+    return or(
+      and(isNull(ingestItems.ingested_at), gt(ingestItems.id, cursor.id)),
+      isNotNull(ingestItems.ingested_at)
+    );
+  }
+
+  return or(
+    gt(ingestItems.ingested_at, cursor.ingestedAt),
+    and(eq(ingestItems.ingested_at, cursor.ingestedAt), gt(ingestItems.id, cursor.id))
+  );
+}
+
+export function ingestOrderCursor(row: {
+  ingested_at: number | null;
+  id: number;
+}): IngestOrderCursor {
+  return { ingestedAt: row.ingested_at, id: row.id };
+}
+
 export class SessionIngestDO extends DurableObject<Env> {
   private db: DrizzleSqliteDODatabase;
 
@@ -295,25 +318,31 @@ export class SessionIngestDO extends DurableObject<Env> {
           // --- messages ---
           const CURSOR_BATCH = 10;
           controller.enqueue(encoder.encode(',"messages":['));
-          let msgCursor = 0;
+          let msgCursor: IngestOrderCursor | undefined;
           let firstMsg = true;
 
           while (true) {
             const msgBatch = db
               .select({
                 id: ingestItems.id,
+                ingested_at: ingestItems.ingested_at,
                 item_id: ingestItems.item_id,
                 item_data: ingestItems.item_data,
                 item_data_r2_key: ingestItems.item_data_r2_key,
               })
               .from(ingestItems)
-              .where(and(eq(ingestItems.item_type, 'message'), gt(ingestItems.id, msgCursor)))
-              .orderBy(ingestItems.id)
+              .where(
+                and(
+                  eq(ingestItems.item_type, 'message'),
+                  msgCursor ? afterIngestOrderCursor(msgCursor) : undefined
+                )
+              )
+              .orderBy(ingestItems.ingested_at, ingestItems.id)
               .limit(CURSOR_BATCH)
               .all();
 
             if (msgBatch.length === 0) break;
-            msgCursor = msgBatch[msgBatch.length - 1].id;
+            msgCursor = ingestOrderCursor(msgBatch[msgBatch.length - 1]);
 
             for (const msgRow of msgBatch) {
               if (!firstMsg) controller.enqueue(encoder.encode(','));
@@ -327,13 +356,14 @@ export class SessionIngestDO extends DurableObject<Env> {
               const msgId = msgRow.item_id.slice('message/'.length);
               const partRange = getPartItemIdentityRange(msgId);
               controller.enqueue(encoder.encode(',"parts":['));
-              let partCursor = 0;
+              let partCursor: IngestOrderCursor | undefined;
               let firstPart = true;
 
               while (true) {
                 const partBatch = db
                   .select({
                     id: ingestItems.id,
+                    ingested_at: ingestItems.ingested_at,
                     item_data: ingestItems.item_data,
                     item_data_r2_key: ingestItems.item_data_r2_key,
                   })
@@ -343,15 +373,15 @@ export class SessionIngestDO extends DurableObject<Env> {
                       eq(ingestItems.item_type, 'part'),
                       gte(ingestItems.item_id, partRange.start),
                       lt(ingestItems.item_id, partRange.end),
-                      gt(ingestItems.id, partCursor)
+                      partCursor ? afterIngestOrderCursor(partCursor) : undefined
                     )
                   )
-                  .orderBy(ingestItems.id)
+                  .orderBy(ingestItems.ingested_at, ingestItems.id)
                   .limit(CURSOR_BATCH)
                   .all();
 
                 if (partBatch.length === 0) break;
-                partCursor = partBatch[partBatch.length - 1].id;
+                partCursor = ingestOrderCursor(partBatch[partBatch.length - 1]);
 
                 for (const partRow of partBatch) {
                   if (!firstPart) controller.enqueue(encoder.encode(','));
@@ -404,7 +434,7 @@ export class SessionIngestDO extends DurableObject<Env> {
       })
       .from(ingestItems)
       .where(ne(ingestItems.item_type, 'session_diff'))
-      .orderBy(ingestItems.id)
+      .orderBy(ingestItems.ingested_at, ingestItems.id)
       .all();
 
     if (rows.length === 0) {

services/session-ingest/src/queue-consumer.test.ts

@@ -143,6 +143,42 @@ describe('createItemExtractor', () => {
 });
 
 describe('queue', () => {
+  it('delays failed queue message retries to avoid immediately hammering hot DOs', async () => {
+    const limit = vi.fn(async () => [{ session_id: 'ses_retry' }]);
+    const where = vi.fn(() => ({ limit }));
+    const from = vi.fn(() => ({ where }));
+    vi.mocked(getWorkerDb).mockReturnValue({ select: vi.fn(() => ({ from })) } as never);
+    const env = {
+      HYPERDRIVE: { connectionString: 'postgres://unused' },
+      SESSION_INGEST_R2: { get: vi.fn(async () => null) },
+    } as never;
+    const ack = vi.fn();
+    const retry = vi.fn();
+
+    await queue(
+      {
+        messages: [
+          {
+            body: {
+              r2Key: 'ingest/retry-missing',
+              kiloUserId: 'usr_retry',
+              sessionId: 'ses_retry',
+              ingestVersion: 1,
+              ingestedAt: 1,
+            },
+            ack,
+            retry,
+          },
+        ],
+      } as never,
+      env,
+      { waitUntil: vi.fn() } as unknown as ExecutionContext
+    );
+
+    expect(ack).not.toHaveBeenCalled();
+    expect(retry).toHaveBeenCalledWith({ delaySeconds: 60 });
+  });
+
   it('passes full parsed oversized message data and its R2 reference into ingest', async () => {
     const ingest = vi.fn(async () => ({ changes: [] }));
     vi.mocked(getSessionIngestDO).mockReturnValue({ ingest } as never);

services/session-ingest/src/queue-consumer.ts

@@ -540,7 +540,7 @@ export async function queue(
         sessionId: msg.body.sessionId,
         error: err instanceof Error ? err.message : String(err),
       });
-      msg.retry();
+      msg.retry({ delaySeconds: 60 });
     }
   }
 }