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

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(),
   })

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 = {
   /**