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

package/src/component/loop.test.ts

@@ -1,21 +1,24 @@
 import { convexTest } from "convex-test";
+import { WithoutSystemFields } from "convex/server";
 import {
+  afterEach,
+  assert,
+  beforeEach,
   describe,
   expect,
   it,
-  beforeEach,
-  afterEach,
   vi,
-  assert,
 } from "vitest";
+import { api, internal } from "./_generated/api";
 import { Doc, Id } from "./_generated/dataModel";
-import schema from "./schema";
-import { internal } from "./_generated/api";
-import { currentSegment, nextSegment, toSegment } from "./shared";
-import { api } from "./_generated/api";
-import { DEFAULT_MAX_PARALLELISM } from "./kick";
-import { WithoutSystemFields } from "convex/server";
 import { MutationCtx } from "./_generated/server";
+import schema from "./schema";
+import {
+  currentSegment,
+  DEFAULT_MAX_PARALLELISM,
+  nextSegment,
+  toSegment,
+} from "./shared";
 
 const modules = import.meta.glob("./**/*.ts");
 
@@ -57,17 +60,38 @@ describe("loop", () => {
       workId,
       fnHandle: "test_handle",
       fnArgs: {},
-      logLevel: "DEBUG",
+      logLevel: "WARN",
       attempt: 0,
     });
   }
 
