@convex-dev/workpool 0.3.1-alpha.1 → 0.3.2

package/src/component/kick.ts

@@ -1,11 +1,11 @@
 import { internal } from "./_generated/api.js";
 import { internalMutation, type MutationCtx } from "./_generated/server.js";
-import { createLogger, DEFAULT_LOG_LEVEL } from "./logging.js";
+import { getOrUpdateGlobals } from "./config.js";
+import { createLogger } from "./logging.js";
 import { INITIAL_STATE } from "./loop.js";
 import {
   boundScheduledTime,
   type Config,
-  DEFAULT_MAX_PARALLELISM,
   fromSegment,
   getCurrentSegment,
   getNextSegment,
@@ -20,9 +20,9 @@ import {
 export async function kickMainLoop(
   ctx: MutationCtx,
   source: "enqueue" | "cancel" | "complete" | "kick",
-  config?: Partial<Config>,
+  config?: Config,
 ): Promise<bigint> {
-  const globals = await getOrUpdateGlobals(ctx, config);
+  const globals = config ?? (await getOrUpdateGlobals(ctx, config));
   const console = createLogger(globals.logLevel);
   const runStatus = await getOrCreateRunStatus(ctx);
   const next = getNextSegment();
@@ -98,31 +98,3 @@ async function getOrCreateRunStatus(ctx: MutationCtx) {
   }
   return runStatus;
 }
-
-async function getOrUpdateGlobals(ctx: MutationCtx, config?: Partial<Config>) {
-  const globals = await ctx.db.query("globals").unique();
-  if (!globals) {
-    const id = await ctx.db.insert("globals", {
-      maxParallelism: config?.maxParallelism ?? DEFAULT_MAX_PARALLELISM,
-      logLevel: config?.logLevel ?? DEFAULT_LOG_LEVEL,
-    });
-    return (await ctx.db.get(id))!;
-  } else if (config) {
-    let updated = false;
-    if (
-      config.maxParallelism &&
-      config.maxParallelism !== globals.maxParallelism
-    ) {
-      globals.maxParallelism = config.maxParallelism;
-      updated = true;
-    }
-    if (config.logLevel && config.logLevel !== globals.logLevel) {
-      globals.logLevel = config.logLevel;
-      updated = true;
-    }
-    if (updated) {
-      await ctx.db.replace(globals._id, globals);
-    }
-  }
-  return globals;
-}

package/src/component/lib.test.ts

@@ -66,7 +66,7 @@ describe("lib", () => {
             logLevel: "WARN",
           },
         }),
-      ).rejects.toThrow("maxParallelism must be <= 100");
+      ).rejects.toThrow("maxParallelism must be <= 200");
     });
 
     it("should throw error if maxParallelism is too low", async () => {
@@ -78,11 +78,11 @@ describe("lib", () => {
           fnType: "mutation",
           runAt: Date.now(),
           config: {
-            maxParallelism: 0, // Less than minimum
+            maxParallelism: -1, // Less than minimum
             logLevel: "WARN",
           },
         }),
-      ).rejects.toThrow("maxParallelism must be >= 1");
+      ).rejects.toThrow("maxParallelism must be >= 0");
     });
   });
 

package/src/component/lib.ts

