@convex-dev/workpool 0.1.2 → 0.2.0-alpha.0

package/src/client/index.ts

@@ -1,121 +1,285 @@
 import {
   createFunctionHandle,
   DefaultFunctionArgs,
-  Expand,
   FunctionReference,
   FunctionVisibility,
-  GenericDataModel,
-  GenericMutationCtx,
-  GenericQueryCtx,
   getFunctionName,
 } from "convex/server";
-import { GenericId } from "convex/values";
-import { api } from "../component/_generated/api";
-import { LogLevel } from "../component/logging";
-import { completionStatus, type CompletionStatus } from "../component/schema";
-export { completionStatus, type CompletionStatus };
+import { v, VString } from "convex/values";
+import { Mounts } from "../component/_generated/api.js";
+import {
+  OnComplete,
+  runResult as runResultValidator,
+  RunResult,
+  type RetryBehavior,
+  OnCompleteArgs as SharedOnCompleteArgs,
+  Status,
+  Config,
+} from "../component/shared.js";
+import { type LogLevel, logLevel } from "../component/logging.js";
+import { RunMutationCtx, RunQueryCtx, UseApi } from "./utils.js";
+import { DEFAULT_LOG_LEVEL } from "../component/logging.js";
+import { DEFAULT_MAX_PARALLELISM } from "../component/kick.js";
+export { runResultValidator, type RunResult };
 
