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

package/src/component/config.ts

@@ -0,0 +1,72 @@
+import { mutation, type MutationCtx } from "./_generated/server.js";
+import { vConfig, DEFAULT_MAX_PARALLELISM, type Config } from "./shared.js";
+import { createLogger, DEFAULT_LOG_LEVEL } from "./logging.js";
+import { kickMainLoop } from "./kick.js";
+
+export const MAX_POSSIBLE_PARALLELISM = 200;
+export const MAX_PARALLELISM_SOFT_LIMIT = 100;
+
+export const update = mutation({
+  args: vConfig.partial(),
+  handler: async (ctx, args) => {
+    const { globals, previousValue } = await _getOrUpdateGlobals(ctx, args);
+    if (args.maxParallelism && args.maxParallelism > previousValue) {
+      await kickMainLoop(ctx, "kick", globals);
+    }
+  },
+});
+
+export function validateConfig(config: Partial<Config>) {
+  if (config.maxParallelism !== undefined) {
+    if (config.maxParallelism > MAX_POSSIBLE_PARALLELISM) {
+      throw new Error(`maxParallelism must be <= ${MAX_POSSIBLE_PARALLELISM}`);
+    } else if (config.maxParallelism > MAX_PARALLELISM_SOFT_LIMIT) {
+      createLogger(config.logLevel ?? DEFAULT_LOG_LEVEL).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 < 0) {
+      throw new Error("maxParallelism must be >= 0");
+    }
+  }
+}
+export async function getOrUpdateGlobals(
+  ctx: MutationCtx,
+  config?: Partial<Config>,
+) {
+  const { globals } = await _getOrUpdateGlobals(ctx, config);
+  return globals;
+}
+async function _getOrUpdateGlobals(
+  ctx: MutationCtx,
+  config?: Partial<Config>,
+) {
+  if (config) {
+    validateConfig(config);
+  }
+  const globals = await ctx.db.query("globals").unique();
+  const previousValue = globals?.maxParallelism ?? DEFAULT_MAX_PARALLELISM;
+  if (!globals) {
+    const id = await ctx.db.insert("globals", {
+      maxParallelism: config?.maxParallelism ?? DEFAULT_MAX_PARALLELISM,
+      logLevel: config?.logLevel ?? DEFAULT_LOG_LEVEL,
+    });
+    return { globals: (await ctx.db.get("globals", id))!, previousValue };
+  } else if (config) {
+    let updated = false;
+    if (
+      config.maxParallelism !== undefined &&
+      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", globals._id, globals);
+    }
+  }
+  return { globals, previousValue };
+}

package/src/component/kick.test.ts

@@ -20,6 +20,7 @@ import {
   getNextSegment,
   toSegment,
 } from "./shared.js";
+import { getOrUpdateGlobals } from "./config.js";
 
 describe("kickMainLoop", () => {
   beforeEach(() => {
@@ -47,27 +48,6 @@ describe("kickMainLoop", () => {
     });
   });
 
-  test("it updates the globals when they change", async () => {
-    const t = convexTest(schema, modules);
-    await t.run(async (ctx) => {
-      await kickMainLoop(ctx, "enqueue");
-      const globals = await ctx.db.query("globals").unique();
-      expect(globals).not.toBeNull();
-      assert(globals);
-      expect(globals.maxParallelism).toBe(DEFAULT_MAX_PARALLELISM);
-      expect(globals.logLevel).toBe(DEFAULT_LOG_LEVEL);
-      await kickMainLoop(ctx, "enqueue", {
-        maxParallelism: DEFAULT_MAX_PARALLELISM + 1,
-        logLevel: "ERROR",
-      });
-      const after = await ctx.db.query("globals").unique();
-      expect(after).not.toBeNull();
-      assert(after);
-      expect(after.maxParallelism).toBe(DEFAULT_MAX_PARALLELISM + 1);
-      expect(after.logLevel).toBe("ERROR");
-    });
-  });
-
   test("does not kick when already running", async () => {
     const t = convexTest(schema, modules);
     await t.run(async (ctx) => {
@@ -249,8 +229,7 @@ describe("kickMainLoop", () => {
   test("preserves state between kicks with different sources", async () => {
     const t = convexTest(schema, modules);
     await t.run(async (ctx) => {
-      // Initial kick with custom config
-      await kickMainLoop(ctx, "enqueue", {
+      await getOrUpdateGlobals(ctx, {
         maxParallelism: 5,
         logLevel: "ERROR",
       });

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/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.
@@ -64,7 +64,7 @@ export default defineSchema({
     fnName: v.string(),
     fnArgs: v.any(),
     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 +80,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(),
   })