@convex-dev/workpool 0.2.14 → 0.2.16

package/package.json

@@ -7,7 +7,7 @@
     "email": "support@convex.dev",
     "url": "https://github.com/get-convex/workpool/issues"
   },
-  "version": "0.2.14",
+  "version": "0.2.16",
   "license": "Apache-2.0",
   "keywords": [
     "convex",
@@ -31,7 +31,7 @@
     "test:debug": "vitest --inspect-brk --no-file-parallelism",
     "test:coverage": "vitest run --coverage --coverage.reporter=text",
     "alpha": "rm -rf dist && npm run build && npm run test && npm version prerelease --preid alpha && npm publish --tag alpha && git push --tags",
-    "release": "rm -rf dist && npm run build && npm run test && npm version patch && npm publish && git push --tags",
+    "release": "rm -rf dist && npm run build && npm run test && npm version patch && npm publish && git push --tags && git push",
     "prepare": "npm run build",
     "prepack": "node node10stubs.mjs",
     "postpack": "node node10stubs.mjs --cleanup"
@@ -72,9 +72,10 @@
     "@eslint/js": "^9.9.1",
     "@types/node": "^18.17.0",
     "@vitest/coverage-v8": "^2.1.9",
-    "convex-test": "^0.0.36-alpha.0",
+    "convex-test": "^0.0.38-alpha.0",
     "eslint": "^9.9.1",
     "globals": "^15.9.0",
+    "pkg-pr-new": "^0.0.54",
     "prettier": "3.2.5",
     "typescript": "^5.8.3",
     "typescript-eslint": "^8.4.0",

package/src/client/index.ts

@@ -5,29 +5,45 @@ import {
   FunctionReference,
   FunctionType,
   FunctionVisibility,
+  GenericDataModel,
+  GenericMutationCtx,
   getFunctionName,
+  internalMutationGeneric,
+  RegisteredMutation,
 } from "convex/server";
-import { v, VString } from "convex/values";
+import { Infer, v, Validator, VAny, VString } from "convex/values";
 import { Mounts } from "../component/_generated/api.js";
 import { DEFAULT_LOG_LEVEL, type LogLevel } from "../component/logging.js";
 import {
   Config,
   DEFAULT_MAX_PARALLELISM,
   OnComplete,
-  runResult as resultValidator,
+  vResultValidator,
   type RetryBehavior,
   RunResult,
   OnCompleteArgs as SharedOnCompleteArgs,
   Status,
 } from "../component/shared.js";
 import { RunMutationCtx, RunQueryCtx, UseApi } from "./utils.js";
-export { resultValidator, type RunResult, type RetryBehavior, type OnComplete };
+export {
+  vResultValidator,
+  type RunResult,
+  type RetryBehavior,
+  type OnComplete,
+};
 export {
   retryBehavior as vRetryBehavior,
   onComplete as vOnComplete,
 } from "../component/shared.js";
 export { logLevel as vLogLevel, type LogLevel } from "../component/logging.js";
-export { resultValidator as vResultValidator };
+export type WorkId = string & { __isWorkId: true };
+export const vWorkIdValidator = v.string() as VString<WorkId>;
+export {
+  /** @deprecated Use `vWorkIdValidator` instead. */
+  vWorkIdValidator as workIdValidator,
+  /** @deprecated Use `vResultValidator` instead. */
+  vResultValidator as resultValidator,
+};
 
 // Attempts will run with delay [0, 250, 500, 1000, 2000] (ms)
 export const DEFAULT_RETRY_BEHAVIOR: RetryBehavior = {
@@ -35,56 +51,6 @@ export const DEFAULT_RETRY_BEHAVIOR: RetryBehavior = {
   initialBackoffMs: 250,
   base: 2,
 };
-export type WorkId = string & { __isWorkId: true };
-export const workIdValidator = v.string() as VString<WorkId>;
-export { workIdValidator as vWorkIdValidator };
-
-export type NameOption = {
-  /**
-   * The name of the function. By default, if you pass in api.foo.bar.baz,
-   * it will use "foo/bar:baz" as the name. If you pass in a function handle,
-   * it will use the function handle directly.
-   */
-  name?: string;
-};
-
-export type RetryOption = {
-  /** Whether to retry the action if it fails.
-   * If true, it will use the default retry behavior.
-   * If custom behavior is provided, it will retry using that behavior.
-   * If unset, it will use the Workpool's configured default.
-   */
-  retry?: boolean | RetryBehavior;
-};
-
-export type WorkpoolOptions = {
-  /** How many actions/mutations can be running at once within this pool.
-   * Min 1, Max 300.
-   */
-  maxParallelism?: number;
-  /** How much to log. This is updated on each call to `enqueue*`,
-   * `status`, or `cancel*`.
-   * Default is REPORT, which logs warnings, errors, and a periodic report.
-   * With INFO, you can also see events for started and completed work.
-   * Stats generated can be parsed by tools like
-   * [Axiom](https://axiom.co) for monitoring.
-   * With DEBUG, you can see timers and internal events for work being
-   * scheduled.
-   */
-  logLevel?: LogLevel;
-} & WorkpoolRetryOptions;
-
-export type WorkpoolRetryOptions = {
-  /** Default retry behavior for enqueued actions.
-   * See {@link RetryBehavior}.
-   */
-  defaultRetryBehavior?: RetryBehavior;
-  /** Whether to retry actions that fail by default. Default: false.
-   * NOTE: Only enable this if your actions are idempotent.
-   * See the docs (README.md) for more details.
-   */
-  retryActionsByDefault?: boolean;
-};
 
 export class Workpool {
   /**
@@ -233,8 +199,122 @@ export class Workpool {
   async status(ctx: RunQueryCtx, id: WorkId): Promise<Status> {
     return ctx.runQuery(this.component.lib.status, { id });
   }
+
+  /**
+   * Defines a mutation that will be run after a work item completes.
+   * You can pass this to a call to enqueue* like so:
+   * ```ts
+   * export const myOnComplete = workpool.defineOnComplete({
+   *   context: v.literal("myContextValue"), // optional
+   *   handler: async (ctx, {workId, context, result}) => {
+   *     // ... do something with the result
+   *   },
+   * });
+   *
+   * // in some other function:
+   * const workId = await workpool.enqueueAction(ctx, internal.foo.bar, {
+   *   // ... args to action
+   * }, {
+   *   onComplete: internal.foo.myOnComplete,
+   * });
+   * ```
+   */
+  defineOnComplete<
+    DataModel extends GenericDataModel,
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    V extends Validator<any, "required", any> = VAny,
+  >({
+    context,
+    handler,
+  }: {
+    context?: V;
+    handler: (
+      ctx: GenericMutationCtx<DataModel>,
+      args: {
+        workId: WorkId;
+        context: Infer<V>;
+        result: RunResult;
+      }
+    ) => Promise<void>;
+  }): RegisteredMutation<"internal", OnCompleteArgs, null> {
+    return internalMutationGeneric({
+      args: vOnCompleteValidator(context),
+      handler,
+    });
+  }
+}
+
+/**
+ * Returns a validator to use for the onComplete mutation.
+ * To be used like:
+ * ```ts
+ * export const myOnComplete = internalMutation({
+ *   args: vOnCompleteValidator(v.string()),
+ *   handler: async (ctx, {workId, context, result}) => {
+ *     // context has been validated as a string
+ *     // ... do something with the result
+ *   },
+ * });
+ * @param context - The context validator. If not provided, it will be `v.any()`.
+ * @returns The validator for the onComplete mutation.
+ */
+export function vOnCompleteValidator<
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  V extends Validator<any, "required", any> = VAny,
+>(context?: V) {
+  return v.object({
+    workId: vWorkIdValidator,
+    context: context ?? v.any(),
+    result: vResultValidator,
+  });
 }
 
+export type NameOption = {
+  /**
+   * The name of the function. By default, if you pass in api.foo.bar.baz,
+   * it will use "foo/bar:baz" as the name. If you pass in a function handle,
+   * it will use the function handle directly.
+   */
+  name?: string;
+};
+
+export type RetryOption = {
+  /** Whether to retry the action if it fails.
+   * If true, it will use the default retry behavior.
+   * If custom behavior is provided, it will retry using that behavior.
+   * If unset, it will use the Workpool's configured default.
+   */
+  retry?: boolean | RetryBehavior;
+};
+
+export type WorkpoolOptions = {
+  /** How many actions/mutations can be running at once within this pool.
+   * Min 1, Max 300.
+   */
+  maxParallelism?: number;
+  /** How much to log. This is updated on each call to `enqueue*`,
+   * `status`, or `cancel*`.
+   * Default is REPORT, which logs warnings, errors, and a periodic report.
+   * With INFO, you can also see events for started and completed work.
+   * Stats generated can be parsed by tools like
+   * [Axiom](https://axiom.co) for monitoring.
+   * With DEBUG, you can see timers and internal events for work being
+   * scheduled.
+   */
+  logLevel?: LogLevel;
+} & WorkpoolRetryOptions;
+
+export type WorkpoolRetryOptions = {
+  /** Default retry behavior for enqueued actions.
+   * See {@link RetryBehavior}.
+   */
+  defaultRetryBehavior?: RetryBehavior;
+  /** Whether to retry actions that fail by default. Default: false.
+   * NOTE: Only enable this if your actions are idempotent.
+   * See the docs (README.md) for more details.
+   */
+  retryActionsByDefault?: boolean;
+};
 export type SchedulerOptions =
   | {
       /**
@@ -259,12 +339,18 @@ export type CallbackOptions = {
    * The context type is for your use, feel free to provide a validator for it.
    * e.g.
    * ```ts
+   * export const completion = workpool.defineOnComplete({
+   *   context: v.string(),
+   *   handler: async (ctx, {workId, context, result}) => {
+   *     // context has been validated as a string
+   *     // ... do something with the result
+   *   },
+   * });
+   * ```
+   * or more manually:
+   * ```ts
    * export const completion = internalMutation({
-   *  args: {
-   *    workId: workIdValidator,
-   *    context: v.any(),
-   *    result: resultValidator,
-   *  },
+   *  args: vOnCompleteValidator(v.string()),
    *  handler: async (ctx, args) => {
    *    console.log(args.result, "Got Context back -> ", args.context, Date.now() - args.context);
    *  },

package/src/component/complete.ts

@@ -4,7 +4,7 @@ import { Id } from "./_generated/dataModel.js";
 import { internalMutation, MutationCtx } from "./_generated/server.js";
 import { kickMainLoop } from "./kick.js";
 import { createLogger } from "./logging.js";
-import { OnCompleteArgs, RunResult, runResult } from "./shared.js";
+import { OnCompleteArgs, RunResult, vResultValidator } from "./shared.js";
 import { recordCompleted } from "./stats.js";
 
 export type CompleteJob = Infer<typeof completeArgs.fields.jobs.element>;
@@ -12,7 +12,7 @@ export type CompleteJob = Infer<typeof completeArgs.fields.jobs.element>;
 export const completeArgs = v.object({
   jobs: v.array(
     v.object({
-      runResult: runResult,
+      runResult: vResultValidator,
       workId: v.id("work"),
       attempt: v.number(),
     })

package/src/component/kick.test.ts

@@ -9,7 +9,6 @@ import {
   vi,
 } from "vitest";
 import { internal } from "./_generated/api";
-import { Id } from "./_generated/dataModel.js";
 import { kickMainLoop } from "./kick.js";
 import { DEFAULT_LOG_LEVEL } from "./logging.js";
 import schema from "./schema.js";
@@ -276,13 +275,18 @@ describe("kickMainLoop", () => {
       const runStatus = await ctx.db.query("runStatus").unique();
       assert(runStatus);
       const segment = getNextSegment() + 10n;
+      const scheduledId = await ctx.scheduler.runAfter(
+        10_000,
+        internal.loop.main,
+        { generation: 0n, segment }
+      );
       await ctx.db.patch(runStatus._id, {
         state: {
           generation: 0n,
           saturated: false,
           kind: "scheduled",
           segment,
-          scheduledId: "" as Id<"_scheduled_functions">,
+          scheduledId,
         },
       });
       // await all scheduled functions to run
@@ -291,6 +295,9 @@ describe("kickMainLoop", () => {
       assert(afterStatus);
       expect(afterStatus.state.kind).toBe("running");
       assert(afterStatus.state.kind === "running");
+      const scheduledJob = await ctx.db.system.get(scheduledId);
+      assert(scheduledJob);
+      expect(scheduledJob.state.kind).toBe("canceled");
     });
   });
 });

package/src/component/loop.test.ts

@@ -17,7 +17,6 @@ import {
   DEFAULT_MAX_PARALLELISM,
   getCurrentSegment,
   getNextSegment,
-  HOUR,
   toSegment,
 } from "./shared";
 import { DEFAULT_LOG_LEVEL } from "./logging";

package/src/component/schema.ts

@@ -5,7 +5,7 @@ import {
   config,
   onComplete,
   retryBehavior,
-  runResult,
+  vResultValidator,
 } from "./shared.js";
 
 // Represents a slice of time to process work.
@@ -80,7 +80,7 @@ export default defineSchema({
   // Written by complete, read & deleted by `main`.
   pendingCompletion: defineTable({
     segment,
-    runResult,
+    runResult: vResultValidator,
     workId: v.id("work"),
     retry: v.boolean(),
   })

package/src/component/shared.ts

@@ -64,7 +64,7 @@ export type RetryBehavior = {
 // This ensures that the type satisfies the schema.
 const _ = {} as RetryBehavior satisfies Infer<typeof retryBehavior>;
 
-export const runResult = v.union(
+export const vResultValidator = v.union(
   v.object({
     kind: v.literal("success"),
     returnValue: v.any(),
@@ -77,7 +77,7 @@ export const runResult = v.union(
     kind: v.literal("canceled"),
   })
 );
-export type RunResult = Infer<typeof runResult>;
+export type RunResult = Infer<typeof vResultValidator>;
 
 export const onComplete = v.object({
   fnHandle: v.string(), // mutation

package/src/component/stats.ts

@@ -15,8 +15,6 @@ import { internal } from "./_generated/api.js";
 import schema from "./schema.js";
 import { paginator } from "convex-helpers/server/pagination";
 
-const BACKLOG_BATCH_SIZE = 100;
-
 /**
  * Record stats about work execution. Intended to be queried by Axiom or Datadog.
  * See the [README](https://github.com/get-convex/workpool) for example queries.
@@ -114,36 +112,15 @@ export const calculateBacklogAndReport = internalMutation({
     logLevel,
   },
   handler: async (ctx, args) => {
-    const pendingStart = await paginator(ctx.db, schema)
-      .query("pendingStart")
-      .withIndex("segment", (q) =>
-        q.gte("segment", args.startSegment).lt("segment", args.endSegment)
-      )
-      .paginate({
-        numItems: BACKLOG_BATCH_SIZE,
-        cursor: args.cursor,
-      });
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    const pendingStart = await (ctx.db.query("pendingStart") as any).count();
+
     const console = createLogger(args.logLevel);
-    if (pendingStart.isDone) {
-      recordReport(console, {
-        ...args.report,
-        running: args.running,
-        backlog: pendingStart.page.length,
-      });
-    } else {
-      await ctx.scheduler.runAfter(
-        0,
-        internal.stats.calculateBacklogAndReport,
-        {
-          startSegment: args.startSegment,
-          endSegment: args.endSegment,
-          cursor: pendingStart.continueCursor,
-          report: args.report,
-          running: args.running,
-          logLevel: args.logLevel,
-        }
-      );
-    }
+    recordReport(console, {
+      ...args.report,
+      running: args.running,
+      backlog: pendingStart,
+    });
   },
 });