+  async function insertInternalState(
+    ctx: MutationCtx,
+    overrides: Partial<WithoutSystemFields<Doc<"internalState">>> = {}
+  ) {
+    await ctx.db.insert("internalState", {
+      generation: 1n,
+      segmentCursors: { incoming: 0n, completion: 0n, cancelation: 0n },
+      lastRecovery: currentSegment(),
+      report: {
+        completed: 0,
+        succeeded: 0,
+        failed: 0,
+        retries: 0,
+        canceled: 0,
+        lastReportTs: Date.now(),
+      },
+      running: [],
+      ...overrides,
+    });
+  }
+
   beforeEach(async () => {
     vi.useFakeTimers();
     t = await setupTest();
     await t.run(async (ctx) => {
       await ctx.db.insert("globals", {
-        logLevel: "DEBUG",
+        logLevel: "WARN",
         maxParallelism: DEFAULT_MAX_PARALLELISM,
       });
     });
@@ -82,20 +106,7 @@ describe("loop", () => {
       // Setup initial state
       const workId = await t.run<Id<"work">>(async (ctx) => {
         // Create internal state
-        await ctx.db.insert("internalState", {
-          generation: 1n,
-          segmentCursors: { incoming: 0n, completion: 0n, cancelation: 0n },
-          lastRecovery: 0n,
-          report: {
-            completed: 0,
-            succeeded: 0,
-            failed: 0,
-            retries: 0,
-            canceled: 0,
-            lastReportTs: Date.now(),
-          },
-          running: [],
-        });
+        await insertInternalState(ctx);
 
         // Create running runStatus
         await ctx.db.insert("runStatus", {
@@ -158,20 +169,7 @@ describe("loop", () => {
       // Setup initial state
       const workId = await t.run<Id<"work">>(async (ctx) => {
         // Create internal state
-        await ctx.db.insert("internalState", {
-          generation: 1n,
-          segmentCursors: { incoming: 0n, completion: 0n, cancelation: 0n },
-          lastRecovery: 0n,
-          report: {
-            completed: 0,
-            succeeded: 0,
-            failed: 0,
-            retries: 0,
-            canceled: 0,
-            lastReportTs: Date.now(),
-          },
-          running: [],
-        });
+        await insertInternalState(ctx);
 
         // Create running runStatus
         await ctx.db.insert("runStatus", {
@@ -228,20 +226,7 @@ describe("loop", () => {
       // Setup initial state with a running job that will need retry
       const workId = await t.run<Id<"work">>(async (ctx) => {
         // Create internal state
-        await ctx.db.insert("internalState", {
-          generation: 1n,
-          segmentCursors: { incoming: 0n, completion: 0n, cancelation: 0n },
-          lastRecovery: 0n,
-          report: {
-            completed: 0,
-            succeeded: 0,
-            failed: 0,
-            retries: 0,
-            canceled: 0,
-            lastReportTs: Date.now(),
-          },
-          running: [],
-        });
+        await insertInternalState(ctx);
 
         // Create running runStatus
         await ctx.db.insert("runStatus", {
@@ -325,20 +310,7 @@ describe("loop", () => {
       // Setup initial idle state
       await t.run(async (ctx) => {
         // Create internal state
-        await ctx.db.insert("internalState", {
-          generation: 1n,
-          segmentCursors: { incoming: 0n, completion: 0n, cancelation: 0n },
-          lastRecovery: 0n,
-          report: {
-            completed: 0,
-            succeeded: 0,
-            failed: 0,
-            retries: 0,
-            canceled: 0,
-            lastReportTs: Date.now(),
-          },
-          running: [],
-        });
+        await insertInternalState(ctx);
 
         // Create idle runStatus
         await ctx.db.insert("runStatus", {
@@ -372,20 +344,7 @@ describe("loop", () => {
       // Setup initial running state with work
       await t.run(async (ctx) => {
         // Create internal state
-        await ctx.db.insert("internalState", {
-          generation: 1n,
-          segmentCursors: { incoming: 0n, completion: 0n, cancelation: 0n },
-          lastRecovery: 0n,
-          report: {
-            completed: 0,
-            succeeded: 0,
-            failed: 0,
-            retries: 0,
-            canceled: 0,
-            lastReportTs: Date.now(),
-          },
-          running: [],
-        });
+        await insertInternalState(ctx);
 
         // Create running runStatus
         await ctx.db.insert("runStatus", {
@@ -428,6 +387,7 @@ describe("loop", () => {
     it("should transition from running to saturated when maxed out", async () => {
       // Setup initial running state with max capacity
       await setMaxParallelism(1);
+      const segment = currentSegment();
       await t.run(async (ctx) => {
         // Create work item
         const workId = await makeDummyWork(ctx);
@@ -436,18 +396,7 @@ describe("loop", () => {
         const scheduledId = await makeDummyScheduledFunction(ctx, workId);
 
         // Create internal state with running job
-        await ctx.db.insert("internalState", {
-          generation: 1n,
-          segmentCursors: { incoming: 0n, completion: 0n, cancelation: 0n },
-          lastRecovery: 0n,
-          report: {
-            completed: 0,
-            succeeded: 0,
-            failed: 0,
-            retries: 0,
-            canceled: 0,
-            lastReportTs: Date.now(),
-          },
+        await insertInternalState(ctx, {
           running: [{ workId, scheduledId, started: Date.now() }],
         });
 
@@ -461,14 +410,14 @@ describe("loop", () => {
 
         await ctx.db.insert("pendingStart", {
           workId: anotherWorkId,
-          segment: 1n,
+          segment,
         });
       });
 
       // Run updateRunStatus to transition to scheduled with saturated=true
       await t.mutation(internal.loop.updateRunStatus, {
         generation: 1n,
-        segment: 1n,
+        segment,
       });
 
       // Verify state transition to scheduled with saturated=true
@@ -484,45 +433,30 @@ describe("loop", () => {
 
     it("should transition from scheduled to running when new work is enqueued", async () => {
       // Setup initial scheduled state
-      const scheduledId = await t.run<Id<"_scheduled_functions">>(
-        async (ctx) => {
-          // Create internal state
-          await ctx.db.insert("internalState", {
+      await t.run<Id<"_scheduled_functions">>(async (ctx) => {
+        // Create internal state
+        await insertInternalState(ctx);
+
+        // Schedule main loop
+        const scheduledId = await ctx.scheduler.runAfter(
+          1000,
+          internal.loop.main,
+          { generation: 1n, segment: nextSegment() + 10n }
+        );
+
+        // Create scheduled runStatus
+        await ctx.db.insert("runStatus", {
+          state: {
+            kind: "scheduled",
+            segment: nextSegment() + 10n,
+            scheduledId,
+            saturated: false,
             generation: 1n,
-            segmentCursors: { incoming: 0n, completion: 0n, cancelation: 0n },
-            lastRecovery: 0n,
-            report: {
-              completed: 0,
-              succeeded: 0,
-              failed: 0,
-              retries: 0,
-              canceled: 0,
-              lastReportTs: Date.now(),
-            },
-            running: [],
-          });
-
-          // Schedule main loop
-          const scheduledId = await ctx.scheduler.runAfter(
-            1000,
-            internal.loop.main,
-            { generation: 1n, segment: nextSegment() + 10n }
-          );
-
-          // Create scheduled runStatus
-          await ctx.db.insert("runStatus", {
-            state: {
-              kind: "scheduled",
-              segment: nextSegment() + 10n,
-              scheduledId,
-              saturated: false,
-              generation: 1n,
-            },
-          });
+          },
+        });
 
-          return scheduledId;
-        }
-      );
+        return scheduledId;
+      });
 
       // Enqueue work to trigger transition to running
       await t.mutation(api.lib.enqueue, {
@@ -551,20 +485,7 @@ describe("loop", () => {
       // Setup initial running state with work
       const workId = await t.run<Id<"work">>(async (ctx) => {
         // Create internal state
-        await ctx.db.insert("internalState", {
-          generation: 1n,
-          segmentCursors: { incoming: 0n, completion: 0n, cancelation: 0n },
-          lastRecovery: 0n,
-          report: {
-            completed: 0,
-            succeeded: 0,
-            failed: 0,
-            retries: 0,
-            canceled: 0,
-            lastReportTs: Date.now(),
-          },
-          running: [],
-        });
+        await insertInternalState(ctx);
 
         // Create running runStatus
         await ctx.db.insert("runStatus", {
@@ -615,6 +536,38 @@ describe("loop", () => {
         assert(runStatus.state.kind === "idle");
       });
     });
+    it("should transition from scheduled to running when main loop runs", async () => {
+      const segment = nextSegment();
+      await t.run(async (ctx) => {
+        await insertInternalState(ctx);
+
+        const scheduledId = await ctx.scheduler.runAfter(
+          1000,
+          internal.loop.main,
+          { generation: 1n, segment }
+        );
+
+        await ctx.db.insert("runStatus", {
+          state: {
+            kind: "scheduled",
+            scheduledId,
+            generation: 1n,
+            segment,
+            saturated: false,
+          },
+        });
+      });
+      // Run main loop
+      await t.mutation(internal.loop.main, { generation: 1n, segment });
+
+      // Verify state transition to running
+      await t.run(async (ctx) => {
+        const runStatus = await ctx.db.query("runStatus").unique();
+        expect(runStatus).toBeDefined();
+        assert(runStatus);
+        expect(runStatus.state.kind).toBe("running");
+      });
+    });
   });
 
   describe("main function", () => {
@@ -1106,7 +1059,7 @@ describe("loop", () => {
         await ctx.db.insert("internalState", {
           generation: 1n,
           segmentCursors: { incoming: 0n, completion: 0n, cancelation: 0n },
-          lastRecovery: 0n,
+          lastRecovery: now,
           report: {
             completed: 0,
             succeeded: 0,
@@ -1156,6 +1109,7 @@ describe("loop", () => {
         const workId = await makeDummyWork(ctx, {
           attempts: 0,
           onComplete: {
+            // TODO: make this a real handle
             fnHandle: "onComplete_handle",
             context: { data: "test" },
           },

package/src/component/loop.ts

@@ -3,7 +3,7 @@ import { v } from "convex/values";
 import { internal } from "./_generated/api.js";
 import { Doc, Id } from "./_generated/dataModel.js";
 import { internalMutation, MutationCtx } from "./_generated/server.js";
-import { DEFAULT_MAX_PARALLELISM } from "./kick.js";
+import type { CompleteJob } from "./complete.js";
 import {
   createLogger,
   DEFAULT_LOG_LEVEL,
@@ -14,6 +14,7 @@ import {
   boundScheduledTime,
   Config,
   currentSegment,
+  DEFAULT_MAX_PARALLELISM,
   fromSegment,
   max,
   nextSegment,
@@ -21,7 +22,6 @@ import {
   toSegment,
 } from "./shared.js";
 import { recordCompleted, recordReport, recordStarted } from "./stats.js";
-import type { CompleteJob } from "./complete.js";
 
 const CANCELLATION_BATCH_SIZE = 64; // the only queue that can get unbounded.
 const SECOND = 1000;
@@ -47,41 +47,49 @@ export const INITIAL_STATE: WithoutSystemFields<Doc<"internalState">> = {
 // There should only ever be at most one of these scheduled or running.
 export const main = internalMutation({
   args: { generation: v.int64(), segment: v.int64() },
-  handler: async (ctx, args) => {
+  handler: async (ctx, { generation, segment }) => {
     // State will be modified and patched at the end of the function.
     const state = await getOrCreateState(ctx);
-    if (args.generation !== state.generation) {
+    if (generation !== state.generation) {
       throw new Error(
-        `generation mismatch: ${args.generation} !== ${state.generation}`
+        `generation mismatch: ${generation} !== ${state.generation}`
       );
     }
     state.generation++;
+    const runStatus = await getOrCreateRunningStatus(ctx);
+    if (runStatus.state.kind !== "running") {
+      await ctx.db.patch(runStatus._id, {
+        state: { kind: "running" },
+      });
+    }
 
     const globals = await getGlobals(ctx);
     const console = createLogger(globals.logLevel);
+    const delayMs = Date.now() - fromSegment(segment);
+    console.debug(`[main] generation ${generation} behind: ${delayMs}ms`);
 
     // Read pendingCompletions, including retry handling.
     console.time("[main] pendingCompletion");
-    const toCancel = await handleCompletions(ctx, state, args.segment, console);
+    const toCancel = await handleCompletions(ctx, state, segment, console);
     console.timeEnd("[main] pendingCompletion");
 
     // Read pendingCancelation, deleting from pendingStart. If it's still running, queue to cancel.
     console.time("[main] pendingCancelation");
-    await handleCancelation(ctx, state, args.segment, console, toCancel);
+    await handleCancelation(ctx, state, segment, console, toCancel);
     console.timeEnd("[main] pendingCancelation");
 
     if (state.running.length === 0) {
       // If there's nothing active, reset lastRecovery.
-      state.lastRecovery = args.segment;
-    } else if (args.segment - state.lastRecovery >= RECOVERY_PERIOD_SEGMENTS) {
+      state.lastRecovery = segment;
+    } else if (segment - state.lastRecovery >= RECOVERY_PERIOD_SEGMENTS) {
       // Otherwise schedule recovery for any old jobs.
       await handleRecovery(ctx, state, console);
-      state.lastRecovery = args.segment;
+      state.lastRecovery = segment;
     }
 
     // Read pendingStart up to max capacity. Update the config, and incomingSegmentCursor.
     console.time("[main] pendingStart");
-    await handleStart(ctx, state, args.segment, console, globals);
+    await handleStart(ctx, state, segment, console, globals);
     console.timeEnd("[main] pendingStart");
 
     if (Date.now() - state.report.lastReportTs >= MINUTE) {
@@ -92,7 +100,7 @@ export const main = internalMutation({
         // It's been a while, let's start fresh.
         lastReportTs = Date.now();
       }
-      console.info(recordReport(state));
+      recordReport(console, state);
       state.report = {
         completed: 0,
         succeeded: 0,
@@ -106,41 +114,42 @@ export const main = internalMutation({
     await ctx.db.replace(state._id, state);
     await ctx.scheduler.runAfter(0, internal.loop.updateRunStatus, {
       generation: state.generation,
-      segment: args.segment,
+      segment,
     });
+    // TODO: if there were more cancellations, schedule main directly.
   },
 });
 
 export const updateRunStatus = internalMutation({
   args: { generation: v.int64(), segment: v.int64() },
-  handler: async (ctx, args) => {
+  handler: async (ctx, { generation, segment }) => {
     const globals = await getGlobals(ctx);
     const console = createLogger(globals.logLevel);
     const maxParallelism = globals.maxParallelism;
     const state = await getOrCreateState(ctx);
-    if (args.generation !== state.generation) {
+    if (generation !== state.generation) {
       throw new Error(
-        `generation mismatch: ${args.generation} !== ${state.generation}`
+        `generation mismatch: ${generation} !== ${state.generation}`
       );
     }
 
     console.time("[updateRunStatus] outstandingCancelations");
     const outstandingCancelations = await getNextUp(ctx, "pendingCancelation", {
       start: state.segmentCursors.cancelation,
-      end: args.segment,
+      end: segment,
     });
     console.timeEnd("[updateRunStatus] outstandingCancelations");
     if (outstandingCancelations) {
       await ctx.scheduler.runAfter(0, internal.loop.main, {
-        generation: args.generation,
-        segment: args.segment,
+        generation,
+        segment,
       });
       return;
     }
 
     // TODO: check for current segment (or from args) first, to avoid OCCs.
     console.time("[updateRunStatus] nextSegmentIsActionable");
-    const next = max(args.segment + 1n, currentSegment());
+    const next = max(segment + 1n, currentSegment());
     const nextIsActionable = await nextSegmentIsActionable(
       ctx,
       state,
@@ -154,7 +163,7 @@ export const updateRunStatus = internalMutation({
         boundScheduledTime(fromSegment(next), console),
         internal.loop.main,
         {
-          generation: args.generation,
+          generation,
           segment: next,
         }
       );
@@ -177,7 +186,7 @@ export const updateRunStatus = internalMutation({
         },
       });
       await ctx.scheduler.runAfter(0, internal.loop.main, {
-        generation: args.generation,
+        generation,
         segment: currentSegment(),
       });
       return;
@@ -199,29 +208,29 @@ export const updateRunStatus = internalMutation({
       )
     );
     console.timeEnd("[updateRunStatus] findNextSegment");
-    let segment = docs.map((d) => d?.segment).sort()[0];
+    let targetSegment = docs.map((d) => d?.segment).sort()[0];
     const runStatus = await getOrCreateRunningStatus(ctx);
     const saturated = state.running.length >= maxParallelism;
-    if (segment !== undefined || state.running.length > 0) {
+    if (targetSegment !== undefined || state.running.length > 0) {
       // If there's something to do, schedule for next actionable segment.
       // Or the next recovery, whichever comes first.
       const nextRecoverySegment = state.lastRecovery + RECOVERY_PERIOD_SEGMENTS;
-      if (!segment || segment > nextRecoverySegment) {
-        segment = nextRecoverySegment;
+      if (!targetSegment || targetSegment > nextRecoverySegment) {
+        targetSegment = nextRecoverySegment;
       }
       const scheduledId = await ctx.scheduler.runAt(
-        boundScheduledTime(fromSegment(segment), console),
+        boundScheduledTime(fromSegment(targetSegment), console),
         internal.loop.main,
-        { generation: args.generation, segment }
+        { generation, segment: targetSegment }
       );
-      if (segment > nextSegment()) {
+      if (targetSegment > nextSegment()) {
         await ctx.db.patch(runStatus._id, {
           state: {
             kind: "scheduled",
             scheduledId,
             saturated,
-            generation: args.generation,
-            segment,
+            generation,
+            segment: targetSegment,
           },
         });
       } else {
@@ -233,7 +242,7 @@ export const updateRunStatus = internalMutation({
     }
     // There seems to be nothing in the future to do, so go idle.
     await ctx.db.patch(runStatus._id, {
-      state: { kind: "idle", generation: args.generation },
+      state: { kind: "idle", generation },
     });
   },
 });
@@ -370,7 +379,7 @@ async function handleCompletions(
         const retried = await rescheduleJob(ctx, work, console);
         if (retried) {
           state.report.retries++;
-          console.info(recordCompleted(work, "retrying"));
+          recordCompleted(console, work, "retrying");
         } else {
           // We don't retry if it's been canceled in the mean time.
           state.report.canceled++;
@@ -492,9 +501,8 @@ async function handleStart(
   state: Doc<"internalState">,
   segment: bigint,
   console: Logger,
-  globals: Config
+  { maxParallelism, logLevel }: Config
 ) {
-  const maxParallelism = globals.maxParallelism;
   // Schedule as many as needed to reach maxParallelism.
   const toSchedule = maxParallelism - state.running.length;
 
@@ -513,12 +521,13 @@ async function handleStart(
   state.running.push(
     ...(
       await Promise.all(
-        pending.map(async ({ _id, workId }) => {
+        pending.map(async ({ _id, workId, segment }) => {
           if (state.running.some((r) => r.workId === workId)) {
             console.error(`[main] ${workId} already running (skipping start)`);
             return null;
           }
-          const scheduledId = await beginWork(ctx, workId, globals.logLevel);
+          const lagMs = Date.now() - fromSegment(segment);
+          const scheduledId = await beginWork(ctx, workId, logLevel, lagMs);
           await ctx.db.delete(_id);
           return { scheduledId, workId, started: Date.now() };
         })
@@ -530,14 +539,15 @@ async function handleStart(
 async function beginWork(
   ctx: MutationCtx,
   workId: Id<"work">,
-  logLevel: LogLevel
+  logLevel: LogLevel,
+  lagMs: number
 ): Promise<Id<"_scheduled_functions">> {
   const console = createLogger(logLevel);
   const work = await ctx.db.get(workId);
   if (!work) {
     throw new Error("work not found");
   }
-  console.info(recordStarted(work));
+  recordStarted(console, work, lagMs);
   const { attempts: attempt, fnHandle, fnArgs } = work;
   const args = { workId, fnHandle, fnArgs, logLevel, attempt };
   if (work.fnType === "action") {

package/src/component/recovery.test.ts

@@ -54,7 +54,7 @@ describe("recovery", () => {
       workId,
       fnHandle: "test_handle",
       fnArgs: {},
-      logLevel: "DEBUG",
+      logLevel: "WARN",
       attempt: 0,
     });
   }
@@ -67,7 +67,7 @@ describe("recovery", () => {
     await t.run(async (ctx) => {
       await ctx.db.insert("globals", {
         maxParallelism: 10,
-        logLevel: "INFO",
+        logLevel: "WARN",
       });
     });
   });
@@ -261,7 +261,7 @@ describe("recovery", () => {
                   workId,
                   fnHandle: "test_handle",
                   fnArgs: {},
-                  logLevel: "DEBUG",
+                  logLevel: "WARN",
                   attempt: 0,
                 },
               ],
@@ -327,7 +327,7 @@ describe("recovery", () => {
                   workId,
                   fnHandle: "test_handle",
                   fnArgs: {},
-                  logLevel: "DEBUG",
+                  logLevel: "WARN",
                   attempt: 0,
                 },
               ],
@@ -396,7 +396,7 @@ describe("recovery", () => {
                   workId: workId1,
                   fnHandle: "test_handle",
                   fnArgs: { test: 1 },
-                  logLevel: "DEBUG",
+                  logLevel: "WARN",
                   attempt: 0,
                 },
               ],
@@ -416,7 +416,7 @@ describe("recovery", () => {
                   workId: workId2,
                   fnHandle: "test_handle",
                   fnArgs: { test: 2 },
-                  logLevel: "DEBUG",
+                  logLevel: "WARN",
                   attempt: 0,
                 },
               ],
@@ -503,7 +503,7 @@ describe("recovery", () => {
                   workId,
                   fnHandle: "test_handle",
                   fnArgs: {},
-                  logLevel: "DEBUG",
+                  logLevel: "WARN",
                   attempt: 0,
                 },
               ],