Record pendingStartId on work and upgrade batch-worker to 0.3

CHANGELOG.md

@@ -1,5 +1,17 @@
 # Changelog
 
+## 0.4.15 (Unreleased)
+
+- Upgrades @convex-dev/batch-worker to 0.3. Requires Convex 1.43 or later.
+- Records each pendingStart ID on its work document, so the next minor version
+  can report status accurately for work already in flight. **Not backwards
+  compatible due to this schema change.**
+- Adds a one-off migration that records the queue entry on work enqueued before
+  this version and removes queue entries whose work is gone. Before upgrading to
+  0.5, deploy this version and run
+  `npx convex run --component workpool migrations:backfillPendingStartIds`,
+  replacing `workpool` with each component instance's name.
+
 ## 0.4.14
 
 - Updates the /test entrypoint for compatibility with convex-test's new

package.json

@@ -72,7 +72,7 @@
     }
   },
   "peerDependencies": {
-    "convex": "^1.36.1",
+    "convex": "^1.43.0",
     "convex-helpers": "^0.1.94"
   },
   "devDependencies": {
@@ -108,6 +108,6 @@
   "types": "./dist/client/index.d.ts",
   "module": "./dist/client/index.js",
   "dependencies": {
-    "@convex-dev/batch-worker": "^0.2.2"
+    "@convex-dev/batch-worker": "^0.3.6"
   }
 }

src/component/_generated/api.ts

@@ -16,6 +16,7 @@ import type * as kick from "../kick.js";
 import type * as lib from "../lib.js";
 import type * as logging from "../logging.js";
 import type * as loop from "../loop.js";
+import type * as migrations from "../migrations.js";
 import type * as recovery from "../recovery.js";
 import type * as shared from "../shared.js";
 import type * as stats from "../stats.js";
