@convex-dev/workpool 0.2.0-beta.0 → 0.2.0

package/src/component/recovery.test.ts

@@ -0,0 +1,541 @@
+import { convexTest } from "convex-test";
+import {
+  describe,
+  expect,
+  it,
+  beforeEach,
+  afterEach,
+  vi,
+  assert,
+} from "vitest";
+import schema from "./schema";
+import { internal } from "./_generated/api";
+import { Id, Doc } from "./_generated/dataModel";
+import { MutationCtx } from "./_generated/server";
+import { WithoutSystemFields } from "convex/server";
+import { recover as recoverMutation } from "./recovery";
+
+// We call it directly so we can mock the system.get. A hack for now.
+// But this avoids the warning about calling directly.
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+const recover = (recoverMutation as any)._handler as typeof recoverMutation;
+
+const modules = import.meta.glob("./**/*.ts");
+
+describe("recovery", () => {
+  async function setupTest() {
+    const t = convexTest(schema, modules);
+    return t;
+  }
+
+  let t: Awaited<ReturnType<typeof setupTest>>;
+
+  // Helper function to create a work item
+  async function makeDummyWork(
+    ctx: MutationCtx,
+    overrides: Partial<WithoutSystemFields<Doc<"work">>> = {}
+  ) {
+    return ctx.db.insert("work", {
+      fnType: "action",
+      fnHandle: "test_handle",
+      fnName: "test_handle",
+      fnArgs: {},
+      attempts: 0,
+      ...overrides,
+    });
+  }
+
+  // Helper function to create a scheduled function
+  async function makeDummyScheduledFunction(
+    ctx: MutationCtx,
+    workId: Id<"work">
+  ) {
+    return ctx.scheduler.runAfter(0, internal.worker.runActionWrapper, {
+      workId,
+      fnHandle: "test_handle",
+      fnArgs: {},
+      logLevel: "WARN",
+      attempt: 0,
+    });
+  }
+
+  beforeEach(async () => {
+    vi.useFakeTimers();
+    t = await setupTest();
+
+    // Set up globals for logging
+    await t.run(async (ctx) => {
+      await ctx.db.insert("globals", {
+        maxParallelism: 10,
+        logLevel: "WARN",
+      });
+    });
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  describe("recover", () => {
+    it("should skip jobs that already have a pendingCompletion", async () => {
+      // Create work and scheduled function
+
+      const [workId, scheduledId] = await t.run(async (ctx) => {
+        const workId = await makeDummyWork(ctx);
+        const scheduledId = await makeDummyScheduledFunction(ctx, workId);
+
+        // Create a pendingCompletion for this work
+        await ctx.db.insert("pendingCompletion", {
+          segment: BigInt(1),
+          workId,
+          runResult: { kind: "failed", error: "test error" },
+          retry: true,
+        });
+
+        return [workId, scheduledId];
+      });
+
+      // Run recovery
+      await t.mutation(internal.recovery.recover, {
+        jobs: [
+          {
+            scheduledId,
+            workId,
+            attempt: 0,
+            started: Date.now(),
+          },
+        ],
+      });
+
+      // Verify no additional pendingCompletion was created
+      await t.run(async (ctx) => {
+        const pendingCompletions = await ctx.db
+          .query("pendingCompletion")
+          .withIndex("workId", (q) => q.eq("workId", workId))
+          .collect();
+        expect(pendingCompletions).toHaveLength(1);
+      });
+    });
+
+    it("should skip jobs where work is not found", async () => {
+      // Create a non-existent work ID and a valid scheduled function ID
+      const [workId, scheduledId] = await t.run(async (ctx) => {
+        // Create a temporary work ID that we'll delete
+        const workId = await makeDummyWork(ctx);
+        const scheduledId = await makeDummyScheduledFunction(ctx, workId);
+
+        // Delete the work to simulate it not being found
+        await ctx.db.delete(workId);
+
+        return [workId, scheduledId];
+      });
+
+      // Run recovery
+      await t.mutation(internal.recovery.recover, {
+        jobs: [
+          {
+            scheduledId,
+            workId,
+            attempt: 0,
+            started: Date.now(),
+          },
+        ],
+      });
+
+      // Verify no pendingCompletion was created
+      await t.run(async (ctx) => {
+        const pendingCompletions = await ctx.db
+          .query("pendingCompletion")
+          .withIndex("workId", (q) => q.eq("workId", workId))
+          .collect();
+        expect(pendingCompletions).toHaveLength(0);
+      });
+    });
+
+    it("should skip jobs where work attempts mismatch", async () => {
+      // Create work and scheduled function
+      const [workId, scheduledId] = await t.run(async (ctx) => {
+        const workId = await makeDummyWork(ctx);
+        const scheduledId = await makeDummyScheduledFunction(ctx, workId);
+
+        // Update the work to have a different attempt number
+        const work = await ctx.db.get(workId);
+        if (work) {
+          await ctx.db.patch(work._id, { attempts: 5 });
+        }
+
+        return [workId, scheduledId];
+      });
+
+      // Run recovery
+      await t.mutation(internal.recovery.recover, {
+        jobs: [
+          {
+            scheduledId,
+            workId,
+            attempt: 0, // Mismatched with the work's attempt number (5)
+            started: Date.now(),
+          },
+        ],
+      });
+
+      // Verify no pendingCompletion was created
+      await t.run(async (ctx) => {
+        const pendingCompletions = await ctx.db
+          .query("pendingCompletion")
+          .withIndex("workId", (q) => q.eq("workId", workId))
+          .collect();
+        expect(pendingCompletions).toHaveLength(0);
+      });
+    });
+
+    it("should handle scheduled job not found", async () => {
+      // Create work but use a non-existent scheduled ID
+      const [workId, scheduledId] = await t.run(async (ctx) => {
+        const workId = await makeDummyWork(ctx);
+        const scheduledId = await makeDummyScheduledFunction(ctx, workId);
+
+        return [workId, scheduledId];
+      });
+
+      // Run recovery with mocked system.get
+      await t.run(async (ctx) => {
+        // Mock the system.get to return null for our scheduledId
+        const originalGet = ctx.db.system.get;
+        ctx.db.system.get = async (id) => {
+          if (id === scheduledId) {
+            return null;
+          }
+          return await originalGet(id);
+        };
+
+        await recover(ctx, {
+          jobs: [
+            {
+              scheduledId,
+              workId,
+              attempt: 0,
+              started: Date.now(),
+            },
+          ],
+        });
+      });
+
+      // Verify pendingCompletion was created with failure
+      await t.run(async (ctx) => {
+        const pendingCompletions = await ctx.db
+          .query("pendingCompletion")
+          .withIndex("workId", (q) => q.eq("workId", workId))
+          .collect();
+        expect(pendingCompletions).toHaveLength(1);
+        expect(pendingCompletions[0].runResult.kind).toBe("failed");
+        assert(pendingCompletions[0].runResult.kind === "failed");
+        expect(pendingCompletions[0].runResult.error).toContain(
+          "Scheduled job not found"
+        );
+      });
+    });
+
+    it("should handle failed scheduled jobs", async () => {
+      // Create work and scheduled function
+
+      const [workId, scheduledId] = await t.run(async (ctx) => {
+        const workId = await makeDummyWork(ctx);
+        const scheduledId = await makeDummyScheduledFunction(ctx, workId);
+
+        return [workId, scheduledId];
+      });
+
+      // Run recovery with mocked failed state
+      await t.run(async (ctx) => {
+        // Mock the system.get to return a failed state
+        const originalGet = ctx.db.system.get;
+        ctx.db.system.get = async (id) => {
+          if (id === scheduledId) {
+            return {
+              _id: scheduledId,
+              _creationTime: Date.now(),
+              name: "internal/worker.runActionWrapper",
+              args: [
+                {
+                  workId,
+                  fnHandle: "test_handle",
+                  fnArgs: {},
+                  logLevel: "WARN",
+                  attempt: 0,
+                },
+              ],
+              scheduledTime: Date.now(),
+              state: {
+                kind: "failed",
+                error: "Function execution failed",
+              },
+            };
+          }
+          return await originalGet(id);
+        };
+
+        await recover(ctx, {
+          jobs: [
+            {
+              scheduledId,
+              workId,
+              attempt: 0,
+              started: Date.now(),
+            },
+          ],
+        });
+      });
+
+      // Verify pendingCompletion was created with the same failure
+      await t.run(async (ctx) => {
+        const pendingCompletions = await ctx.db
+          .query("pendingCompletion")
+          .withIndex("workId", (q) => q.eq("workId", workId))
+          .collect();
+        expect(pendingCompletions).toHaveLength(1);
+        expect(pendingCompletions[0].runResult.kind).toBe("failed");
+        assert(pendingCompletions[0].runResult.kind === "failed");
+        expect(pendingCompletions[0].runResult.error).toBe(
+          "Function execution failed"
+        );
+      });
+    });
+
+    it("should handle canceled scheduled jobs", async () => {
+      // Create work and scheduled function
+      let workId: Id<"work">;
+      let scheduledId: Id<"_scheduled_functions">;
+
+      await t.run(async (ctx) => {
+        workId = await makeDummyWork(ctx);
+        scheduledId = await makeDummyScheduledFunction(ctx, workId);
+      });
+
+      // Run recovery with mocked system.get
+      await t.run(async (ctx) => {
+        // Mock the system.get to return a canceled state
+        const originalGet = ctx.db.system.get;
+        ctx.db.system.get = async (id) => {
+          if (id === scheduledId) {
+            return {
+              _id: scheduledId,
+              _creationTime: Date.now(),
+              name: "internal/worker.runActionWrapper",
+              args: [
+                {
+                  workId,
+                  fnHandle: "test_handle",
+                  fnArgs: {},
+                  logLevel: "WARN",
+                  attempt: 0,
+                },
+              ],
+              scheduledTime: Date.now(),
+              state: {
+                kind: "canceled",
+              },
+            };
+          }
+          return await originalGet(id);
+        };
+
+        await recover(ctx, {
+          jobs: [
+            {
+              scheduledId,
+              workId,
+              attempt: 0,
+              started: Date.now(),
+            },
+          ],
+        });
+      });
+
+      // Verify pendingCompletion was created with failure due to cancelation
+      await t.run(async (ctx) => {
+        const pendingCompletions = await ctx.db
+          .query("pendingCompletion")
+          .withIndex("workId", (q) => q.eq("workId", workId))
+          .collect();
+        expect(pendingCompletions).toHaveLength(1);
+        expect(pendingCompletions[0].runResult.kind).toBe("failed");
+        assert(pendingCompletions[0].runResult.kind === "failed");
+        expect(pendingCompletions[0].runResult.error).toBe(
+          "Canceled via scheduler"
+        );
+      });
+    });
+
+    it("should handle multiple jobs in a single call", async () => {
+      // Create multiple work items and scheduled functions
+      let workId1: Id<"work">;
+      let workId2: Id<"work">;
+      let scheduledId1: Id<"_scheduled_functions">;
+      let scheduledId2: Id<"_scheduled_functions">;
+
+      await t.run(async (ctx) => {
+        workId1 = await makeDummyWork(ctx, { fnArgs: { test: 1 } });
+        workId2 = await makeDummyWork(ctx, { fnArgs: { test: 2 } });
+        scheduledId1 = await makeDummyScheduledFunction(ctx, workId1);
+        scheduledId2 = await makeDummyScheduledFunction(ctx, workId2);
+      });
+
+      // Run recovery with mocked system.get
+      await t.run(async (ctx) => {
+        // Mock the system.get to return different states for each scheduled function
+        const originalGet = ctx.db.system.get;
+        ctx.db.system.get = async (id) => {
+          if (id === scheduledId1) {
+            return {
+              _id: scheduledId1,
+              _creationTime: Date.now(),
+              name: "internal/worker.runActionWrapper",
+              args: [
+                {
+                  workId: workId1,
+                  fnHandle: "test_handle",
+                  fnArgs: { test: 1 },
+                  logLevel: "WARN",
+                  attempt: 0,
+                },
+              ],
+              scheduledTime: Date.now(),
+              state: {
+                kind: "failed",
+                error: "Function 1 failed",
+              },
+            };
+          } else if (id === scheduledId2) {
+            return {
+              _id: scheduledId2,
+              _creationTime: Date.now(),
+              name: "internal/worker.runActionWrapper",
+              args: [
+                {
+                  workId: workId2,
+                  fnHandle: "test_handle",
+                  fnArgs: { test: 2 },
+                  logLevel: "WARN",
+                  attempt: 0,
+                },
+              ],
+              scheduledTime: Date.now(),
+              state: {
+                kind: "canceled",
+              },
+            };
+          }
+          return await originalGet(id);
+        };
+
+        await recover(ctx, {
+          jobs: [
+            {
+              scheduledId: scheduledId1,
+              workId: workId1,
+              attempt: 0,
+              started: Date.now(),
+            },
+            {
+              scheduledId: scheduledId2,
+              workId: workId2,
+              attempt: 0,
+              started: Date.now(),
+            },
+          ],
+        });
+      });
+
+      // Verify both jobs were processed correctly
+      await t.run(async (ctx) => {
+        const pendingCompletions = await ctx.db
+          .query("pendingCompletion")
+          .collect();
+        expect(pendingCompletions).toHaveLength(2);
+
+        // Find completions for each work ID
+        const completion1 = pendingCompletions.find(
+          (pc) => pc.workId === workId1
+        );
+        const completion2 = pendingCompletions.find(
+          (pc) => pc.workId === workId2
+        );
+
+        expect(completion1).toBeDefined();
+        expect(completion2).toBeDefined();
+
+        if (completion1) {
+          expect(completion1.runResult.kind).toBe("failed");
+          assert(completion1.runResult.kind === "failed");
+          expect(completion1.runResult.error).toBe("Function 1 failed");
+        }
+
+        if (completion2) {
+          expect(completion2.runResult.kind).toBe("failed");
+          assert(completion2.runResult.kind === "failed");
+          expect(completion2.runResult.error).toBe("Canceled via scheduler");
+        }
+      });
+    });
+
+    it("should not process jobs with other scheduled states", async () => {
+      // Create work and scheduled function
+      const [workId, scheduledId] = await t.run(async (ctx) => {
+        const workId = await makeDummyWork(ctx);
+        const scheduledId = await makeDummyScheduledFunction(ctx, workId);
+
+        return [workId, scheduledId];
+      });
+
+      // Run recovery with mocked system.get
+      await t.run(async (ctx) => {
+        // Mock the system.get to return a pending state
+        const originalGet = ctx.db.system.get;
+        ctx.db.system.get = async (id) => {
+          if (id === scheduledId) {
+            return {
+              _id: scheduledId,
+              _creationTime: Date.now(),
+              name: "internal/worker.runActionWrapper",
+              args: [
+                {
+                  workId,
+                  fnHandle: "test_handle",
+                  fnArgs: {},
+                  logLevel: "WARN",
+                  attempt: 0,
+                },
+              ],
+              scheduledTime: Date.now(),
+              state: {
+                kind: "pending",
+              },
+            };
+          }
+          return await originalGet(id);
+        };
+
+        await recover(ctx, {
+          jobs: [
+            {
+              scheduledId,
+              workId,
+              attempt: 0,
+              started: Date.now(),
+            },
+          ],
+        });
+      });
+
+      // Verify no pendingCompletion was created
+      await t.run(async (ctx) => {
+        const pendingCompletions = await ctx.db
+          .query("pendingCompletion")
+          .withIndex("workId", (q) => q.eq("workId", workId))
+          .collect();
+        expect(pendingCompletions).toHaveLength(0);
+      });
+    });
+  });
+});

package/src/component/recovery.ts

@@ -1,79 +1,96 @@
-import { Id } from "./_generated/dataModel.js";
+import { Infer, v } from "convex/values";
 import { internalMutation } from "./_generated/server.js";
-import { kickMainLoop } from "./kick.js";
+import { completeArgs, completeHandler } from "./complete.js";
 import { createLogger } from "./logging.js";
-import schema from "./schema.js";
-import { RunResult, nextSegment } from "./shared.js";
 
+/**
+ * This can run when things fail because of server failures / restarts, or when
+ * the user cancels scheduled jobs (from the dashboard).
+ * Possible states it could be in at the moment this executes:
+ * - in internalState.running and complete was never called
+ *   -> we should call completeHandler with failure.
+ * - complete already called, no action needed (only possible for actions):
+ *  - In pendingCompletion still and internalState.running.
+ *    -> check for pendingCompletion.
+ *  - pendingCompletion already processed.
+ *   - No retry: work was deleted, not in internalState.running.
+ *     -> check for work.
+ *   - Retry: attempts will mismatch
+ *     -> check work.attempts
+ */
 export const recover = internalMutation({
   args: {
-    jobs: schema.tables.internalState.validator.fields.running,
+    jobs: v.array(
+      v.object({
+        scheduledId: v.id("_scheduled_functions"),
+        workId: v.id("work"),
+        attempt: v.number(),
+        started: v.number(),
+      })
+    ),
   },
   handler: async (ctx, { jobs }) => {
     const globals = await ctx.db.query("globals").unique();
     const console = createLogger(globals?.logLevel);
-    const completed: { workId: Id<"work">; runResult: RunResult }[] = [];
-    let didAnything = false;
-    const segment = nextSegment();
-    await Promise.all(
-      jobs.map(async (job) => {
-        const scheduled = await ctx.db.system.get(job.scheduledId);
-        const preamble = `[recovery] Scheduled job ${job.scheduledId} for work ${job.workId}`;
-        if (scheduled === null) {
-          console.warn(`${preamble} not found`);
-          completed.push({
+    const toComplete: Infer<typeof completeArgs.fields.jobs> = [];
+    for (let i = 0; i < jobs.length; i++) {
+      const job = jobs[i];
+      const preamble = `[recovery] Scheduled job ${job.scheduledId} for work ${job.workId}`;
+      const pendingCompletion = await ctx.db
+        .query("pendingCompletion")
+        .withIndex("workId", (q) => q.eq("workId", job.workId))
+        .first();
+      if (pendingCompletion) {
+        // Completion already pending, no need to do anything.
+        console.debug(`${preamble} already in pendingCompletion, skipping`);
+        continue;
+      }
+      const work = await ctx.db.get(job.workId);
+      if (work === null) {
+        // Completion already executed w/o retries, no need to do anything.
+        console.warn(`${preamble} work not found, skipping`);
+        continue;
+      }
+      if (work.attempts !== job.attempt) {
+        // Retry already started, no need to do anything.
+        console.warn(`${preamble} attempts mismatch, skipping`);
+        continue;
+      }
+      const scheduled = await ctx.db.system.get(job.scheduledId);
+      if (scheduled === null) {
+        console.warn(`${preamble} not found in _scheduled_functions`);
+        toComplete.push({
+          workId: job.workId,
+          runResult: { kind: "failed", error: `Scheduled job not found` },
+          attempt: job.attempt,
+        });
+        continue;
+      }
+      // This will find everything that timed out, failed ungracefully, was
+      // canceled, or succeeded without a return value.
+      switch (scheduled.state.kind) {
+        case "failed": {
+          console.debug(`${preamble} failed and detected in recovery`);
+          toComplete.push({
             workId: job.workId,
-            runResult: { kind: "failed", error: `Scheduled job not found` },
+            runResult: scheduled.state,
+            attempt: job.attempt,
           });
-          return;
+          break;
         }
-        // This will find everything that timed out, failed ungracefully, was
-        // canceled, or succeeded without a return value.
-        switch (scheduled.state.kind) {
-          case "failed": {
-            console.debug(`${preamble} failed and detected in recovery`);
-            const pendingCompletion = await ctx.db
-              .query("pendingCompletion")
-              .withIndex("workId", (q) => q.eq("workId", job.workId))
-              .first();
-            if (pendingCompletion) {
-              console.debug(
-                `${preamble} already in pendingCompletion, not reporting`
-              );
-            } else {
-              await ctx.db.insert("pendingCompletion", {
-                runResult: scheduled.state,
-                workId: job.workId,
-                segment,
-              });
-              didAnything = true;
-            }
-            break;
-          }
-          case "canceled": {
-            console.debug(`${preamble} was canceled and detected in recovery`);
-            const pendingCancelation = await ctx.db
-              .query("pendingCancelation")
-              .withIndex("workId", (q) => q.eq("workId", job.workId))
-              .first();
-            if (pendingCancelation) {
-              console.debug(
-                `${preamble} already in pendingCancelation, not reporting`
-              );
-            } else {
-              await ctx.db.insert("pendingCancelation", {
-                workId: job.workId,
-                segment,
-              });
-              didAnything = true;
-            }
-            break;
-          }
+        case "canceled": {
+          console.debug(`${preamble} was canceled and detected in recovery`);
+          toComplete.push({
+            workId: job.workId,
+            runResult: { kind: "failed", error: "Canceled via scheduler" },
+            attempt: job.attempt,
+          });
+          break;
         }
-      })
-    );
-    if (didAnything) {
-      await kickMainLoop(ctx, "recovery");
+      }
+    }
+    if (toComplete.length > 0) {
+      await completeHandler(ctx, { jobs: toComplete });
     }
   },
 });