@convex-dev/workpool 0.2.18-alpha.2 → 0.2.18

package/src/client/index.ts

@@ -40,7 +40,7 @@ export {
   /** @deprecated Use `vResultValidator` instead. */
   vResultValidator as resultValidator,
 };
-/** @deprecated Use `vOnCompleteArgs(<your-context-validator>)` instead. */
+/** Equivalent to `vOnCompleteArgs(<your-context-validator>)`. */
 export const vOnComplete = vOnCompleteArgs(v.any());
 /** @deprecated Use `vOnCompleteArgs()` instead. */
 export const vOnCompleteValidator = vOnCompleteArgs;
@@ -52,6 +52,9 @@ export const DEFAULT_RETRY_BEHAVIOR: RetryBehavior = {
   base: 2,
 };
 
+// UseApi<api> for jump to definition
+export type WorkpoolComponent = UseApi<Mounts>;
+
 export class Workpool {
   /**
    * Initializes a Workpool.
@@ -65,9 +68,10 @@ export class Workpool {
    * @param options - The {@link WorkpoolOptions} for the Workpool.
    */
   constructor(
-    private component: UseApi<Mounts>, // UseApi<api> for jump to definition
+    public component: WorkpoolComponent,
     public options: WorkpoolOptions
   ) {}
+
   /**
    * Enqueues an action to be run.
    *
@@ -82,28 +86,48 @@ export class Workpool {
     ctx: RunMutationCtx,
     fn: FunctionReference<"action", FunctionVisibility, Args, ReturnType>,
     fnArgs: Args,
-    options?: RetryOption & CallbackOptions & SchedulerOptions & NameOption
+    options?: RetryOption & EnqueueOptions
   ): Promise<WorkId> {
     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, {
-      ...(await defaultEnqueueArgs(fn, options?.name, this.options)),
-      fnArgs,
-      fnType: "action",
-      runAt: getRunAt(options),
-      onComplete,
+    return enqueue(this.component, ctx, "action", fn, fnArgs, {
+      retryBehavior,
+      ...this.options,
+      ...options,
+    });
+  }
+
+  /**
+   * Enqueues a batch of actions to be run.
+   * Each action will be run independently, and the onComplete handler will
+   * be called for each action.
+   *
+   * @param ctx - The mutation or action ctx that can call ctx.runMutation.
+   * @param fn - The action to run, like `internal.example.myAction`.
+   * @param argsArray - The arguments to pass to the action.
+   * @param options - The options for the actions to specify retry behavior,
+   *   onComplete handling, and scheduling via `runAt` or `runAfter`.
+   * @returns The IDs of the work that was enqueued.
+   */
+  async enqueueActionBatch<Args extends DefaultFunctionArgs, ReturnType>(
+    ctx: RunMutationCtx,
+    fn: FunctionReference<"action", FunctionVisibility, Args, ReturnType>,
+    argsArray: Array<Args>,
+    options?: RetryOption & EnqueueOptions
+  ): Promise<WorkId[]> {
+    const retryBehavior = getRetryBehavior(
+      this.options.defaultRetryBehavior,
+      this.options.retryActionsByDefault,
+      options?.retry
+    );
+    return enqueueBatch(this.component, ctx, "action", fn, argsArray, {
       retryBehavior,
+      ...this.options,
+      ...options,
     });
-    return id as WorkId;
   }
 
   /**
@@ -123,44 +147,81 @@ export class Workpool {
     ctx: RunMutationCtx,
     fn: FunctionReference<"mutation", FunctionVisibility, Args, ReturnType>,
     fnArgs: Args,
-    options?: CallbackOptions & SchedulerOptions & NameOption
+    options?: EnqueueOptions
   ): Promise<WorkId> {
-    const onComplete: OnComplete | undefined = options?.onComplete
-      ? {
-          fnHandle: await createFunctionHandle(options.onComplete),
-          context: options.context,
-        }
-      : undefined;
-    const id = await ctx.runMutation(this.component.lib.enqueue, {
-      ...(await defaultEnqueueArgs(fn, options?.name, this.options)),
-      fnArgs,
-      fnType: "mutation",
-      runAt: getRunAt(options),
-      onComplete,
+    return enqueue(this.component, ctx, "mutation", fn, fnArgs, {
+      ...this.options,
+      ...options,
+    });
+  }
+  /**
+   * Enqueues a batch of mutations to be run.
+   * Each mutation will be run independently, and the onComplete handler will
+   * be called for each mutation.
+   *
+   * @param ctx - The mutation or action context that can call ctx.runMutation.
+   * @param fn - The mutation to run, like `internal.example.myMutation`.
+   * @param argsArray - The arguments to pass to the mutations.
+   * @param options - The options for the mutations to specify onComplete handling
+   *   and scheduling via `runAt` or `runAfter`.
+   */
+  async enqueueMutationBatch<Args extends DefaultFunctionArgs, ReturnType>(
+    ctx: RunMutationCtx,
+    fn: FunctionReference<"mutation", FunctionVisibility, Args, ReturnType>,
+    argsArray: Array<Args>,
+    options?: EnqueueOptions
+  ): Promise<WorkId[]> {
+    return enqueueBatch(this.component, ctx, "mutation", fn, argsArray, {
+      ...this.options,
+      ...options,
     });
-    return id as WorkId;
   }
 
+  /**
+   * Enqueues a query to be run.
+   * Usually not what you want, but it can be useful during workflows.
+   * The query is run in a mutation and the result is returned to the caller,
+   * so it can conflict if other mutations are writing the value.
+   *
+   * @param ctx - The mutation or action context that can call ctx.runMutation.
+   * @param fn - The query to run, like `internal.example.myQuery`.
+   * @param fnArgs - The arguments to pass to the query.
+   * @param options - The options for the query to specify onComplete handling
+   *   and scheduling via `runAt` or `runAfter`.
+   */
   async enqueueQuery<Args extends DefaultFunctionArgs, ReturnType>(
     ctx: RunMutationCtx,
     fn: FunctionReference<"query", FunctionVisibility, Args, ReturnType>,
     fnArgs: Args,
-    options?: CallbackOptions & SchedulerOptions & NameOption
+    options?: EnqueueOptions
   ): Promise<WorkId> {
-    const onComplete: OnComplete | undefined = options?.onComplete
-      ? {
-          fnHandle: await createFunctionHandle(options.onComplete),
-          context: options.context,
-        }
-      : undefined;
-    const id = await ctx.runMutation(this.component.lib.enqueue, {
-      ...(await defaultEnqueueArgs(fn, options?.name, this.options)),
-      fnArgs,
-      fnType: "query",
-      runAt: getRunAt(options),
-      onComplete,
+    return enqueue(this.component, ctx, "query", fn, fnArgs, {
+      ...this.options,
+      ...options,
+    });
+  }
+
+  /**
+   * Enqueues a batch of queries to be run.
+   * Each query will be run independently, and the onComplete handler will
+   * be called for each query.
+   *
+   * @param ctx - The mutation or action context that can call ctx.runMutation.
+   * @param fn - The query to run, like `internal.example.myQuery`.
+   * @param argsArray - The arguments to pass to the queries.
+   * @param options - The options for the queries to specify onComplete handling
+   *   and scheduling via `runAt` or `runAfter`.
+   */
+  async enqueueQueryBatch<Args extends DefaultFunctionArgs, ReturnType>(
+    ctx: RunMutationCtx,
+    fn: FunctionReference<"query", FunctionVisibility, Args, ReturnType>,
+    argsArray: Array<Args>,
+    options?: EnqueueOptions
+  ): Promise<WorkId[]> {
+    return enqueueBatch(this.component, ctx, "query", fn, argsArray, {
+      ...this.options,
+      ...options,
     });
-    return id as WorkId;
   }
 
   /**
@@ -200,6 +261,17 @@ export class Workpool {
     return ctx.runQuery(this.component.lib.status, { id });
   }
 
+  /**
+   * Gets the status of a batch of work items.
+   *
+   * @param ctx - The query context that can call ctx.runQuery.
+   * @param ids - The IDs of the work to get the status of.
+   * @returns The status of the work items.
+   */
+  async statusBatch(ctx: RunQueryCtx, ids: WorkId[]): Promise<Status[]> {
+    return ctx.runQuery(this.component.lib.statusBatch, { ids });
+  }
+
   /**
    * Defines a mutation that will be run after a work item completes.
    * You can pass this to a call to enqueue* like so:
@@ -264,20 +336,11 @@ export function vOnCompleteArgs<
 >(context?: V) {
   return v.object({
     workId: vWorkIdValidator,
-    context: context ?? v.optional(v.any()),
+    context: (context ?? v.optional(v.any())) as V,
     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.
@@ -315,25 +378,14 @@ export type WorkpoolRetryOptions = {
    */
   retryActionsByDefault?: boolean;
 };
-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 = {
+export type EnqueueOptions = {
+  /**
+   * 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;
   /**
    * 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.
@@ -368,7 +420,24 @@ export type CallbackOptions = {
    * Useful for passing data from the enqueue site to the onComplete site.
    */
   context?: unknown;
-};
+} & (
+  | {
+      /**
+       * 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 OnCompleteArgs = {
   /**
@@ -412,36 +481,105 @@ function getRetryBehavior(
   return retryOverride ?? (retryByDefault ? defaultRetry : undefined);
 }
 
-async function defaultEnqueueArgs(
+async function enqueueArgs(
   fn:
     | FunctionReference<FunctionType, FunctionVisibility>
     | FunctionHandle<FunctionType, DefaultFunctionArgs>,
-  name: string | undefined,
-  { logLevel, maxParallelism }: Partial<Config>
+  opts:
+    | (EnqueueOptions & Partial<Config> & { retryBehavior?: RetryBehavior })
+    | undefined
 ) {
   const [fnHandle, fnName] =
     typeof fn === "string" && fn.startsWith("function://")
-      ? [fn, name ?? fn]
-      : [await createFunctionHandle(fn), name ?? safeFunctionName(fn)];
+      ? [fn, opts?.name ?? fn]
+      : [await createFunctionHandle(fn), opts?.name ?? safeFunctionName(fn)];
+  const onComplete: OnComplete | undefined = opts?.onComplete
+    ? {
+        fnHandle: await createFunctionHandle(opts.onComplete),
+        context: opts.context,
+      }
+    : undefined;
   return {
     fnHandle,
     fnName,
+    onComplete,
+    runAt: getRunAt(opts),
+    retryBehavior: opts?.retryBehavior,
     config: {
-      logLevel: logLevel ?? DEFAULT_LOG_LEVEL,
-      maxParallelism: maxParallelism ?? DEFAULT_MAX_PARALLELISM,
+      logLevel: opts?.logLevel ?? DEFAULT_LOG_LEVEL,
+      maxParallelism: opts?.maxParallelism ?? DEFAULT_MAX_PARALLELISM,
     },
   };
 }
 
-function getRunAt(options?: SchedulerOptions): number {
+function getRunAt(
+  options:
+    | {
+        runAt?: number;
+        runAfter?: number;
+      }
+    | undefined
+): number {
   if (!options) {
     return Date.now();
   }
-  if ("runAt" in options && options.runAt !== undefined) {
+  if (options.runAt !== undefined) {
     return options.runAt;
   }
-  if ("runAfter" in options && options.runAfter !== undefined) {
+  if (options.runAfter !== undefined) {
     return Date.now() + options.runAfter;
   }
   return Date.now();
 }
+
+export async function enqueueBatch<
+  FnType extends FunctionType,
+  Args extends DefaultFunctionArgs,
+  ReturnType,
+>(
+  component: UseApi<Mounts>,
+  ctx: RunMutationCtx,
+  fnType: FnType,
+  fn: FunctionReference<FnType, FunctionVisibility, Args, ReturnType>,
+  fnArgsArray: Array<Args>,
+  options: EnqueueOptions & {
+    retryBehavior?: RetryBehavior;
+    maxParallelism?: number;
+    logLevel?: LogLevel;
+  }
+): Promise<WorkId[]> {
+  const { config, ...defaults } = await enqueueArgs(fn, options);
+  const ids = await ctx.runMutation(component.lib.enqueueBatch, {
+    items: fnArgsArray.map((fnArgs) => ({
+      ...defaults,
+      fnArgs,
+      fnType,
+    })),
+    config,
+  });
+  return ids as WorkId[];
+}
+
+export async function enqueue<
+  FnType extends FunctionType,
+  Args extends DefaultFunctionArgs,
+  ReturnType,
+>(
+  component: UseApi<Mounts>,
+  ctx: RunMutationCtx,
+  fnType: FnType,
+  fn: FunctionReference<FnType, FunctionVisibility, Args, ReturnType>,
+  fnArgs: Args,
+  options: EnqueueOptions & {
+    retryBehavior?: RetryBehavior;
+    maxParallelism?: number;
+    logLevel?: LogLevel;
+  }
+): Promise<WorkId> {
+  const id = await ctx.runMutation(component.lib.enqueue, {
+    ...(await enqueueArgs(fn, options)),
+    fnArgs,
+    fnType,
+  });
+  return id as WorkId;
+}

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

@@ -89,6 +89,30 @@ export type Mounts = {
       },
       string
     >;
+    enqueueBatch: FunctionReference<
+      "mutation",
+      "public",
+      {
+        config: {
+          logLevel: "DEBUG" | "TRACE" | "INFO" | "REPORT" | "WARN" | "ERROR";
+          maxParallelism: number;
+        };
+        items: Array<{
+          fnArgs: any;
+          fnHandle: string;
+          fnName: string;
+          fnType: "action" | "mutation" | "query";
+          onComplete?: { context?: any; fnHandle: string };
+          retryBehavior?: {
+            base: number;
+            initialBackoffMs: number;
+            maxAttempts: number;
+          };
+          runAt: number;
+        }>;
+      },
+      Array<string>
+    >;
     status: FunctionReference<
       "query",
       "public",
@@ -97,6 +121,16 @@ export type Mounts = {
       | { previousAttempts: number; state: "running" }
       | { state: "finished" }
     >;
+    statusBatch: FunctionReference<
+      "query",
+      "public",
+      { ids: Array<string> },
+      Array<
+        | { previousAttempts: number; state: "pending" }
+        | { previousAttempts: number; state: "running" }
+        | { state: "finished" }
+      >
+    >;
   };
 };
 // For now fullApiWithMounts is only fullApi which provides

package/src/component/kick.ts

@@ -9,6 +9,8 @@ import {
   fromSegment,
   getCurrentSegment,
   getNextSegment,
+  SECOND,
+  toSegment,
 } from "./shared.js";
 
 /**
@@ -40,7 +42,7 @@ export async function kickMainLoop(
       );
       return next;
     }
-    if (runStatus.state.segment <= next) {
+    if (runStatus.state.segment <= toSegment(Date.now() + SECOND)) {
       console.debug(
         `[${source}] main is scheduled to run soon enough, so we don't need to kick it`
       );

package/src/component/lib.ts

@@ -1,10 +1,10 @@
-import { v } from "convex/values";
+import { Infer, ObjectType, v } from "convex/values";
 import { api } from "./_generated/api.js";
 import { fnType } from "./shared.js";
 import { Id } from "./_generated/dataModel.js";
-import { mutation, MutationCtx, query } from "./_generated/server.js";
+import { mutation, MutationCtx, query, QueryCtx } from "./_generated/server.js";
 import { kickMainLoop } from "./kick.js";
-import { createLogger, LogLevel, logLevel } from "./logging.js";
+import { createLogger, Logger, LogLevel, logLevel } from "./logging.js";
 import {
   boundScheduledTime,
   config,
@@ -20,44 +20,75 @@ import { recordEnqueued } from "./stats.js";
 const MAX_POSSIBLE_PARALLELISM = 100;
 const MAX_PARALLELISM_SOFT_LIMIT = 50;
 
+const itemArgs = {
+  fnHandle: v.string(),
+  fnName: v.string(),
+  fnArgs: v.any(),
+  fnType,
+  runAt: v.number(),
+  // TODO: annotation?
+  onComplete: v.optional(onComplete),
+  retryBehavior: v.optional(retryBehavior),
+};
+const enqueueArgs = {
+  ...itemArgs,
+  config,
+};
 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);
+    return await enqueueHandler(ctx, console, kickSegment, itemArgs);
+  },
+});
+async function enqueueHandler(
+  ctx: MutationCtx,
+  console: Logger,
+  kickSegment: bigint,
+  { runAt, ...workArgs }: ObjectType<typeof itemArgs>
+) {
+  runAt = boundScheduledTime(runAt, console);
+  const workId = await ctx.db.insert("work", {
+    ...workArgs,
+    attempts: 0,
+  });
+  await ctx.db.insert("pendingStart", {
+    workId,
+    segment: max(toSegment(runAt), kickSegment),
+  });
+  recordEnqueued(console, { workId, fnName: workArgs.fnName, runAt });
+  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: {
-    fnHandle: v.string(),
-    fnName: v.string(),
-    fnArgs: v.any(),
-    fnType,
-    runAt: v.number(),
-    // TODO: annotation?
-    onComplete: v.optional(onComplete),
-    retryBehavior: v.optional(retryBehavior),
+    items: v.array(v.object(itemArgs)),
     config,
   },
-  returns: v.id("work"),
-  handler: async (ctx, { config, runAt, ...workArgs }) => {
+  returns: v.array(v.id("work")),
+  handler: async (ctx, { config, items }) => {
+    validateConfig(config);
     const console = createLogger(config.logLevel);
-    if (config.maxParallelism > MAX_POSSIBLE_PARALLELISM) {
-      throw new Error(
-        `maxParallelism must be <= ${MAX_PARALLELISM_SOFT_LIMIT}`
-      );
-    } else if (config.maxParallelism > MAX_PARALLELISM_SOFT_LIMIT) {
-      console.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");
-    }
-    runAt = boundScheduledTime(runAt, console);
-    const workId = await ctx.db.insert("work", {
-      ...workArgs,
-      attempts: 0,
-    });
-    const limit = await kickMainLoop(ctx, "enqueue", config);
-    await ctx.db.insert("pendingStart", {
-      workId,
-      segment: max(toSegment(runAt), limit),
-    });
-    recordEnqueued(console, { workId, fnName: workArgs.fnName, runAt });
-    return workId;
+    const kickSegment = await kickMainLoop(ctx, "enqueue", config);
+    return Promise.all(
+      items.map((item) => enqueueHandler(ctx, console, kickSegment, item))
+    );
   },
 });
 
@@ -119,27 +150,38 @@ export const cancelAll = mutation({
 export const status = query({
   args: { id: v.id("work") },
   returns: statusValidator,
-  handler: async (ctx, { id }) => {
-    const work = await ctx.db.get(id);
-    if (!work) {
-      return { state: "finished" } as const;
-    }
-    const pendingStart = await ctx.db
-      .query("pendingStart")
-      .withIndex("workId", (q) => q.eq("workId", id))
-      .unique();
-    if (pendingStart) {
-      return { state: "pending", previousAttempts: work.attempts } as const;
-    }
-    const pendingCompletion = await ctx.db
-      .query("pendingCompletion")
-      .withIndex("workId", (q) => q.eq("workId", id))
-      .unique();
-    if (pendingCompletion?.retry) {
-      return { state: "pending", previousAttempts: work.attempts } as const;
-    }
-    // Assume it's in progress. It could be pending cancelation
-    return { state: "running", previousAttempts: work.attempts } as const;
+  handler: statusHandler,
+});
+async function statusHandler(ctx: QueryCtx, { id }: { id: Id<"work"> }) {
+  const work = await ctx.db.get(id);
+  if (!work) {
+    return { state: "finished" } as const;
+  }
+  const pendingStart = await ctx.db
+    .query("pendingStart")
+    .withIndex("workId", (q) => q.eq("workId", id))
+    .unique();
+  if (pendingStart) {
+    return { state: "pending", previousAttempts: work.attempts } as const;
+  }
+  const pendingCompletion = await ctx.db
+    .query("pendingCompletion")
+    .withIndex("workId", (q) => q.eq("workId", id))
+    .unique();
+  if (pendingCompletion?.retry) {
+    return { state: "pending", previousAttempts: work.attempts } as const;
+  }
+  // Assume it's in progress. It could be pending cancelation
+  return { state: "running", previousAttempts: work.attempts } as const;
+}
+
+export const statusBatch = query({
+  args: { ids: v.array(v.id("work")) },
+  returns: v.array(statusValidator),
+  handler: async (ctx, { ids }) => {
+    return await Promise.all(
+      ids.map(async (id) => await statusHandler(ctx, { id }))
+    );
   },
 });