Add an env var to schedule the batch worker ping

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## Unreleased
+
+- Add experimental configuration options for configuring the Aggregate's batch
+  worker: `WORKER_IDLE_COOLDOWN_MS`, `WORKER_POLL_INTERVAL_MS`, and
+  `WORKER_SCHEDULE_PING`.
+
 ## 0.3.3
 
 - Bump `convex-helpers` to 0.1.128, which isn’t affected by

src/component/_generated/server.ts

@@ -41,6 +41,7 @@ type Env = {
   readonly CONVEX_SITE_URL: string;
   readonly WORKER_IDLE_COOLDOWN_MS: string | undefined;
   readonly WORKER_POLL_INTERVAL_MS: string | undefined;
+  readonly WORKER_SCHEDULE_PING: "true" | "false" | undefined;
 };
 
 /**

src/component/convex.config.ts

@@ -6,6 +6,9 @@ const component = defineComponent("aggregate", {
   env: {
     WORKER_IDLE_COOLDOWN_MS: v.optional(v.string()),
     WORKER_POLL_INTERVAL_MS: v.optional(v.string()),
+    WORKER_SCHEDULE_PING: v.optional(
+      v.union(v.literal("true"), v.literal("false")),
+    ),
   },
 });
 component.use(batchWorker);

src/component/worker.test.ts

@@ -990,6 +990,108 @@ describe("staying awake between bursts", () => {
   });
 });
 
+describe("scheduling the ping", () => {
+  afterEach(() => {
+    vi.unstubAllEnvs();
+  });
+
+  // Pending only: convex-test keeps completed runs in the table too.
+  async function pendingPings(t: TestConvex<typeof schema>) {
+    return await t.run(async (ctx) => {
+      const scheduled = await ctx.db.system
+        .query("_scheduled_functions")
+        .collect();
+      return scheduled.filter(
+        (f) => f.name.includes("pingWorker") && f.state.kind === "pending",
+      );
+    });
+  }
+
+  test("the ping is inline unless asked for", async () => {
+    const t = initConvexTest();
+    await enqueue(t, { type: "insert", key: 1, value: "a" });
+    expect(await pendingPings(t)).toHaveLength(0);
+    expect(await workerStatus(t)).toEqual({ kind: "running" });
+  });
+
+  test("a parked worker is woken from its own transaction", async () => {
+    vi.stubEnv("WORKER_SCHEDULE_PING", "true");
+    const t = initConvexTest();
+    await enqueue(t, { type: "insert", key: 1, value: "a" });
+    // The enqueue left the worker alone; the scheduled ping registers it.
+    expect(await workerStatus(t)).toBeNull();
+    expect(await pendingPings(t)).toHaveLength(1);
+
+    await drainViaWorker(t);
+    await t.run(async (ctx) => {
+      expect(await getHandler(ctx, { key: 1 })).toEqual({
+        k: 1,
+        v: "a",
+        s: 0,
+      });
+      expect(await ctx.db.query("pendingOperations").collect()).toEqual([]);
+    });
+  });
+
+  test("each transaction schedules one, and it no-ops on a running worker", async () => {
+    vi.stubEnv("WORKER_SCHEDULE_PING", "true");
+    const t = initConvexTest();
+    await enqueue(t, { type: "insert", key: 1, value: "a" });
+    vi.advanceTimersByTime(1);
+    await t.finishInProgressScheduledFunctions();
+    expect(await workerStatus(t)).toEqual({ kind: "running" });
+
+    await enqueue(t, { type: "insert", key: 2, value: "b" });
+    expect(await pendingPings(t)).toHaveLength(1);
+    await drainViaWorker(t);
+    await t.run(async (ctx) => {
+      expect(await getHandler(ctx, { key: 2 })).toEqual({
+        k: 2,
+        v: "b",
+        s: 0,
+      });
+    });
+  });
+
+  test("a transaction that enqueues repeatedly schedules one ping", async () => {
+    vi.stubEnv("WORKER_SCHEDULE_PING", "true");
+    const t = initConvexTest();
+    await enqueue(
+      t,
+      { type: "insert", key: 1, value: "a" },
+      { type: "insert", key: 2, value: "b" },
+      { type: "insert", key: 3, value: "c" },
+    );
+    expect(await pendingPings(t)).toHaveLength(1);
+
+    await drainViaWorker(t);
+    await t.run(async (ctx) => {
+      expect(await getHandler(ctx, { key: 3 })).toEqual({ k: 3, v: "c", s: 0 });
+      expect(await ctx.db.query("pendingOperations").collect()).toEqual([]);
+    });
+  });
+
+  test("work enqueued after the worker parks still drains", async () => {
+    vi.stubEnv("WORKER_SCHEDULE_PING", "true");
+    const t = initConvexTest();
+    await enqueue(t, { type: "insert", key: 1, value: "a" });
+    await drainViaWorker(t);
+    expect(await workerStatus(t)).toEqual({ kind: "idle" });
+
+    await enqueue(t, { type: "insert", key: 2, value: "b" });
+    expect(await pendingPings(t)).toHaveLength(1);
+    await drainViaWorker(t);
+    await t.run(async (ctx) => {
+      expect(await getHandler(ctx, { key: 2 })).toEqual({
+        k: 2,
+        v: "b",
+        s: 0,
+      });
+      expect(await ctx.db.query("pendingOperations").collect()).toEqual([]);
+    });
+  });
+});
+
 describe("pinging the worker", () => {
   beforeEach(() => {
     vi.mocked(ping).mockClear();

src/component/worker.ts

@@ -95,6 +95,14 @@ export async function enqueueOperations(
   if (newestEntry) {
     return;
   }
+  if (env.WORKER_SCHEDULE_PING === "true") {
+    await ctx.scheduler.runAfter(0, internal.worker.pingWorker, {});
+    return;
+  }
+  await pingBatchWorker(ctx);
+}
+
+async function pingBatchWorker(ctx: MutationCtx) {
   await ping(ctx, components.batchWorker, {
     // TODO: explore separate queues by namespace
     name: OPS_WORKER_NAME,
@@ -103,6 +111,14 @@ export async function enqueueOperations(
   });
 }
 
+export const pingWorker = internalMutation({
+  args: {},
+  returns: v.null(),
+  handler: async (ctx) => {
+    await pingBatchWorker(ctx);
+  },
+});
+
 const vBatchEntry = v.object({
   id: v.id("pendingOperations"),
   commitTs: v.int64(),