-export type WorkId = string;
+// Attempts will run with delay [0, 250, 500, 1000, 2000] (ms)
+export const DEFAULT_RETRY_BEHAVIOR: RetryBehavior = {
+  maxAttempts: 5,
+  initialBackoffMs: 250,
+  base: 2,
+};
+export type WorkId = string & { __isWorkId: true };
+export const workIdValidator = v.string() as VString<WorkId>;
 
 export class Workpool {
+  /**
+   * Initializes a Workpool.
+   *
+   * Note: if you want different pools, you need to *create different instances*
+   * of Workpool in convex.config.ts. It isn't sufficient to have different
+   * instances of this class.
+   *
+   * @param component - The component to use, like `components.workpool` from
+   *   `./_generated/api.ts`.
+   * @param options - The options for the Workpool.
+   */
   constructor(
-    private component: UseApi<typeof api>,
+    private component: UseApi<Mounts>, // UseApi<api> for jump to definition
     private options: {
       /** How many actions/mutations can be running at once within this pool.
        * Min 1, Max 300.
        */
-      maxParallelism: number;
-      /** How much to log.
-       * Default WARN.
+      maxParallelism?: number;
+      /** How much to log. This is updated on each call to `enqueue*`,
+       * `status`, or `cancel*`.
+       * Default is WARN.
        * With INFO, you can see events for started and completed work, which can
-       * be parsed.
+       * 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;
-      /** How long to keep completed work in the database, for access by `status`.
-       * Default 1 day.
+      /** Default retry behavior for enqueued actions. */
+      defaultRetryBehavior?: RetryBehavior;
+      /** Whether to retry actions that fail by default. Default: false.
+       * NOTE: Only do this if your actions are idempotent.
+       * See the docs (README.md) for more details.
        */
-      statusTtl?: number;
+      retryActionsByDefault?: boolean;
     }
   ) {}
+  /**
+   * Enqueues an action to be run.
+   *
+   * @param ctx - The mutation or action context that can call ctx.runMutation.
+   * @param fn - The action to run, like `internal.example.myAction`.
+   * @param fnArgs - The arguments to pass to the action.
+   * @param options - The options for the action to specify retry behavior,
+   *   onComplete handling, and scheduling via `runAt` or `runAfter`.
+   * @returns The ID of the work that was enqueued.
+   */
   async enqueueAction<Args extends DefaultFunctionArgs, ReturnType>(
     ctx: RunMutationCtx,
     fn: FunctionReference<"action", FunctionVisibility, Args, ReturnType>,
-    fnArgs: Args
+    fnArgs: Args,
+    options?: {
+      /** 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;
+    } & CallbackOptions &
+      SchedulerOptions
   ): Promise<WorkId> {
-    const fnHandle = await createFunctionHandle(fn);
+    const retryBehavior = getRetryBehavior(
+      this.options.defaultRetryBehavior,
+      this.options.retryActionsByDefault,
+      options?.retry
+    );
+    const onComplete: OnComplete | undefined = options?.onComplete
+      ? {
+          fnHandle: await createFunctionHandle(options.onComplete),
+          context: options.context,
+        }
+      : undefined;
     const id = await ctx.runMutation(this.component.lib.enqueue, {
-      fnHandle,
-      fnName: getFunctionName(fn),
+      ...(await defaultEnqueueArgs(fn, this.options)),
       fnArgs,
       fnType: "action",
-      options: this.options,
+      runAt: getRunAt(options),
+      onComplete,
+      retryBehavior,
     });
     return id as WorkId;
   }
+
+  /**
+   * Enqueues a mutation to be run.
+   *
+   * Note: mutations are not retried by the workpool. Convex automatically
+   * retries them on database conflicts and transient failures.
+   * Because they're deterministic, external retries don't provide any benefit.
+   *
+   * @param ctx - The mutation or action context that can call ctx.runMutation.
+   * @param fn - The mutation to run, like `internal.example.myMutation`.
+   * @param fnArgs - The arguments to pass to the mutation.
+   * @param options - The options for the mutation to specify onComplete handling
+   *   and scheduling via `runAt` or `runAfter`.
+   */
   async enqueueMutation<Args extends DefaultFunctionArgs, ReturnType>(
     ctx: RunMutationCtx,
     fn: FunctionReference<"mutation", FunctionVisibility, Args, ReturnType>,
-    fnArgs: Args
+    fnArgs: Args,
+    options?: CallbackOptions & SchedulerOptions
   ): Promise<WorkId> {
-    const fnHandle = await createFunctionHandle(fn);
     const id = await ctx.runMutation(this.component.lib.enqueue, {
-      fnHandle,
-      fnName: getFunctionName(fn),
+      ...(await defaultEnqueueArgs(fn, this.options)),
       fnArgs,
       fnType: "mutation",
-      options: this.options,
+      runAt: getRunAt(options),
     });
     return id as WorkId;
   }
   async cancel(ctx: RunMutationCtx, id: WorkId): Promise<void> {
-    await ctx.runMutation(this.component.lib.cancel, { id });
+    await ctx.runMutation(this.component.lib.cancel, {
+      id,
+      logLevel: this.options.logLevel ?? getDefaultLogLevel(),
+    });
+  }
+  async cancelAll(ctx: RunMutationCtx): Promise<void> {
+    await ctx.runMutation(this.component.lib.cancelAll, {
+      logLevel: this.options.logLevel ?? getDefaultLogLevel(),
+    });
   }
-  async status(
-    ctx: RunQueryCtx,
-    id: WorkId
-  ): Promise<
-    | { kind: "pending" }
-    | { kind: "inProgress" }
-    | { kind: "completed"; completionStatus: CompletionStatus }
-  > {
-    return await ctx.runQuery(this.component.lib.status, { id });
+  async status(ctx: RunQueryCtx, id: WorkId): Promise<Status> {
+    return ctx.runQuery(this.component.lib.status, { id });
   }
 }
 
-/* Type utils follow */
+function getRetryBehavior(
+  defaultRetryBehavior: RetryBehavior | undefined,
+  retryActionsByDefault: boolean | undefined,
+  retryOverride: boolean | RetryBehavior | undefined
+): RetryBehavior | undefined {
+  const defaultRetry = defaultRetryBehavior ?? DEFAULT_RETRY_BEHAVIOR;
+  const retryByDefault = retryActionsByDefault ?? false;
+  if (retryOverride === true) {
+    return defaultRetry;
+  }
+  if (retryOverride === false) {
+    return undefined;
+  }
+  return retryOverride ?? (retryByDefault ? defaultRetry : undefined);
+}
 
-type RunQueryCtx = {
-  runQuery: GenericQueryCtx<GenericDataModel>["runQuery"];
+async function defaultEnqueueArgs(
+  fn: FunctionReference<"action" | "mutation", FunctionVisibility>,
+  { logLevel, maxParallelism }: Partial<Config>
+) {
+  return {
+    fnHandle: await createFunctionHandle(fn),
+    fnName: getFunctionName(fn),
+    config: {
+      logLevel: logLevel ?? getDefaultLogLevel(),
+      maxParallelism: maxParallelism ?? DEFAULT_MAX_PARALLELISM,
+    },
+  };
+}
+
+export type SchedulerOptions =
+  | {
+      /**
+       * The time (ms since epoch) to run the action at.
+       * If not provided, the action will be run as soon as possible.
+       * Note: this is advisory only. It may run later.
+       */
+      runAt?: number;
+    }
+  | {
+      /**
+       * The number of milliseconds to run the action after.
+       * If not provided, the action will be run as soon as possible.
+       * Note: this is advisory only. It may run later.
+       */
+      runAfter?: number;
+    };
+
+export type CallbackOptions = {
+  /**
+   * A mutation to run after the function succeeds, fails, or is canceled.
+   * The context type is for your use, feel free to provide a validator for it.
+   * e.g.
+   * ```ts
+   * export const completion = internalMutation({
+   *  args: {
+   *    workId: workIdValidator,
+   *    context: v.any(),
+   *    result: runResult,
+   *  },
+   *  handler: async (ctx, args) => {
+   *    console.log(args.result, "Got Context back -> ", args.context, Date.now() - args.context);
+   *  },
+   * });
+   * ```
+   */
+  onComplete?: FunctionReference<
+    "mutation",
+    FunctionVisibility,
+    OnCompleteArgs
+  > | null;
+
+  /**
+   * A context object to pass to the `onComplete` mutation.
+   * Useful for passing data from the enqueue site to the onComplete site.
+   */
+  context?: unknown;
 };
-type RunMutationCtx = {
-  runMutation: GenericMutationCtx<GenericDataModel>["runMutation"];
+
+export type OnCompleteArgs = {
+  /**
+   * The ID of the work that completed.
+   */
+  workId: WorkId;
+  /**
+   * The context object passed when enqueuing the work.
+   * Useful for passing data from the enqueue site to the onComplete site.
+   */
+  context: unknown;
+  /**
+   * The result of the run that completed.
+   */
+  result: RunResult;
 };
+// ensure OnCompleteArgs satisfies SharedOnCompleteArgs
+const _ = {} as OnCompleteArgs satisfies SharedOnCompleteArgs;
 
-export type OpaqueIds<T> =
-  T extends GenericId<infer _T>
-    ? string
-    : T extends (infer U)[]
-      ? OpaqueIds<U>[]
-      : T extends object
-        ? { [K in keyof T]: OpaqueIds<T[K]> }
-        : T;
+function getRunAt(options?: SchedulerOptions): number {
+  if (!options) {
+    return Date.now();
+  }
+  if ("runAt" in options && options.runAt !== undefined) {
+    return options.runAt;
+  }
+  if ("runAfter" in options && options.runAfter !== undefined) {
+    return Date.now() + options.runAfter;
+  }
+  return Date.now();
+}
 
-export type UseApi<API> = Expand<{
-  [mod in keyof API]: API[mod] extends FunctionReference<
-    infer FType,
-    "public",
-    infer FArgs,
-    infer FReturnType,
-    infer FComponentPath
-  >
-    ? FunctionReference<
-        FType,
-        "internal",
-        OpaqueIds<FArgs>,
-        OpaqueIds<FReturnType>,
-        FComponentPath
-      >
-    : UseApi<API[mod]>;
-}>;
+function getDefaultLogLevel(): LogLevel {
+  if (process.env.WORKPOOL_LOG_LEVEL) {
+    if (
+      !logLevel.members
+        .map((m) => m.value as string)
+        .includes(process.env.WORKPOOL_LOG_LEVEL)
+    ) {
+      console.warn(
+        `Invalid log level (${process.env.WORKPOOL_LOG_LEVEL}), defaulting to "INFO"`
+      );
+    } else {
+      return process.env.WORKPOOL_LOG_LEVEL as LogLevel;
+    }
+  }
+  return DEFAULT_LOG_LEVEL;
+}

package/src/client/utils.ts

@@ -0,0 +1,45 @@
+import {
+  Expand,
+  FunctionReference,
+  GenericMutationCtx,
+  GenericQueryCtx,
+} from "convex/server";
+import { GenericId } from "convex/values";
+
+import { GenericDataModel } from "convex/server";
+
+/* Type utils follow */
+
+export type RunQueryCtx = {
+  runQuery: GenericQueryCtx<GenericDataModel>["runQuery"];
+};
+export type RunMutationCtx = {
+  runMutation: GenericMutationCtx<GenericDataModel>["runMutation"];
+};
+
+export type OpaqueIds<T> =
+  T extends GenericId<infer _T>
+    ? string
+    : T extends (infer U)[]
+      ? OpaqueIds<U>[]
+      : T extends object
+        ? { [K in keyof T]: OpaqueIds<T[K]> }
+        : T;
+
+export type UseApi<API> = Expand<{
+  [mod in keyof API]: API[mod] extends FunctionReference<
+    infer FType,
+    "public",
+    infer FArgs,
+    infer FReturnType,
+    infer FComponentPath
+  >
+    ? FunctionReference<
+        FType,
+        "internal",
+        OpaqueIds<FArgs>,
+        OpaqueIds<FReturnType>,
+        FComponentPath
+      >
+    : UseApi<API[mod]>;
+}>;

package/src/component/README.md

@@ -0,0 +1,73 @@
+# Workpool: implementation notes and high-level architecture
+
+Concepts:
+
+- `segment`: A slice of time to process work. All work is bucketed into one.
+  This enables us to batch work and avoid database conflicts.
+- `generation`: A monotonically increasing counter to ensure the loop is
+  only running one instance. If two loops start with the same generation,
+  one will successfully increase it, the other will retry and find that the
+  generation has changed and fail out.
+- "Retention" is used to refer to situations where a query might have to read
+  over a lot of "tombstones" - deleted data that hasn't been vacuumed from the
+  underlying database yet. If there are frequent deletions, scanning across them
+  can delay a query. Because of our delete-heavy queuing strategy, we have to be
+  careful. Strategies are below.
+- Cursors: A pointer to the last processed place in a table. In our case, they
+  might allow data to be written before them if out-of-order writes happen, so
+  we need to account for finding those "missed" writes on some granularity. We
+  choose to wait until there isn't any immediate work to do before those scans.
+  They help avoid retention issues.
+
+## Data state machine
+
+```mermaid
+flowchart LR
+    Client -->|enqueue| pendingStart
+    Client -->|cancel| pendingCancelation
+    complete --> |success or failure| pendingCompletion
+    pendingCompletion -->|retry| pendingStart
+    pendingStart --> workerRunning["worker running"]
+    workerRunning -->|worker finished| complete
+    workerRunning --> |recovery| complete
+    successfulCancel["AND"]@{shape: delay} --> |canceled| complete
+    pendingStart --> successfulCancel
+    pendingCancelation --> successfulCancel
+```
+
+Notably:
+
+- The pending\* states are written by outside sources.
+- The main loop federates changes to/from "running"
+- Canceling only impacts pending and retrying jobs.
+
+## Loop state machine
+
+```mermaid
+flowchart TD
+    idle -->|enqueue| running
+    running-->|"all started, leftover capacity"| scheduled
+    scheduled -->|"enqueue, cancel, saveResult, recovery"| running
+    running -->|"maxed out"| saturated
+    saturated -->|"cancel, saveResult, recovery"| running
+    running-->|"all done"| idle
+```
+
+- While the loop is running, the runStatus doesn't change, making it safer to
+  read from clients without database conflicts.
+- The "saturated" state is concretely "running" or "scheduled" at max
+  parallelism. There is a boolean set on "scheduled" to avoid clients from
+  kicking the main loop on enqueueing, which is unlikely to be productive, since
+  the next action needs to be something terminating.
+
+## Retention optimization strategy
+
+- Producers (Client, Worker, Recovery) write to a future "segment".
+- Consumers (main) read the current segment.
+  - On conflicts, producers will write to progressively higher segments, while
+    the main loop will continue to read the segment originally called with.
+    This means conflicts are less likely on each retry.
+- Patch singletons to avoid tombstones.
+- Use segements & cursors to bound reads to latest data.
+  - Do scans outside of the critical path (during load).
+- Do point reads otherwise.

package/src/component/_generated/api.d.ts

@@ -8,9 +8,15 @@
  * @module
  */
 
+import type * as complete from "../complete.js";
+import type * as kick from "../kick.js";
 import type * as lib from "../lib.js";
 import type * as logging from "../logging.js";
+import type * as loop from "../loop.js";
+import type * as recovery from "../recovery.js";
+import type * as shared from "../shared.js";
 import type * as stats from "../stats.js";
+import type * as worker from "../worker.js";
 
 import type {
   ApiFromModules,
@@ -26,27 +32,49 @@ import type {
  * ```
  */
 declare const fullApi: ApiFromModules<{
+  complete: typeof complete;
+  kick: typeof kick;
   lib: typeof lib;
   logging: typeof logging;
+  loop: typeof loop;
+  recovery: typeof recovery;
+  shared: typeof shared;
   stats: typeof stats;
+  worker: typeof worker;
 }>;
 export type Mounts = {
   lib: {
-    cancel: FunctionReference<"mutation", "public", { id: string }, any>;
-    cleanup: FunctionReference<"mutation", "public", { maxAgeMs: number }, any>;
+    cancel: FunctionReference<
+      "mutation",
+      "public",
+      { id: string; logLevel: "DEBUG" | "INFO" | "WARN" | "ERROR" },
+      any
+    >;
+    cancelAll: FunctionReference<
+      "mutation",
+      "public",
+      { before?: number; logLevel: "DEBUG" | "INFO" | "WARN" | "ERROR" },
+      any
+    >;
     enqueue: FunctionReference<
       "mutation",
       "public",
       {
+        config: {
+          logLevel: "DEBUG" | "INFO" | "WARN" | "ERROR";
+          maxParallelism: number;
+        };
         fnArgs: any;
         fnHandle: string;
         fnName: string;
         fnType: "action" | "mutation";
-        options: {
-          logLevel?: "DEBUG" | "INFO" | "WARN" | "ERROR";
-          maxParallelism: number;
-          statusTtl?: number;
+        onComplete?: { context?: any; fnHandle: string };
+        retryBehavior?: {
+          base: number;
+          initialBackoffMs: number;
+          maxAttempts: number;
         };
+        runAt: number;
       },
       string
     >;
@@ -54,14 +82,10 @@ export type Mounts = {
       "query",
       "public",
       { id: string },
-      | { kind: "pending" }
-      | { kind: "inProgress" }
-      | {
-          completionStatus: "success" | "error" | "canceled" | "timeout";
-          kind: "completed";
-        }
+      | { previousAttempts: number; state: "pending" }
+      | { previousAttempts: number; state: "running" }
+      | { state: "finished" }
     >;
-    stopCleanup: FunctionReference<"mutation", "public", {}, any>;
   };
 };
 // For now fullApiWithMounts is only fullApi which provides
@@ -78,56 +102,4 @@ export declare const internal: FilterApi<
   FunctionReference<any, "internal">
 >;
 
-export declare const components: {
-  crons: {
-    public: {
-      del: FunctionReference<
-        "mutation",
-        "internal",
-        { identifier: { id: string } | { name: string } },
-        null
-      >;
-      get: FunctionReference<
-        "query",
-        "internal",
-        { identifier: { id: string } | { name: string } },
-        {
-          args: Record<string, any>;
-          functionHandle: string;
-          id: string;
-          name?: string;
-          schedule:
-            | { kind: "interval"; ms: number }
-            | { cronspec: string; kind: "cron" };
-        } | null
-      >;
-      list: FunctionReference<
-        "query",
-        "internal",
-        {},
-        Array<{
-          args: Record<string, any>;
-          functionHandle: string;
-          id: string;
-          name?: string;
-          schedule:
-            | { kind: "interval"; ms: number }
-            | { cronspec: string; kind: "cron" };
-        }>
-      >;
-      register: FunctionReference<
-        "mutation",
-        "internal",
-        {
-          args: Record<string, any>;
-          functionHandle: string;
-          name?: string;
-          schedule:
-            | { kind: "interval"; ms: number }
-            | { cronspec: string; kind: "cron" };
-        },
-        string
-      >;
-    };
-  };
-};
+export declare const components: {};