@@ -37,6 +38,7 @@ const fullApi: ApiFromModules<{
   lib: typeof lib;
   logging: typeof logging;
   loop: typeof loop;
+  migrations: typeof migrations;
   recovery: typeof recovery;
   shared: typeof shared;
   stats: typeof stats;

src/component/lib.test.ts

@@ -129,6 +129,24 @@ describe("lib", () => {
       expect(status).toEqual({ state: "pending", previousAttempts: 0 });
     });
 
+    it("records the queue entry on the work", async () => {
+      const id = await t.mutation(api.lib.enqueue, {
+        fnHandle: "testHandle",
+        fnName: "testFunction",
+        fnArgs: {},
+        fnType: "mutation",
+        runAt: Date.now() + 60_000,
+        config: {},
+      });
+      await t.run(async (ctx) => {
+        const pending = await ctx.db.query("pendingStart").unique();
+        assert(pending);
+        expect((await ctx.db.get("work", id))?.pendingStartId).toBe(
+          pending._id,
+        );
+      });
+    });
+
     it("should throw error if maxParallelism is too high", async () => {
       await expect(
         t.mutation(api.lib.enqueue, {

src/component/lib.ts

@@ -112,10 +112,11 @@ async function enqueueHandler(
   // Store the work item
   const workId = await ctx.db.insert("work", workItem);
 
-  await ctx.db.insert("pendingStart", {
+  const pendingStartId = await ctx.db.insert("pendingStart", {
     workId,
     segment: max(toSegment(runAt), getCurrentSegment()),
   });
+  await ctx.db.patch("work", workId, { pendingStartId });
   recordEnqueued(console, { workId, fnName: workArgs.fnName, runAt });
   return workId;
 }

src/component/loop.test.ts

@@ -188,6 +188,17 @@ describe("loop", () => {
       expect(await statusOf(workId)).toMatchObject({ state: "running" });
     });
 
+    it("records the queue entry on older work when starting it", async () => {
+      await initialize();
+      const workId = await enqueueWork();
+      const [pending] = (await observe()).pendingStart;
+
+      await runLoop();
+
+      const work = await t.run((ctx) => ctx.db.get("work", workId));
+      expect(work?.pendingStartId).toBe(pending._id);
+    });
+
     it("removes work from running once a successful completion is processed", async () => {
       await initialize();
       const workId = await enqueueWork();
@@ -391,6 +402,8 @@ describe("loop", () => {
       const o = await observe();
       expect(o.pendingStart).toHaveLength(1);
       expect(o.pendingStart[0].workId).toBe(workId);
+      const work = await t.run((ctx) => ctx.db.get("work", workId));
+      expect(work?.pendingStartId).toBe(o.pendingStart[0]._id);
       expect(await statusOf(workId)).toMatchObject({
         state: "pending",
         previousAttempts: 1,

src/component/loop.ts

@@ -1,6 +1,9 @@
 import type { WithoutSystemFields } from "convex/server";
 import { type Infer, v } from "convex/values";
-import { type BatchResult, vBatchResult } from "@convex-dev/batch-worker";
+import {
+  type BatchResult,
+  defineBatchWorkerValidators,
+} from "@convex-dev/batch-worker";
 import { internal } from "./_generated/api.js";
 import { kickMainLoop } from "./kick.js";
 import type { Doc, Id } from "./_generated/dataModel.js";
@@ -115,6 +118,7 @@ const batchFields = {
   starts: v.array(vStart),
 };
 type Batch = Infer<ReturnType<typeof v.object<typeof batchFields>>>;
+const batchValidators = defineBatchWorkerValidators({ batch: batchFields });
 
 /**
  * The work query (batch-worker contract). Decides whether there's work to do,
@@ -126,8 +130,8 @@ type Batch = Infer<ReturnType<typeof v.object<typeof batchFields>>>;
  * a real dependency before going idle, so we just read the tables directly.
  */
 export const getBatch = internalQuery({
-  args: { name: v.string() },
-  returns: vBatchResult(v.object(batchFields)),
+  args: batchValidators.vQueryArgs,
+  returns: batchValidators.vQueryReturns,
   handler: async (ctx): Promise<BatchResult<Batch>> => {
     const globals = await getGlobals(ctx);
     const state = await ctx.db.query("internalState").unique();
@@ -232,7 +236,7 @@ export const getBatch = internalQuery({
  * Returning `null` tells batch-worker to re-run immediately to keep draining.
  */
 export const run = internalMutation({
-  args: batchFields,
+  args: batchValidators.vMutationArgs,
   returns: v.null(),
   handler: async (ctx, batch) => {
     const state = await getOrCreateState(ctx);
@@ -562,6 +566,9 @@ async function handleStart(
           console.error(`Trying to start, but work not found: ${workId}`);
           return null;
         }
+        if (work.pendingStartId === undefined) {
+          await ctx.db.patch("work", workId, { pendingStartId: _id });
+        }
         return {
           work,
           lagMs: Date.now() - fromSegment(segment),
@@ -698,10 +705,11 @@ async function rescheduleJob(
   const nextAttempt = withJitter(backoffMs);
   const startTime = Date.now() + nextAttempt;
   const segment = toSegment(startTime);
-  await ctx.db.insert("pendingStart", {
+  const pendingStartId = await ctx.db.insert("pendingStart", {
     workId: work._id,
     segment,
   });
+  await ctx.db.patch("work", work._id, { pendingStartId });
   return true;
 }
 

src/component/migrations.test.ts

@@ -0,0 +1,94 @@
+import { convexTest } from "convex-test";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import batchWorker from "@convex-dev/batch-worker/test";
+import { internal } from "./_generated/api.js";
+import type { Id } from "./_generated/dataModel.js";
+import schema from "./schema.js";
+import { getCurrentSegment } from "./shared.js";
+
+const modules = import.meta.glob("./**/*.ts");
+
+describe("backfillPendingStartIds", () => {
+  let t: ReturnType<typeof convexTest>;
+
+  beforeEach(() => {
+    vi.useFakeTimers();
+    t = convexTest(schema, modules);
+    batchWorker.register(t);
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  /** Insert work the way versions before the pointer did. */
+  async function insertLegacy(count = 1) {
+    return t.run(async (ctx) => {
+      const inserted = [];
+      for (let i = 0; i < count; i++) {
+        const workId = await ctx.db.insert("work", {
+          fnType: "action",
+          fnHandle: "handle",
+          fnName: "fn",
+          fnArgs: {},
+          attempts: 0,
+        });
+        const pendingId = await ctx.db.insert("pendingStart", {
+          workId,
+          segment: getCurrentSegment() + BigInt(i % 3) - 1n,
+        });
+        inserted.push({ workId, pendingId });
+      }
+      return inserted;
+    });
+  }
+
+  async function pointerOf(workId: Id<"work">) {
+    return t.run(
+      async (ctx) => (await ctx.db.get("work", workId))?.pendingStartId ?? null,
+    );
+  }
+
+  async function migrate(before?: number) {
+    await t.mutation(internal.migrations.backfillPendingStartIds, { before });
+    await t.finishAllScheduledFunctions(vi.runAllTimers);
+  }
+
+  it("records entries created before the run across pages", async () => {
+    const legacy = await insertLegacy(600);
+    // Under frozen time, convex-test spaces creation times 1µs apart.
+    vi.advanceTimersByTime(1000);
+    const before = Date.now();
+    vi.advanceTimersByTime(1000);
+    const [later] = await insertLegacy();
+
+    await migrate(before);
+
+    for (const { workId, pendingId } of legacy) {
+      expect(await pointerOf(workId)).toBe(pendingId);
+    }
+    expect(await pointerOf(later.workId)).toBeNull();
+  });
+
+  it("keeps existing pointers and deletes entries whose work is gone", async () => {
+    const [recorded, orphan] = await insertLegacy(2);
+    // Visited first, but the work already points at its own entry.
+    await t.run(async (ctx) => {
+      await ctx.db.insert("pendingStart", {
+        workId: recorded.workId,
+        segment: getCurrentSegment(),
+      });
+      await ctx.db.patch("work", recorded.workId, {
+        pendingStartId: recorded.pendingId,
+      });
+      await ctx.db.delete("work", orphan.workId);
+    });
+
+    await migrate();
+
+    expect(await pointerOf(recorded.workId)).toBe(recorded.pendingId);
+    expect(
+      await t.run((ctx) => ctx.db.get("pendingStart", orphan.pendingId)),
+    ).toBeNull();
+  });
+});