@@ -1,4 +1,4 @@
-import { type Infer, type ObjectType, v } from "convex/values";
+import { type ObjectType, v } from "convex/values";
 import { api } from "./_generated/api.js";
 import type { Id } from "./_generated/dataModel.js";
 import {
@@ -16,19 +16,17 @@ import {
 } from "./logging.js";
 import {
   boundScheduledTime,
-  config,
+  vConfig,
   fnType,
   getNextSegment,
   max,
-  onComplete,
+  vOnCompleteFnContext,
   retryBehavior,
   status as statusValidator,
   toSegment,
 } from "./shared.js";
 import { recordEnqueued } from "./stats.js";
-
-const MAX_POSSIBLE_PARALLELISM = 200;
-const MAX_PARALLELISM_SOFT_LIMIT = 100;
+import { getOrUpdateGlobals } from "./config.js";
 
 const itemArgs = {
   fnHandle: v.string(),
@@ -37,20 +35,20 @@ const itemArgs = {
   fnType,
   runAt: v.number(),
   // TODO: annotation?
-  onComplete: v.optional(onComplete),
+  onComplete: v.optional(vOnCompleteFnContext),
   retryBehavior: v.optional(retryBehavior),
 };
 const enqueueArgs = {
   ...itemArgs,
-  config,
+  config: vConfig.partial(),
 };
 export const enqueue = mutation({
   args: enqueueArgs,
   returns: v.id("work"),
   handler: async (ctx, { config, ...itemArgs }) => {
-    validateConfig(config);
-    const console = createLogger(config.logLevel);
-    const kickSegment = await kickMainLoop(ctx, "enqueue", config);
+    const globals = await getOrUpdateGlobals(ctx, config);
+    const console = createLogger(globals.logLevel);
+    const kickSegment = await kickMainLoop(ctx, "enqueue", globals);
     return await enqueueHandler(ctx, console, kickSegment, itemArgs);
   },
 });
@@ -73,29 +71,16 @@ async function enqueueHandler(
   return workId;
 }
 
-type Config = Infer<typeof config>;
-function validateConfig(config: Config) {
-  if (config.maxParallelism > MAX_POSSIBLE_PARALLELISM) {
-    throw new Error(`maxParallelism must be <= ${MAX_PARALLELISM_SOFT_LIMIT}`);
-  } else if (config.maxParallelism > MAX_PARALLELISM_SOFT_LIMIT) {
-    createLogger(config.logLevel).warn(
-      `maxParallelism should be <= ${MAX_PARALLELISM_SOFT_LIMIT}, but is set to ${config.maxParallelism}. This will be an error in a future version.`,
-    );
-  } else if (config.maxParallelism < 1) {
-    throw new Error("maxParallelism must be >= 1");
-  }
-}
-
 export const enqueueBatch = mutation({
   args: {
     items: v.array(v.object(itemArgs)),
-    config,
+    config: vConfig.partial(),
   },
   returns: v.array(v.id("work")),
   handler: async (ctx, { config, items }) => {
-    validateConfig(config);
-    const console = createLogger(config.logLevel);
-    const kickSegment = await kickMainLoop(ctx, "enqueue", config);
+    const globals = await getOrUpdateGlobals(ctx, config);
+    const console = createLogger(globals.logLevel);
+    const kickSegment = await kickMainLoop(ctx, "enqueue", globals);
     return Promise.all(
       items.map((item) => enqueueHandler(ctx, console, kickSegment, item)),
     );
@@ -105,12 +90,13 @@ export const enqueueBatch = mutation({
 export const cancel = mutation({
   args: {
     id: v.id("work"),
-    logLevel,
+    logLevel: v.optional(logLevel),
   },
   handler: async (ctx, { id, logLevel }) => {
-    const shouldCancel = await shouldCancelWorkItem(ctx, id, logLevel);
+    const globals = await getOrUpdateGlobals(ctx, { logLevel });
+    const shouldCancel = await shouldCancelWorkItem(ctx, id, globals.logLevel);
     if (shouldCancel) {
-      const segment = await kickMainLoop(ctx, "cancel", { logLevel });
+      const segment = await kickMainLoop(ctx, "cancel", globals);
       await ctx.db.insert("pendingCancelation", {
         workId: id,
         segment,
@@ -122,7 +108,7 @@ export const cancel = mutation({
 const PAGE_SIZE = 64;
 export const cancelAll = mutation({
   args: {
-    logLevel,
+    logLevel: v.optional(logLevel),
     before: v.optional(v.number()),
     limit: v.optional(v.number()),
   },
@@ -134,14 +120,15 @@ export const cancelAll = mutation({
       .withIndex("by_creation_time", (q) => q.lte("_creationTime", beforeTime))
       .order("desc")
       .take(pageSize);
+    const globals = await getOrUpdateGlobals(ctx, { logLevel });
     const shouldCancel = await Promise.all(
       pageOfWork.map(async ({ _id }) =>
-        shouldCancelWorkItem(ctx, _id, logLevel),
+        shouldCancelWorkItem(ctx, _id, globals.logLevel),
       ),
     );
     let segment = getNextSegment();
     if (shouldCancel.some((c) => c)) {
-      segment = await kickMainLoop(ctx, "cancel", { logLevel });
+      segment = await kickMainLoop(ctx, "cancel", globals);
     }
     await Promise.all(
       pageOfWork.map(({ _id }, index) => {

package/src/component/loop.ts

@@ -23,7 +23,7 @@ import {
 } from "./shared.js";
 import { generateReport, recordCompleted, recordStarted } from "./stats.js";
 
-const CANCELATION_BATCH_SIZE = 64; // the only queue that can get unbounded.
+const CANCELLATION_BATCH_SIZE = 64; // the only queue that can get unbounded.
 const SECOND = 1000;
 const MINUTE = 60 * SECOND;
 const RECOVERY_THRESHOLD_MS = 5 * MINUTE; // attempt to recover jobs this old.
@@ -75,7 +75,7 @@ export const main = internalMutation({
 
     // Read pendingCancelation, deleting from pendingStart. If it's still running, queue to cancel.
     console.time("[main] pendingCancelation");
-    await handleCancelation(ctx, state, segment, console, toCancel, globals);
+    await handleCancelation(ctx, state, segment, console, toCancel);
     console.timeEnd("[main] pendingCancelation");
 
     if (state.running.length === 0) {
@@ -415,7 +415,6 @@ async function handleCancelation(
   segment: bigint,
   console: Logger,
   toCancel: CompleteJob[],
-  { cancelationBatchSize }: Config,
 ) {
   const start = state.segmentCursors.cancelation - CURSOR_BUFFER_SEGMENTS;
   const canceled = await ctx.db
@@ -423,7 +422,7 @@ async function handleCancelation(
     .withIndex("segment", (q) =>
       q.gte("segment", start).lte("segment", segment),
     )
-    .take(cancelationBatchSize ?? CANCELATION_BATCH_SIZE);
+    .take(CANCELLATION_BATCH_SIZE);
   state.segmentCursors.cancelation = canceled.at(-1)?.segment ?? segment;
   if (canceled.length) {
     console.debug(`[main] attempting to cancel ${canceled.length}`);
@@ -562,7 +561,12 @@ async function beginWork(
     throw new Error("work not found");
   }
   recordStarted(console, work, lagMs);
-  const { attempts: attempt, fnHandle, fnArgs } = work;
+  let fnArgs = work.fnArgs;
+  if (fnArgs === undefined && work.payloadId) {
+    const payload = await ctx.db.get(work.payloadId);
+    fnArgs = payload?.args ?? {};
+  }
+  const { attempts: attempt, fnHandle } = work;
   const args = { workId, fnHandle, fnArgs, logLevel, attempt };
   if (work.fnType === "action") {
     return ctx.scheduler.runAfter(0, internal.worker.runActionWrapper, args);

package/src/component/recovery.test.ts

@@ -1,5 +1,12 @@
 import { convexTest } from "convex-test";
-import type { WithoutSystemFields } from "convex/server";
+import type {
+  DocumentByName,
+  GenericDatabaseReader,
+  GenericDataModel,
+  SystemDataModel,
+  SystemTableNames,
+  WithoutSystemFields,
+} from "convex/server";
 import {
   afterEach,
   assert,
@@ -14,8 +21,7 @@ import type { Doc, Id } from "./_generated/dataModel.js";
 import type { MutationCtx } from "./_generated/server.js";
 import { recoveryHandler } from "./recovery.js";
 import schema from "./schema.js";
-
-const modules = import.meta.glob("./**/*.ts");
+import { modules } from "./setup.test.js";
 
 describe("recovery", () => {
   async function setupTest() {
@@ -196,13 +202,7 @@ describe("recovery", () => {
       // 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);
-        };
+        ctx.db.system.get = patchedSystemGet(ctx.db, { [scheduledId]: null });
 
         await recoveryHandler(ctx, {
           jobs: [
@@ -244,31 +244,27 @@ describe("recovery", () => {
       // 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",
+        ctx.db.system.get = patchedSystemGet(ctx.db, {
+          [scheduledId]: {
+            _id: scheduledId,
+            _creationTime: Date.now(),
+            name: "internal/worker.runActionWrapper",
+            args: [
+              {
+                workId,
+                fnHandle: "test_handle",
+                fnArgs: {},
+                logLevel: "WARN",
+                attempt: 0,
               },
-            };
-          }
-          return await originalGet(id);
-        };
+            ],
+            scheduledTime: Date.now(),
+            state: {
+              kind: "failed",
+              error: "Function execution failed",
+            },
+          },
+        });
 
         await recoveryHandler(ctx, {
           jobs: [
@@ -310,30 +306,26 @@ describe("recovery", () => {
       // 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",
+        ctx.db.system.get = patchedSystemGet(ctx.db, {
+          [scheduledId]: {
+            _id: scheduledId,
+            _creationTime: Date.now(),
+            name: "internal/worker.runActionWrapper",
+            args: [
+              {
+                workId,
+                fnHandle: "test_handle",
+                fnArgs: {},
+                logLevel: "WARN",
+                attempt: 0,
               },
-            };
-          }
-          return await originalGet(id);
-        };
+            ],
+            scheduledTime: Date.now(),
+            state: {
+              kind: "canceled",
+            },
+          },
+        });
 
         await recoveryHandler(ctx, {
           jobs: [
@@ -379,50 +371,45 @@ describe("recovery", () => {
       // 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",
+        ctx.db.system.get = patchedSystemGet(ctx.db, {
+          [scheduledId1]: {
+            _id: scheduledId1,
+            _creationTime: Date.now(),
+            name: "internal/worker.runActionWrapper",
+            args: [
+              {
+                workId: workId1,
+                fnHandle: "test_handle",
+                fnArgs: { test: 1 },
+                logLevel: "WARN",
+                attempt: 0,
               },
-            };
-          } 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",
+            ],
+            scheduledTime: Date.now(),
+            state: {
+              kind: "failed",
+              error: "Function 1 failed",
+            },
+          },
+          [scheduledId2]: {
+            _id: scheduledId2,
+            _creationTime: Date.now(),
+            name: "internal/worker.runActionWrapper",
+            args: [
+              {
+                workId: workId2,
+                fnHandle: "test_handle",
+                fnArgs: { test: 2 },
+                logLevel: "WARN",
+                attempt: 0,
               },
-            };
-          }
-          return await originalGet(id);
-        };
+            ],
+            scheduledTime: Date.now(),
+            state: {
+              kind: "canceled",
+            },
+          },
+        });
 
         await recoveryHandler(ctx, {
           jobs: [
@@ -486,30 +473,26 @@ describe("recovery", () => {
       // 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",
+        ctx.db.system.get = patchedSystemGet(ctx.db, {
+          [scheduledId]: {
+            _id: scheduledId,
+            _creationTime: Date.now(),
+            name: "internal/worker.runActionWrapper",
+            args: [
+              {
+                workId,
+                fnHandle: "test_handle",
+                fnArgs: {},
+                logLevel: "WARN",
+                attempt: 0,
               },
-            };
-          }
-          return await originalGet(id);
-        };
+            ],
+            scheduledTime: Date.now(),
+            state: {
+              kind: "pending",
+            },
+          },
+        });
 
         await recoveryHandler(ctx, {
           jobs: [
@@ -534,3 +517,20 @@ describe("recovery", () => {
     });
   });
 });
+
+function patchedSystemGet(
+  db: GenericDatabaseReader<GenericDataModel>,
+  overrides: Record<
+    string,
+    DocumentByName<SystemDataModel, "_scheduled_functions"> | null
+  >,
+) {
+  const originalGet = db.system.get;
+  return async (
+    tableOrId: SystemTableNames | Id<SystemTableNames>,
+    maybeId?: Id<SystemTableNames>,
+  ) => {
+    const id = (maybeId ?? tableOrId) as Id<"_scheduled_functions">;
+    return id in overrides ? overrides[id] : await originalGet(id);
+  };
+}

package/src/component/schema.ts

@@ -2,10 +2,10 @@ import { defineSchema, defineTable } from "convex/server";
 import { v } from "convex/values";
 import {
   fnType,
-  config,
-  onComplete,
+  vConfig,
+  vOnCompleteFnContext,
   retryBehavior,
-  vResultValidator,
+  vResult,
 } from "./shared.js";
 
 // Represents a slice of time to process work.
@@ -13,7 +13,7 @@ const segment = v.int64();
 
 export default defineSchema({
   // Written from kickLoop, read everywhere.
-  globals: defineTable(config),
+  globals: defineTable(vConfig),
   // Singleton, only read & written by `main`.
   internalState: defineTable({
     // Ensure that only one main is running at a time.
@@ -62,9 +62,12 @@ export default defineSchema({
     fnType,
     fnHandle: v.string(),
     fnName: v.string(),
-    fnArgs: v.any(),
+    fnArgs: v.optional(v.any()),
+    // Reference to large args/onComplete context if stored separately
+    payloadId: v.optional(v.id("payload")),
+    payloadSize: v.optional(v.number()),
     attempts: v.number(), // number of completed attempts
-    onComplete: v.optional(onComplete),
+    onComplete: v.optional(vOnCompleteFnContext),
     retryBehavior: v.optional(retryBehavior),
     canceled: v.optional(v.boolean()),
   }),
@@ -80,7 +83,7 @@ export default defineSchema({
   // Written by complete, read & deleted by `main`.
   pendingCompletion: defineTable({
     segment,
-    runResult: vResultValidator,
+    runResult: vResult,
     workId: v.id("work"),
     retry: v.boolean(),
   })
@@ -94,4 +97,10 @@ export default defineSchema({
   })
     .index("workId", ["workId"])
     .index("segment", ["segment"]),
+
+  // Store large data separately to avoid document size limits
+  payload: defineTable({
+    args: v.optional(v.record(v.string(), v.any())),
+    context: v.optional(v.any()),
+  }),
 });

package/src/component/shared.ts

@@ -33,12 +33,11 @@ export function fromSegment(segment: bigint): number {
   return Number(segment) * SEGMENT_MS;
 }
 
-export const config = v.object({
+export const vConfig = v.object({
   maxParallelism: v.number(),
   logLevel,
-  cancelationBatchSize: v.optional(v.number()),
 });
-export type Config = Infer<typeof config>;
+export type Config = Infer<typeof vConfig>;
 
 export const retryBehavior = v.object({
   maxAttempts: v.number(),
@@ -71,7 +70,7 @@ export const DEFAULT_RETRY_BEHAVIOR: RetryBehavior = {
 // This ensures that the type satisfies the schema.
 const _ = {} as RetryBehavior satisfies Infer<typeof retryBehavior>;
 
-export const vResultValidator = v.union(
+export const vResult = v.union(
   v.object({
     kind: v.literal("success"),
     returnValue: v.any(),
@@ -84,13 +83,12 @@ export const vResultValidator = v.union(
     kind: v.literal("canceled"),
   }),
 );
-export type RunResult = Infer<typeof vResultValidator>;
+export type RunResult = Infer<typeof vResult>;
 
-export const onComplete = v.object({
+export const vOnCompleteFnContext = v.object({
   fnHandle: v.string(), // mutation
   context: v.optional(v.any()),
 });
-export type OnComplete = Infer<typeof onComplete>;
 
 export type OnCompleteArgs = {
   /**

package/src/component/stats.test.ts

@@ -190,7 +190,7 @@ describe("stats", () => {
       });
 
       // Create more pending start items than maxParallelism
-      const maxParallelism = 5;
+      const maxParallelism = 50;
 
       // Create maxParallelism + 1 work items to trigger pagination
       for (let i = 0; i < maxParallelism + 1; i++) {

package/src/component/stats.ts

@@ -81,7 +81,7 @@ export async function generateReport(
         .lt("segment", currentSegment),
     )
     .paginate({
-      numItems: maxParallelism,
+      numItems: Math.max(maxParallelism, 10),
       cursor: null,
     });
   if (pendingStart.isDone) {