restore failed state

Author: conico974

examples/e2e/app-router/wrangler.jsonc

@@ -19,7 +19,7 @@
   "migrations": [
     {
       "tag": "v1",
-      "new_classes": ["DurableObjectQueueHandler"]
+      "new_sqlite_classes": ["DurableObjectQueueHandler"]
     }
   ],
   "kv_namespaces": [

packages/cloudflare/src/api/durable-objects/queue.spec.ts

@@ -26,6 +26,9 @@ const createDurableObjectQueue = ({
     storage: {
       setAlarm: vi.fn(),
       getAlarm: vi.fn(),
+      sql: {
+        exec: vi.fn(),
+      },
     },
   };
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -102,8 +105,9 @@ describe("DurableObjectQueue", () => {
       expect(queue.ongoingRevalidations.has("id6")).toBe(false);
       expect(Array.from(queue.ongoingRevalidations.keys())).toEqual(["id", "id2", "id3", "id4", "id5"]);
 
+      // BlockConcurrencyWhile is called twice here, first time during creation of the object and second time when we try to revalidate
       // @ts-expect-error
-      expect(queue.ctx.blockConcurrencyWhile).toHaveBeenCalledTimes(1);
+      expect(queue.ctx.blockConcurrencyWhile).toHaveBeenCalledTimes(2);
 
       // Here we await the blocked request to ensure it's resolved
       await blockedReq;
@@ -201,9 +205,10 @@ describe("DurableObjectQueue", () => {
 
     it("should add an alarm if there are failed states", async () => {
       const queue = createDurableObjectQueue({ fetchDuration: 10 });
-      queue.routeInFailedState.set("id", { msg: createMessage("id"), retryCount: 0, nextAlarmMs: 1000 });
+      const nextAlarm = Date.now() + 1000;
+      queue.routeInFailedState.set("id", { msg: createMessage("id"), retryCount: 0, nextAlarmMs: nextAlarm });
       await queue.addAlarm();
-      expect(getStorage(queue).setAlarm).toHaveBeenCalledWith(1000);
+      expect(getStorage(queue).setAlarm).toHaveBeenCalledWith(nextAlarm);
     });
 
     it("should not add an alarm if there is already an alarm set", async () => {
@@ -217,10 +222,16 @@ describe("DurableObjectQueue", () => {
 
     it("should set the alarm to the lowest nextAlarm", async () => {
       const queue = createDurableObjectQueue({ fetchDuration: 10 });
-      queue.routeInFailedState.set("id", { msg: createMessage("id"), retryCount: 0, nextAlarmMs: 1000 });
-      queue.routeInFailedState.set("id2", { msg: createMessage("id2"), retryCount: 0, nextAlarmMs: 500 });
+      const nextAlarm = Date.now() + 1000;
+      const firstAlarm = Date.now() + 500;
+      queue.routeInFailedState.set("id", { msg: createMessage("id"), retryCount: 0, nextAlarmMs: nextAlarm });
+      queue.routeInFailedState.set("id2", {
+        msg: createMessage("id2"),
+        retryCount: 0,
+        nextAlarmMs: firstAlarm,
+      });
       await queue.addAlarm();
-      expect(getStorage(queue).setAlarm).toHaveBeenCalledWith(500);
+      expect(getStorage(queue).setAlarm).toHaveBeenCalledWith(firstAlarm);
     });
   });
 

packages/cloudflare/src/api/durable-objects/queue.ts

@@ -15,17 +15,21 @@ interface ExtendedQueueMessage extends QueueMessage {
   previewModeId: string;
 }
 
+interface FailedState {
+  msg: ExtendedQueueMessage;
+  retryCount: number;
+  nextAlarmMs: number;
+}
+
 export class DurableObjectQueueHandler extends DurableObject<CloudflareEnv> {
   // Ongoing revalidations are deduped by the deduplication id
   // Since this is running in waitUntil, we expect the durable object state to persist this during the duration of the revalidation
   // TODO: handle incremental cache with only eventual consistency (i.e. KV or R2/D1 with the optional cache layer on top)
   ongoingRevalidations = new Map<string, Promise<void>>();
 
-  // TODO: restore the state of the failed revalidations - Probably in the next PR where i'll add the storage
-  routeInFailedState = new Map<
-    string,
-    { msg: ExtendedQueueMessage; retryCount: number; nextAlarmMs: number }
-  >();
+  sql: SqlStorage;
+
+  routeInFailedState = new Map<string, FailedState>();
 
   service: NonNullable<CloudflareEnv["NEXT_CACHE_REVALIDATION_WORKER"]>;
 
@@ -37,6 +41,10 @@ export class DurableObjectQueueHandler extends DurableObject<CloudflareEnv> {
     this.service = env.NEXT_CACHE_REVALIDATION_WORKER!;
     // If there is no service binding, we throw an error because we can't revalidate without it
     if (!this.service) throw new IgnorableError("No service binding for cache revalidation worker");
+    this.sql = ctx.storage.sql;
+
+    // We restore the state
+    ctx.blockConcurrencyWhile(() => this.initState());
   }
 
   async revalidate(msg: ExtendedQueueMessage) {
@@ -71,7 +79,6 @@ export class DurableObjectQueueHandler extends DurableObject<CloudflareEnv> {
       } = msg;
       const protocol = host.includes("localhost") ? "http" : "https";
 
-      //TODO: handle the different types of errors that can occur during the fetch (i.e. timeout, network error, etc)
       const response = await this.service.fetch(`${protocol}://${host}${url}`, {
         method: "HEAD",
         headers: {
@@ -133,6 +140,8 @@ export class DurableObjectQueueHandler extends DurableObject<CloudflareEnv> {
   async addToFailedState(msg: ExtendedQueueMessage) {
     const existingFailedState = this.routeInFailedState.get(msg.MessageDeduplicationId);
 
+    let updatedFailedState: FailedState;
+
     if (existingFailedState) {
       if (existingFailedState.retryCount >= 6) {
         // We give up after 6 retries and log the error
@@ -143,18 +152,24 @@ export class DurableObjectQueueHandler extends DurableObject<CloudflareEnv> {
         return;
       }
       const nextAlarmMs = Date.now() + Math.pow(2, existingFailedState.retryCount + 1) * 2_000;
-      this.routeInFailedState.set(msg.MessageDeduplicationId, {
+      updatedFailedState = {
         ...existingFailedState,
         retryCount: existingFailedState.retryCount + 1,
         nextAlarmMs,
-      });
+      };
     } else {
-      this.routeInFailedState.set(msg.MessageDeduplicationId, {
+      updatedFailedState = {
         msg,
         retryCount: 1,
         nextAlarmMs: Date.now() + 2_000,
-      });
+      };
     }
+    this.routeInFailedState.set(msg.MessageDeduplicationId, updatedFailedState);
+    this.sql.exec(
+      "INSERT OR REPLACE INTO failed_state (id, data) VALUES (?, ?)",
+      msg.MessageDeduplicationId,
+      JSON.stringify(updatedFailedState)
+    );
     // We probably want to do something if routeInFailedState is becoming too big, at least log it
     await this.addAlarm();
   }
@@ -164,9 +179,29 @@ export class DurableObjectQueueHandler extends DurableObject<CloudflareEnv> {
     if (existingAlarm) return;
     if (this.routeInFailedState.size === 0) return;
 
-    const nextAlarmToSetup = Math.min(
+    let nextAlarmToSetup = Math.min(
       ...Array.from(this.routeInFailedState.values()).map(({ nextAlarmMs }) => nextAlarmMs)
     );
+    if (nextAlarmToSetup < Date.now()) {
+      // We don't want to set an alarm in the past
+      nextAlarmToSetup = Date.now() + 2_000;
+    }
     await this.ctx.storage.setAlarm(nextAlarmToSetup);
   }
+
+  // This function is used to restore the state of the durable object
+  // We don't restore the ongoing revalidations because we cannot know in which state they are
+  // We only restore the failed state and the alarm
+  async initState() {
+    // We store the failed state as a blob, we don't want to do anything with it anyway besides restoring
+    this.sql.exec("CREATE TABLE IF NOT EXISTS failed_state (id TEXT PRIMARY KEY, data TEXT)");
+
+    const failedStateCursor = this.sql.exec<{ id: string; data: string }>("SELECT * FROM failed_state");
+    for (const row of failedStateCursor) {
+      this.routeInFailedState.set(row.id, JSON.parse(row.data));
+    }
+
+    // Now that we have restored the failed state, we can restore the alarm as well
+    await this.addAlarm();
+  }
 }