Feature: Alarm Flush (#324)

* flush working + tests + jest

* remove unused

* update name

* flushMiniflareDurableObjectAlarms returned nothing

Author: CraigglesO

package-lock.json

@@ -214,17 +214,40 @@ export class DurableObjectsPlugin
       const id = new DurableObjectId(objectName, hexId);
       const state = await this.getObject(storage, id);
       // execute the alarm
-      await new RequestContext({
-        requestDepth: 1,
-        pipelineDepth: 1,
-        durableObject: true,
-        externalSubrequestLimit: usageModelExternalSubrequestLimit(
-          this.ctx.usageModel
-        ),
-      }).runWith(() => state[kAlarm]());
+      await this.#executeAlarm(state);
     });
   }
 
+  async flushAlarms(
+    storageFactory: StorageFactory,
+    ids?: DurableObjectId[]
+  ): Promise<void> {
+    if (ids !== undefined) {
+      for (const id of ids) {
+        const state = await this.getObject(storageFactory, id);
+        await this.#executeAlarm(state);
+        this.#alarmStore.deleteAlarm(`${id[kObjectName]}:${id.toString()}`);
+      }
+    } else {
+      // otherwise flush all alarms
+      for (const [key, state] of this.#objects) {
+        await this.#executeAlarm(await state);
+        this.#alarmStore.deleteAlarm(key);
+      }
+    }
+  }
+
+  async #executeAlarm(state: DurableObjectState): Promise<void> {
+    await new RequestContext({
+      requestDepth: 1,
+      pipelineDepth: 1,
+      durableObject: true,
+      externalSubrequestLimit: usageModelExternalSubrequestLimit(
+        this.ctx.usageModel
+      ),
+    }).runWith(() => state[kAlarm]());
+  }
+
   beforeReload(): void {
     // Clear instance map, this should cause old instances to be GCed
     this.#objects.clear();

packages/durable-objects/src/plugin.ts

@@ -363,3 +363,87 @@ test("DurableObjectsPlugin: set alarm and run list filters out alarm", async (t)
   const res1 = await ns1.get(ns1.newUniqueId()).fetch("/");
   t.is(await res1.text(), "{}");
 });
+
+test("DurableObjectsPlugin: flush alarms", async (t) => {
+  class TestObject implements DurableObject {
+    constructor(private readonly state: DurableObjectState) {}
+
+    fetch = async () => {
+      await this.state.storage.setAlarm(Date.now() + 60 * 1000);
+      return new Response("ok");
+    };
+    alarm = async () => {
+      await this.state.storage.put("a", 1);
+    };
+  }
+
+  const map = new Map<string, StoredValue>();
+  const factory = new MemoryStorageFactory({
+    [`test://map:TEST:${testId.toString()}`]: map,
+  });
+  const plugin = new DurableObjectsPlugin(ctx, {
+    durableObjects: { TEST: "TestObject" },
+    durableObjectsPersist: "test://map",
+    durableObjectsAlarms: true,
+  });
+  await plugin.setup(factory);
+  plugin.beforeReload();
+  plugin.reload({}, { TestObject }, new Map());
+
+  const ns = plugin.getNamespace(factory, "TEST");
+  const res = await ns.get(testId).fetch("/");
+  t.is(await res.text(), "ok");
+  // now flush
+  await plugin.flushAlarms(factory);
+  // check storage "a" is 1
+  const state = await plugin.getObject(factory, testId);
+  t.is(await state.storage.get("a"), 1);
+});
+
+test("DurableObjectsPlugin: flush a specific alarm", async (t) => {
+  class TestObject implements DurableObject {
+    constructor(private readonly state: DurableObjectState) {}
+
+    fetch = async () => {
+      await this.state.storage.setAlarm(Date.now() + 60 * 1000);
+      return new Response("ok");
+    };
+    alarm = async () => {
+      await this.state.storage.put("key", 1);
+    };
+  }
+
+  const mapA = new Map<string, StoredValue>();
+  const mapB = new Map<string, StoredValue>();
+  const factory = new MemoryStorageFactory({
+    [`test://map:TEST:a`]: mapA,
+    [`test://map:TEST:b`]: mapB,
+  });
+  const plugin = new DurableObjectsPlugin(ctx, {
+    durableObjects: { TEST: "TestObject" },
+    durableObjectsPersist: "test://map",
+    durableObjectsAlarms: true,
+  });
+  await plugin.setup(factory);
+  plugin.beforeReload();
+  plugin.reload({}, { TestObject }, new Map());
+
+  const ns = plugin.getNamespace(factory, "TEST");
+  const idA = ns.idFromName("a");
+  const idB = ns.idFromName("b");
+
+  const resA = await ns.get(idA).fetch("/");
+  t.is(await resA.text(), "ok");
+  const resB = await ns.get(idB).fetch("/");
+  t.is(await resB.text(), "ok");
+  // now flush
+  await plugin.flushAlarms(factory, [idA]);
+  // check storage "a" is 1
+  const stateA = await plugin.getObject(factory, idA);
+  t.is(await stateA.storage.get("key"), 1);
+  // check that storage "b" is still empty
+  const stateB = await plugin.getObject(factory, idB);
+  t.is(await stateB.storage.get("key"), undefined);
+  // dispose
+  plugin.dispose();
+});

packages/durable-objects/test/plugin.spec.ts

@@ -29,6 +29,9 @@ declare global {
   function getMiniflareDurableObjectStorage(
     id: DurableObjectId
   ): Promise<DurableObjectStorage>;
+  function flushMiniflareDurableObjectAlarms(
+    ids: DurableObjectId[]
+  ): Promise<void>;
 }
 
 // MiniflareCore will ensure CorePlugin is first and BindingsPlugin is last,
@@ -223,6 +226,13 @@ export default class MiniflareEnvironment implements JestEnvironment<Timer> {
       const state = await plugin.getObject(storage, id);
       return state.storage;
     };
+    global.flushMiniflareDurableObjectAlarms = async (
+      ids?: DurableObjectId[]
+    ): Promise<void> => {
+      const plugin = (await mf.getPlugins()).DurableObjectsPlugin;
+      const storage = mf.getPluginStorage("DurableObjectsPlugin");
+      return plugin.flushAlarms(storage, ids);
+    };
   }
 
   async teardown(): Promise<void> {