@convex-dev/workpool 0.2.0 → 0.2.1

package/src/component/complete.ts

@@ -1,9 +1,10 @@
 import { FunctionHandle } from "convex/server";
 import { Infer, v } from "convex/values";
+import { Id } from "./_generated/dataModel.js";
 import { internalMutation, MutationCtx } from "./_generated/server.js";
 import { kickMainLoop } from "./kick.js";
 import { createLogger } from "./logging.js";
-import { nextSegment, OnCompleteArgs, runResult } from "./shared.js";
+import { OnCompleteArgs, RunResult, runResult } from "./shared.js";
 import { recordCompleted } from "./stats.js";
 
 export type CompleteJob = Infer<typeof completeArgs.fields.jobs.element>;
@@ -23,7 +24,11 @@ export async function completeHandler(
 ) {
   const globals = await ctx.db.query("globals").unique();
   const console = createLogger(globals?.logLevel);
-  let anyPendingCompletions = false;
+  const pendingCompletions: {
+    runResult: RunResult;
+    workId: Id<"work">;
+    retry: boolean;
+  }[] = [];
   await Promise.all(
     args.jobs.map(async (job) => {
       const work = await ctx.db.get(job.workId);
@@ -77,18 +82,24 @@ export async function completeHandler(
         await ctx.db.delete(job.workId);
       }
       if (job.runResult.kind !== "canceled") {
-        await ctx.db.insert("pendingCompletion", {
+        pendingCompletions.push({
           runResult: job.runResult,
           workId: job.workId,
-          segment: nextSegment(),
           retry,
         });
-        anyPendingCompletions = true;
       }
     })
   );
-  if (anyPendingCompletions) {
-    await kickMainLoop(ctx, "complete");
+  if (pendingCompletions.length > 0) {
+    const segment = await kickMainLoop(ctx, "complete");
+    await Promise.all(
+      pendingCompletions.map((completion) =>
+        ctx.db.insert("pendingCompletion", {
+          ...completion,
+          segment,
+        })
+      )
+    );
   }
 }
 

package/src/component/kick.test.ts

@@ -14,7 +14,12 @@ import { kickMainLoop } from "./kick.js";
 import { DEFAULT_LOG_LEVEL } from "./logging.js";
 import schema from "./schema.js";
 import { modules } from "./setup.test.js";
-import { fromSegment, nextSegment, toSegment } from "./shared";
+import {
+  fromSegment,
+  getCurrentSegment,
+  getNextSegment,
+  toSegment,
+} from "./shared";
 import { DEFAULT_MAX_PARALLELISM } from "./shared.js";
 
 describe("kickMainLoop", () => {
@@ -74,11 +79,12 @@ describe("kickMainLoop", () => {
       expect(runStatus.state.kind).toBe("running");
 
       // Second kick should not change state
-      await kickMainLoop(ctx, "enqueue");
+      const segment = await kickMainLoop(ctx, "enqueue");
       const afterStatus = await ctx.db.query("runStatus").unique();
       assert(afterStatus);
       expect(afterStatus.state.kind).toBe("running");
       expect(afterStatus._id).toBe(runStatus._id);
+      expect(segment).toBe(getNextSegment());
     });
   });
 
@@ -115,7 +121,8 @@ describe("kickMainLoop", () => {
       });
 
       // Kick should reschedule to run sooner
-      await kickMainLoop(ctx, "enqueue");
+      const segment = await kickMainLoop(ctx, "enqueue");
+      expect(segment).toBe(getCurrentSegment());
 
       const afterStatus = await ctx.db.query("runStatus").unique();
       assert(afterStatus);
@@ -156,7 +163,8 @@ describe("kickMainLoop", () => {
       });
 
       // Kick should not change state when saturated
-      await kickMainLoop(ctx, "enqueue");
+      const segment = await kickMainLoop(ctx, "enqueue");
+      expect(segment).toBe(getNextSegment());
       const afterStatus = await ctx.db.query("runStatus").unique();
       assert(afterStatus);
       expect(afterStatus.state.kind).toBe("scheduled");
@@ -209,13 +217,15 @@ describe("kickMainLoop", () => {
   test("handles race conditions between multiple kicks", async () => {
     const t = convexTest(schema, modules);
     // Run kicks in separate transactions to simulate concurrent access
-    await Promise.all(
+    const segments = await Promise.all(
       Array.from({ length: 10 }, () =>
         t.run(async (ctx) => {
-          await kickMainLoop(ctx, "enqueue");
+          const segment = await kickMainLoop(ctx, "enqueue");
+          return segment;
         })
       )
     );
+    expect(segments.filter((s) => s === getCurrentSegment())).toHaveLength(1);
 
     // Check final state in a new transaction
     await t.run(async (ctx) => {
@@ -265,7 +275,7 @@ describe("kickMainLoop", () => {
       await kickMainLoop(ctx, "enqueue");
       const runStatus = await ctx.db.query("runStatus").unique();
       assert(runStatus);
-      const segment = nextSegment() + 10n;
+      const segment = getNextSegment() + 10n;
       await ctx.db.patch(runStatus._id, {
         state: {
           generation: 0n,

package/src/component/kick.ts

@@ -7,43 +7,44 @@ import {
   Config,
   DEFAULT_MAX_PARALLELISM,
   fromSegment,
-  nextSegment,
+  getCurrentSegment,
+  getNextSegment,
 } from "./shared.js";
 
 /**
- * Called from outside the loop:
+ * Called from outside the loop.
+ * Returns the soonest segment to enqueue work for the main loop.
  */
-
 export async function kickMainLoop(
   ctx: MutationCtx,
   source: "enqueue" | "cancel" | "complete",
   config?: Partial<Config>
-): Promise<void> {
+): Promise<bigint> {
   const globals = await getOrUpdateGlobals(ctx, config);
   const console = createLogger(globals.logLevel);
   const runStatus = await getOrCreateRunStatus(ctx);
+  const next = getNextSegment();
 
   // Only kick to run now if we're scheduled or idle.
   if (runStatus.state.kind === "running") {
     console.debug(
       `[${source}] main is actively running, so we don't need to kick it`
     );
-    return;
+    return next;
   }
-  const segment = nextSegment();
   // main is scheduled to run later, so we should cancel it and reschedule.
   if (runStatus.state.kind === "scheduled") {
     if (source === "enqueue" && runStatus.state.saturated) {
       console.debug(
         `[${source}] main is saturated, so we don't need to kick it`
       );
-      return;
+      return next;
     }
-    if (runStatus.state.segment <= segment) {
+    if (runStatus.state.segment <= next) {
       console.debug(
         `[${source}] main is scheduled to run soon enough, so we don't need to kick it`
       );
-      return;
+      return next;
     }
     console.debug(
       `[${source}] main is scheduled to run later, so reschedule it to run now`
@@ -60,11 +61,13 @@ export async function kickMainLoop(
     console.debug(`[${source}] main was idle, so run it now`);
   }
   await ctx.db.patch(runStatus._id, { state: { kind: "running" } });
-  const scheduledTime = boundScheduledTime(fromSegment(segment), console);
+  const current = getCurrentSegment();
+  const scheduledTime = boundScheduledTime(fromSegment(current), console);
   await ctx.scheduler.runAt(scheduledTime, internal.loop.main, {
     generation: runStatus.state.generation,
-    segment,
+    segment: current,
   });
+  return current;
 }
 
 export const forceKick = internalMutation({

package/src/component/lib.ts

@@ -1,20 +1,19 @@
 import { v } from "convex/values";
+import { api } from "./_generated/api.js";
+import { Id } from "./_generated/dataModel.js";
 import { mutation, MutationCtx, query } from "./_generated/server.js";
+import { kickMainLoop } from "./kick.js";
+import { createLogger, LogLevel, logLevel } from "./logging.js";
 import {
-  nextSegment,
+  boundScheduledTime,
+  config,
+  getNextSegment,
+  max,
   onComplete,
   retryBehavior,
-  config,
   status as statusValidator,
   toSegment,
-  boundScheduledTime,
-  max,
 } from "./shared.js";
-import { LogLevel, logLevel } from "./logging.js";
-import { kickMainLoop } from "./kick.js";
-import { api } from "./_generated/api.js";
-import { createLogger } from "./logging.js";
-import { Id } from "./_generated/dataModel.js";
 import { recordEnqueued } from "./stats.js";
 
 const MAX_POSSIBLE_PARALLELISM = 100;
@@ -45,11 +44,11 @@ export const enqueue = mutation({
       ...workArgs,
       attempts: 0,
     });
+    const limit = await kickMainLoop(ctx, "enqueue", config);
     await ctx.db.insert("pendingStart", {
       workId,
-      segment: max(toSegment(runAt), nextSegment()),
+      segment: max(toSegment(runAt), limit),
     });
-    await kickMainLoop(ctx, "enqueue", config);
     recordEnqueued(console, { workId, fnName: workArgs.fnName, runAt });
     return workId;
   },
@@ -61,11 +60,14 @@ export const cancel = mutation({
     logLevel,
   },
   handler: async (ctx, { id, logLevel }) => {
-    const canceled = await cancelWorkItem(ctx, id, nextSegment(), logLevel);
-    if (canceled) {
-      await kickMainLoop(ctx, "cancel", { logLevel });
+    const shouldCancel = await shouldCancelWorkItem(ctx, id, logLevel);
+    if (shouldCancel) {
+      const segment = await kickMainLoop(ctx, "cancel", { logLevel });
+      await ctx.db.insert("pendingCancelation", {
+        workId: id,
+        segment,
+      });
     }
-    // TODO: stats event
   },
 });
 
@@ -74,20 +76,30 @@ export const cancelAll = mutation({
   args: { logLevel, before: v.optional(v.number()) },
   handler: async (ctx, { logLevel, before }) => {
     const beforeTime = before ?? Date.now();
-    const segment = nextSegment();
     const pageOfWork = await ctx.db
       .query("work")
       .withIndex("by_creation_time", (q) => q.lte("_creationTime", beforeTime))
       .order("desc")
       .take(PAGE_SIZE);
-    const canceled = await Promise.all(
+    const shouldCancel = await Promise.all(
       pageOfWork.map(async ({ _id }) =>
-        cancelWorkItem(ctx, _id, segment, logLevel)
+        shouldCancelWorkItem(ctx, _id, logLevel)
       )
     );
-    if (canceled.some((c) => c)) {
-      await kickMainLoop(ctx, "cancel", { logLevel });
+    let segment = getNextSegment();
+    if (shouldCancel.some((c) => c)) {
+      segment = await kickMainLoop(ctx, "cancel", { logLevel });
     }
+    await Promise.all(
+      pageOfWork.map(({ _id }, index) => {
+        if (shouldCancel[index]) {
+          return ctx.db.insert("pendingCancelation", {
+            workId: _id,
+            segment,
+          });
+        }
+      })
+    );
     if (pageOfWork.length === PAGE_SIZE) {
       await ctx.scheduler.runAfter(0, api.lib.cancelAll, {
         logLevel,
@@ -124,10 +136,9 @@ export const status = query({
   },
 });
 
-async function cancelWorkItem(
+async function shouldCancelWorkItem(
   ctx: MutationCtx,
   workId: Id<"work">,
-  segment: bigint,
   logLevel: LogLevel
 ) {
   const console = createLogger(logLevel);
@@ -145,10 +156,6 @@ async function cancelWorkItem(
     console.warn(`[cancel] work ${workId} has already been canceled`);
     return false;
   }
-  await ctx.db.insert("pendingCancelation", {
-    workId,
-    segment,
-  });
   return true;
 }
 

package/src/component/logging.test.ts

@@ -0,0 +1,16 @@
+import { describe, expect, it } from "vitest";
+import { shouldLog } from "./logging";
+
+describe("logging", () => {
+  describe("shouldLog", () => {
+    it("should return true if the log level is above the config level", () => {
+      expect(shouldLog("INFO", "DEBUG")).toBe(false);
+    });
+    it("should return false if the log level is below the config level", () => {
+      expect(shouldLog("INFO", "WARN")).toBe(true);
+    });
+    it("should return true if the log level is equal to the config level", () => {
+      expect(shouldLog("INFO", "INFO")).toBe(true);
+    });
+  });
+});

package/src/component/logging.ts

@@ -1,6 +1,17 @@
 import { v, Infer } from "convex/values";
 
-export const DEFAULT_LOG_LEVEL: LogLevel = "WARN";
+export const DEFAULT_LOG_LEVEL: LogLevel = "REPORT";
+
+// NOTE: the ordering here is important! A config level of "INFO" will log
+// "INFO", "REPORT", "WARN",and "ERROR" events.
+export const logLevel = v.union(
+  v.literal("DEBUG"),
+  v.literal("INFO"),
+  v.literal("REPORT"),
+  v.literal("WARN"),
+  v.literal("ERROR")
+);
+export type LogLevel = Infer<typeof logLevel>;
 
 export type Logger = {
   debug: (...args: unknown[]) => void;
@@ -12,59 +23,70 @@ export type Logger = {
   event: (event: string, payload: Record<string, unknown>) => void;
 };
 
+const logLevelOrder = logLevel.members.map((l) => l.value);
+const logLevelByName = logLevelOrder.reduce(
+  (acc, l, i) => {
+    acc[l] = i;
+    return acc;
+  },
+  {} as Record<LogLevel, number>
+);
+export function shouldLog(config: LogLevel, level: LogLevel) {
+  return logLevelByName[config] <= logLevelByName[level];
+}
+const DEBUG = logLevelByName["DEBUG"];
+const INFO = logLevelByName["INFO"];
+const REPORT = logLevelByName["REPORT"];
+const WARN = logLevelByName["WARN"];
+const ERROR = logLevelByName["ERROR"];
+
 export function createLogger(level?: LogLevel): Logger {
-  const levelIndex = ["DEBUG", "INFO", "WARN", "ERROR"].indexOf(
-    level ?? DEFAULT_LOG_LEVEL
-  );
-  if (levelIndex === -1) {
+  const levelIndex = logLevelByName[level ?? DEFAULT_LOG_LEVEL];
+  if (levelIndex === undefined) {
     throw new Error(`Invalid log level: ${level}`);
   }
   return {
     debug: (...args: unknown[]) => {
-      if (levelIndex <= 0) {
+      if (levelIndex <= DEBUG) {
         console.debug(...args);
       }
     },
     info: (...args: unknown[]) => {
-      if (levelIndex <= 1) {
+      if (levelIndex <= INFO) {
         console.info(...args);
       }
     },
     warn: (...args: unknown[]) => {
-      if (levelIndex <= 2) {
+      if (levelIndex <= WARN) {
         console.warn(...args);
       }
     },
     error: (...args: unknown[]) => {
-      if (levelIndex <= 3) {
+      if (levelIndex <= ERROR) {
         console.error(...args);
       }
     },
     time: (label: string) => {
-      if (levelIndex <= 0) {
+      if (levelIndex <= DEBUG) {
         console.time(label);
       }
     },
     timeEnd: (label: string) => {
-      if (levelIndex <= 0) {
+      if (levelIndex <= DEBUG) {
         console.timeEnd(label);
       }
     },
     event: (event: string, payload: Record<string, unknown>) => {
-      if (levelIndex <= 1) {
-        const fullPayload = {
-          event,
-          ...payload,
-        };
+      const fullPayload = {
+        component: "workpool",
+        event,
+        ...payload,
+      };
+      if (levelIndex === REPORT && event === "report") {
+        console.info(JSON.stringify(fullPayload));
+      } else if (levelIndex <= INFO) {
         console.info(JSON.stringify(fullPayload));
       }
     },
   };
 }
-export const logLevel = v.union(
-  v.literal("DEBUG"),
-  v.literal("INFO"),
-  v.literal("WARN"),
-  v.literal("ERROR")
-);
-export type LogLevel = Infer<typeof logLevel>;

package/src/component/loop.test.ts

@@ -14,9 +14,9 @@ import { Doc, Id } from "./_generated/dataModel";
 import { MutationCtx } from "./_generated/server";
 import schema from "./schema";
 import {
-  currentSegment,
   DEFAULT_MAX_PARALLELISM,
-  nextSegment,
+  getCurrentSegment,
+  getNextSegment,
   toSegment,
 } from "./shared";
 
@@ -72,7 +72,7 @@ describe("loop", () => {
     await ctx.db.insert("internalState", {
       generation: 1n,
       segmentCursors: { incoming: 0n, completion: 0n, cancelation: 0n },
-      lastRecovery: currentSegment(),
+      lastRecovery: getCurrentSegment(),
       report: {
         completed: 0,
         succeeded: 0,
@@ -281,7 +281,7 @@ describe("loop", () => {
       // Run main loop to process pendingCompletion -> pendingStart
       await t.mutation(internal.loop.main, {
         generation: 1n,
-        segment: nextSegment(),
+        segment: getNextSegment(),
       });
 
       // Verify work is now in pendingStart for retry
@@ -364,13 +364,13 @@ describe("loop", () => {
       // Run main loop to process the work
       await t.mutation(internal.loop.main, {
         generation: 1n,
-        segment: nextSegment(),
+        segment: getNextSegment(),
       });
 
       // Run updateRunStatus to transition to scheduled
       await t.mutation(internal.loop.updateRunStatus, {
         generation: 2n,
-        segment: nextSegment(),
+        segment: getNextSegment(),
       });
 
       // Verify state transition to scheduled
@@ -387,7 +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();
+      const segment = getCurrentSegment();
       await t.run(async (ctx) => {
         // Create work item
         const workId = await makeDummyWork(ctx);
@@ -441,14 +441,14 @@ describe("loop", () => {
         const scheduledId = await ctx.scheduler.runAfter(
           1000,
           internal.loop.main,
-          { generation: 1n, segment: nextSegment() + 10n }
+          { generation: 1n, segment: getNextSegment() + 10n }
         );
 
         // Create scheduled runStatus
         await ctx.db.insert("runStatus", {
           state: {
             kind: "scheduled",
-            segment: nextSegment() + 10n,
+            segment: getNextSegment() + 10n,
             scheduledId,
             saturated: false,
             generation: 1n,
@@ -481,7 +481,7 @@ describe("loop", () => {
     });
 
     it("should transition from running to idle when all work is done", async () => {
-      const segment = nextSegment();
+      const segment = getNextSegment();
       // Setup initial running state with work
       const workId = await t.run<Id<"work">>(async (ctx) => {
         // Create internal state
@@ -537,7 +537,7 @@ describe("loop", () => {
       });
     });
     it("should transition from scheduled to running when main loop runs", async () => {
-      const segment = nextSegment();
+      const segment = getNextSegment();
       await t.run(async (ctx) => {
         await insertInternalState(ctx);
 
@@ -1037,7 +1037,7 @@ describe("loop", () => {
 
     it("should set saturated flag when at max capacity", async () => {
       // Setup state with running jobs at max capacity
-      const now = currentSegment();
+      const now = getCurrentSegment();
       const later = now + 10n;
       await setMaxParallelism(10);
       await t.run(async (ctx) => {

package/src/component/loop.ts

@@ -13,15 +13,15 @@ import {
 import {
   boundScheduledTime,
   Config,
-  currentSegment,
   DEFAULT_MAX_PARALLELISM,
   fromSegment,
+  getCurrentSegment,
+  getNextSegment,
   max,
-  nextSegment,
   RunResult,
   toSegment,
 } from "./shared.js";
-import { recordCompleted, recordReport, recordStarted } from "./stats.js";
+import { recordCompleted, generateReport, recordStarted } from "./stats.js";
 
 const CANCELLATION_BATCH_SIZE = 64; // the only queue that can get unbounded.
 const SECOND = 1000;
@@ -100,7 +100,7 @@ export const main = internalMutation({
         // It's been a while, let's start fresh.
         lastReportTs = Date.now();
       }
-      recordReport(console, state);
+      await generateReport(ctx, console, state, globals);
       state.report = {
         completed: 0,
         succeeded: 0,
@@ -149,22 +149,22 @@ export const updateRunStatus = internalMutation({
 
     // TODO: check for current segment (or from args) first, to avoid OCCs.
     console.time("[updateRunStatus] nextSegmentIsActionable");
-    const next = max(segment + 1n, currentSegment());
+    const nextSegment = max(segment + 1n, getCurrentSegment());
     const nextIsActionable = await nextSegmentIsActionable(
       ctx,
       state,
       maxParallelism,
-      next
+      nextSegment
     );
     console.timeEnd("[updateRunStatus] nextSegmentIsActionable");
 
     if (nextIsActionable) {
       await ctx.scheduler.runAt(
-        boundScheduledTime(fromSegment(next), console),
+        boundScheduledTime(fromSegment(nextSegment), console),
         internal.loop.main,
         {
           generation,
-          segment: next,
+          segment: nextSegment,
         }
       );
       return;
@@ -187,7 +187,7 @@ export const updateRunStatus = internalMutation({
       });
       await ctx.scheduler.runAfter(0, internal.loop.main, {
         generation,
-        segment: currentSegment(),
+        segment: getCurrentSegment(),
       });
       return;
     }
@@ -204,7 +204,7 @@ export const updateRunStatus = internalMutation({
     }
     const docs = await Promise.all(
       actionableTables.map(async (tableName) =>
-        getNextUp(ctx, tableName, { start: next })
+        getNextUp(ctx, tableName, { start: nextSegment })
       )
     );
     console.timeEnd("[updateRunStatus] findNextSegment");
@@ -223,7 +223,7 @@ export const updateRunStatus = internalMutation({
         internal.loop.main,
         { generation, segment: targetSegment }
       );
-      if (targetSegment > nextSegment()) {
+      if (targetSegment > getNextSegment()) {
         await ctx.db.patch(runStatus._id, {
           state: {
             kind: "scheduled",

package/src/component/recovery.test.ts

@@ -13,12 +13,7 @@ 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;
+import { recoveryHandler } from "./recovery";
 
 const modules = import.meta.glob("./**/*.ts");
 
@@ -209,7 +204,7 @@ describe("recovery", () => {
           return await originalGet(id);
         };
 
-        await recover(ctx, {
+        await recoveryHandler(ctx, {
           jobs: [
             {
               scheduledId,
@@ -275,7 +270,7 @@ describe("recovery", () => {
           return await originalGet(id);
         };
 
-        await recover(ctx, {
+        await recoveryHandler(ctx, {
           jobs: [
             {
               scheduledId,
@@ -340,7 +335,7 @@ describe("recovery", () => {
           return await originalGet(id);
         };
 
-        await recover(ctx, {
+        await recoveryHandler(ctx, {
           jobs: [
             {
               scheduledId,
@@ -429,7 +424,7 @@ describe("recovery", () => {
           return await originalGet(id);
         };
 
-        await recover(ctx, {
+        await recoveryHandler(ctx, {
           jobs: [
             {
               scheduledId: scheduledId1,
@@ -516,7 +511,7 @@ describe("recovery", () => {
           return await originalGet(id);
         };
 
-        await recover(ctx, {
+        await recoveryHandler(ctx, {
           jobs: [
             {
               scheduledId,