diff --git a/CHANGELOG.md b/CHANGELOG.md index 3fb64a2..f5e212a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Changelog +## 0.2.1-alpha.0 + +- Provide a worker-persisted cursor to make it easier to implement iterating + through work using commitTs, which optimizes your queries to skip slow regions + of the database full of deleted items. Requires Convex 1.43+ for v.commitTs() + ## 0.2.0 - Removes the heartbeat field (breaking) diff --git a/README.md b/README.md index 0beef3a..5e526f2 100644 --- a/README.md +++ b/README.md @@ -56,14 +56,20 @@ component too: call `component.use(batchWorker)` in your component's ## Usage -Insert work into your own table, then call `ping`. Provide a query (typed with -`vBatchQueryArgs` / `vBatchResult`) that returns the next batch or `idle`, and a -mutation that processes it. The query's `batch` shape must match the mutation's -args. (`kind: "work"` is optional — returning `{ batch }` alone also works.) +Insert work into your own table, then call `ping`. Provide a query that returns +the next batch or `idle`, and a mutation that processes it. To ensure that the +batch shape always matches between the query and the mutation, use +`defineBatchWorkerValidators` to obtain all four argument and return validators. +(`kind: "work"` is optional — returning `{ batch }` alone also works.) + +Alongside the batch, the query returns a **cursor** saying how far it got. The +component commits it with the batch and hands it back as `args.cursor` next +time, so each scan resumes where the last one stopped. See +[why that matters](#why-is-a-cursor-necessary). ```ts import { v } from "convex/values"; -import { ping, vBatchQueryArgs, vBatchResult } from "@convex-dev/batch-worker"; +import { defineBatchWorkerValidators, ping } from "@convex-dev/batch-worker"; import { components, internal } from "./_generated/api"; import { internalMutation, internalQuery, mutation } from "./_generated/server"; @@ -73,7 +79,9 @@ const BATCH_SIZE = 10; export const addEvent = mutation({ args: { value: v.number() }, handler: async (ctx, { value }) => { - await ctx.db.insert("events", { value }); + // `db.vars.commitTs` resolves at commit to an int64 ordered by commit + // order — that's what makes it safe to use as a cursor. See below. + await ctx.db.insert("events", { value, insertedAt: ctx.db.vars.commitTs }); await ping(ctx, components.batchWorker, { name: "events", // distinct names give you independent queues workQuery: internal.example.getBatch, @@ -84,13 +92,21 @@ export const addEvent = mutation({ const vEvents = v.array(v.object({ id: v.id("events"), value: v.number() })); +// One declaration gives the query's args (`{ name, cursor? }`) and returns, +// plus the mutation's args. `cursor` defaults to a commit timestamp. +const { vQueryArgs, vQueryReturns, vMutationArgs } = + defineBatchWorkerValidators({ batch: { events: vEvents } }); + // Return the next batch of work, or `idle` when there's nothing to do. export const getBatch = internalQuery({ - args: vBatchQueryArgs, // { name } — lets one query serve multiple queues - returns: vBatchResult({ events: vEvents }), - handler: async (ctx) => { - // See the "Resume from a cursor" section below for an optimized approach - const events = await ctx.db.query("events").take(BATCH_SIZE); + args: vQueryArgs, // `name` lets one query serve many queues + returns: vQueryReturns, + handler: async (ctx, { cursor }) => { + // Resume where the last batch stopped. + const events = await ctx.db + .query("events") + .withIndex("insertedAt", (q) => q.gte("insertedAt", cursor ?? 0n)) + .take(BATCH_SIZE); if (events.length === 0) { return { kind: "idle" as const }; // Or, if you know when the next item is due: @@ -98,14 +114,16 @@ export const getBatch = internalQuery({ } return { kind: "work" as const, - batch: events.map((e) => ({ id: e._id, value: e.value })), + batch: { events: events.map((e) => ({ id: e._id, value: e.value })) }, + // Rows come back in commit order, so the last one is how far we got. + cursor: events.at(-1)!.insertedAt, }; }, }); // Process one batch. The worker owns cleanup — delete what you process! export const processBatch = internalMutation({ - args: { events: vEvents }, + args: vMutationArgs, handler: async (ctx, { events }) => { for (const { value, id } of events) { // ... do the work (sum, call an API, schedule downstream jobs, etc.) ... @@ -122,9 +140,6 @@ The component **does not clean up your work for you** — your worker mutation i responsible for deleting (or marking complete / advancing past) the rows it processed, otherwise the next query will return them again. -Note: the snippet above keeps that as simple as possible; for a busy worker, use -a cursor as described [below](#use-a-cursor-to-track-worker-progress). - ### Fetch work in the query, process it in the mutation Think of the work query and the worker mutation as **two separate transactions** @@ -154,7 +169,8 @@ return { ### Steering the loop dynamically -Your worker mutation may return `{ debounceMs }` to throttle the loop: +Your worker mutation may return `{ debounceMs }` to throttle the loop, and a +`cursor` to override the query's (see [The cursor](#the-cursor)): ```ts return { @@ -214,80 +230,96 @@ await ping(ctx, components.batchWorker, { }); ``` -### Use a cursor to track worker progress +### The cursor -For best performance, it's advised to use a cursor to track how far you've read, -and use it to focus future queries, even if you're deleting documents as you go. -See below for [why a cursor is necessary](#why-is-a-cursor-necessary). +The component stores the cursor for you: whatever your work query returns +alongside a batch comes back as `args.cursor` on the next call. It's committed +in the same transaction as the batch, so if your worker mutation throws, the +cursor stays where it was and the batch is handed out again. -The best cursor to use is the +The default cursor type is a [commit timestamp](https://docs.convex.dev/database/advanced/commit-timestamp) -(Convex ≥ 1.43), as it avoids having to handle out-of-order commits. Do not use -`_creationTime`. It's assigned when a mutation _starts_, so a row inserted by a -slow mutation can land _behind_ rows that were already committed and queried by -the worker. - -**Step 1**: Add a field to your schema tracking the commit timestamp for work -you'll query, and track the cursor somewhere for the worker to access and -update. +(Convex ≥ 1.43), which avoids having to handle out-of-order commits. Add a +`v.commitTs()` field to the table the worker scans, index it, and write +`ctx.db.vars.commitTs` to it on every insert or patch: ```ts // convex/schema.ts -const schema = defineSchema({ - myWork: defineTable({ - updatedAt: v.commitTs(), - ...fields - }).index("updatedAt", ["updatedAt"]), - workerState: defineTable({ - cursor: v.int64(), // a commit timestamp resolves to a bigint - ... - }) - ... -}); +myWork: defineTable({ + updatedAt: v.commitTs(), // resolves at commit to an int64 in commit order + ...fields +}).index("updatedAt", ["updatedAt"]), ``` -**Step 2:** Set it when inserting or updating a document that needs processing: - ```ts -// On insert/patch, use the placeholder, (resolves to a bigint at commit time) -await ctx.db.insert("myWork", { updatedAt: ctx.db.vars.commitTs, ... }); await -ctx.db.patch("myWork", id, { updatedAt: ctx.db.vars.commitTs, ... }); +await ctx.db.insert("myWork", { updatedAt: ctx.db.vars.commitTs, ... }); +await ctx.db.patch("myWork", id, { updatedAt: ctx.db.vars.commitTs, ... }); ``` -**Step 3:** In the worker query, resume from the cursor and report how far you -got: +Do **not** use `_creationTime`. It's assigned when a mutation _starts_, so a row +inserted by a slow mutation can land _behind_ rows the worker has already +queried past. + +**Scan with `gte`.** Everything a single mutation inserts shares one commit +timestamp, so a batch can end in the middle of such a tie and the next scan has +to start at the tie to pick up the rest. Re-reading from the tie costs at most +one mutation's worth of already-processed rows. + +To end each batch on a tie boundary instead, read the remaining documents at +that commit timestamp into the batch and scan with `gt`. That also processes +work added "at once" together, as long as it fits in one transaction. See +[the aggregate example](./example/convex/aggregates.ts). + +#### Other cursor types + +The cursor can be any Convex value, for instance a `paginator` cursor from +`convex-helpers`, which captures the index range +`[commitTs, _creationTime, _id]` so you know exactly where you left off. Pass +`cursor` to `defineBatchWorkerValidators` and all four validators pick it up: ```ts -const workerState = await getWorkerState(ctx, ...) // however you access it -const work = await ctx.db - .query("myWork") - .withIndex("updatedAt", (q) => q.gte("updatedAt", workerState.cursor)) - .take(BATCH_SIZE); -if (events.length === 0) return { kind: "idle" as const }; -const cursor = work.at(-1)!.updatedAt as bigint; // cast is fine here -return { kind: "work" as const, batch: { cursor, ... } }; +const { vQueryArgs, vQueryReturns, vMutationArgs, vMutationReturns } = + defineBatchWorkerValidators({ + batch: { ids: v.array(v.id("myWork")) }, + cursor: v.string(), + }); ``` -Note: the query uses `gte`, not `gt`, since multiple entries may have been -inserted at the same commit timestamp (they committed transactionally). If our -batch of work lands in the middle of them, we don't want to skip the rest. +`ping` takes the cursor's type from the work query's `cursor` arg and checks +both return types against it, so a mismatch points at one place. + +#### Adjusting the cursor from the worker mutation + +The worker mutation can return a `cursor` that overrides the query's. Use it +when the batch made partial progress and you can work out how far it got, e.g. +from a per-item timestamp carried in the batch: + +```ts +return { cursor: lastProcessed.updatedAt }; +``` -An alternative here is to read the remaining documents from that same commitTs -and include them in this batch. This has the advantage of ensuring all work -items added at "once" get processed together, if you can guarantee you can -process all of that work at once without exceeding transaction limits. See an -example of that in [the aggregate example](./example/convex/aggregates.ts). +If it is not returned, it will defer to the cursor returned by the work query, +if one was returned, or keep the old value. Note: if the cursor does not advance +and you did not update anything the query read, you may end up in an infinite +loop. -Another option is to use `paginator` from `convex-helpers` to track a cursor -that captures the index range `[commitTs, _creationTime, _id]`, so you know -exactly where you left off while still being immune to out-of-order commits. +#### Reading and resetting the cursor -**Step 4:** In the worker mutation, update the cursor +The work query receives the cursor in its args. To read it elsewhere, say from a +query counting pending work, call the component like `status`: ```ts -await ctx.db.patch("workerState", workerStateId, { cursor: args.cursor }); +const from = ((await ctx.runQuery(components.batchWorker.lib.getCursor, { + name: "events", +})) ?? 0n) as bigint; ``` +`lib.setCursor` overwrites it, and clears it when `cursor` is omitted, so the +next scan starts from the front. Use it for migrations and recovery: the loop +writes that document every iteration, so a call made while the worker is busy is +liable to conflict. It throws if the worker doesn't exist yet (`ping` creates +it). + #### Why is a cursor necessary? The reason is due to a technical nuance of how Convex tracks document changes. @@ -388,11 +420,12 @@ Tips for rate-limiting LLM calls: ### Failure handling If your work query or worker mutation throws, the loop dies and the liveness -monitor restarts it after ~`monitorLagMs` — and since the unprocessed rows are -still in your table, the query will hand out the **same batch again**. That -gives you at-least-once processing, but it also means one poison item that -always throws can wedge the queue. For work that can fail per item, catch errors -inside the worker mutation, and isolate bad docs in a table for async debugging. +monitor restarts it after ~`monitorLagMs`. The unprocessed rows are still in +your table and the cursor stayed where it was, so the query hands out the **same +batch again**. That gives you at-least-once processing, but it also means one +poison item that always throws can wedge the queue. For work that can fail per +item, catch errors inside the worker mutation, and isolate bad docs in a table +for async debugging. This is a low-level primitive, relative to components like Workpool or Workflow, so you have to handle exceptional cases yourself. @@ -457,9 +490,8 @@ Note: Use `convex-test` (≥ 0.0.55-alpha.0) for `v.commitTs()` support. See the full working examples in the example app: [example.ts](./example/convex/example.ts) (basic queue), -[aggregates.ts](./example/convex/aggregates.ts) (denormalized aggregates), -[rateLimited.ts](./example/convex/rateLimited.ts) (async LLM batches), and -[cursor.ts](./example/convex/cursor.ts) (the shared `commitTs` cursor helpers). +[aggregates.ts](./example/convex/aggregates.ts) (denormalized aggregates), and +[rateLimited.ts](./example/convex/rateLimited.ts) (async LLM batches). @@ -489,7 +521,8 @@ published to `https://.convex.site`. See | `workerState` | `loop` (every iteration) | `loop`, monitor | The high-churn loop state lives in `workerState` (generation, heartbeat, the -scheduled runner, and the monitor), separate from the rarely-written `workers` +scheduled runner, the monitor, and the cursor, which rides along on a patch the +loop already does every iteration), separate from the rarely-written `workers` doc (which holds the handles, config, and run-status: `idle` / `running` / `stopped`, plus a pointer to its `workerState`). That lets `ping`/`start` — which you call on every insert — read `workers` and return without conflicting diff --git a/example/README.md b/example/README.md index d32b182..27f41ed 100644 --- a/example/README.md +++ b/example/README.md @@ -9,9 +9,9 @@ defines three named workers and the Vite UI ([`src/`](./src)) drives each one: - **Rate-limited LLM batches** ([rateLimited.ts](./convex/rateLimited.ts)) — collect requests and spend a token budget on async calls. -All three drain their queue by resuming from a `v.commitTs()` cursor -([cursor.ts](./convex/cursor.ts)) rather than rescanning the front of the table -each batch. +All three drain their queue by returning a `v.commitTs()` cursor from the work +query. The component commits it with the batch and hands it back on the next +call, so each scan resumes where the last one stopped. Run all commands from the **root of the repo**. diff --git a/example/convex/_generated/api.d.ts b/example/convex/_generated/api.d.ts index c4e7e6b..bfd09a4 100644 --- a/example/convex/_generated/api.d.ts +++ b/example/convex/_generated/api.d.ts @@ -9,7 +9,6 @@ */ import type * as aggregates from "../aggregates.js"; -import type * as cursor from "../cursor.js"; import type * as e2e from "../e2e.js"; import type * as example from "../example.js"; import type * as rateLimited from "../rateLimited.js"; @@ -22,7 +21,6 @@ import type { declare const fullApi: ApiFromModules<{ aggregates: typeof aggregates; - cursor: typeof cursor; e2e: typeof e2e; example: typeof example; rateLimited: typeof rateLimited; diff --git a/example/convex/aggregates.test.ts b/example/convex/aggregates.test.ts index cb209dc..57bda51 100644 --- a/example/convex/aggregates.test.ts +++ b/example/convex/aggregates.test.ts @@ -1,6 +1,6 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { initConvexTest } from "./setup.test"; -import { api } from "./_generated/api"; +import { api, components } from "./_generated/api"; describe("serial aggregate updates", () => { beforeEach(() => { @@ -92,13 +92,14 @@ describe("serial aggregate updates", () => { await t.mutation(api.aggregates.recordScore, { team: "red", points: 2 }); await t.finishAllScheduledFunctions(vi.runAllTimers); - const [scores, cursor] = await t.run(async (ctx) => [ - await ctx.db.query("scoreEvents").withIndex("insertedAt").collect(), - await ctx.db - .query("cursors") - .withIndex("name", (q) => q.eq("name", "aggregates")) - .unique(), - ]); - expect(cursor?.commitTs).toBe(scores.at(-1)!.insertedAt); + const scores = await t.run((ctx) => + ctx.db.query("scoreEvents").withIndex("insertedAt").collect(), + ); + const cursor = await t.run((ctx) => + ctx.runQuery(components.batchWorker.lib.getCursor, { + name: "aggregates", + }), + ); + expect(cursor).toBe(scores.at(-1)!.insertedAt); }); }); diff --git a/example/convex/aggregates.ts b/example/convex/aggregates.ts index 94b5974..7211b17 100644 --- a/example/convex/aggregates.ts +++ b/example/convex/aggregates.ts @@ -1,5 +1,5 @@ import { v } from "convex/values"; -import { ping, vBatchQueryArgs, vBatchResult } from "@convex-dev/batch-worker"; +import { defineBatchWorkerValidators, ping } from "@convex-dev/batch-worker"; import { components, internal } from "./_generated/api.js"; import { internalMutation, @@ -7,7 +7,6 @@ import { mutation, query, } from "./_generated/server.js"; -import { advanceCursor, cursorFor } from "./cursor.js"; // Serial processing to update denormalized aggregates without write conflicts. // @@ -35,7 +34,7 @@ export const recordScore = mutation({ handler: async (ctx, { team, points }) => { // `db.vars.commitTs` resolves to this mutation's commit timestamp. Every row // a mutation writes shares one value — getBatch below reads a whole tie at - // once — and nothing can commit later with a smaller one. See cursor.ts. + // once — and nothing can commit later with a smaller one. await ctx.db.insert("scoreEvents", { team, points, @@ -45,27 +44,28 @@ export const recordScore = mutation({ }, }); -const vBatch = { - events: v.array( - v.object({ - team: v.string(), - points: v.number(), - }), - ), - cursor: v.int64(), -}; +const { vQueryArgs, vQueryReturns, vMutationArgs } = + defineBatchWorkerValidators({ + batch: { + events: v.array( + v.object({ + team: v.string(), + points: v.number(), + }), + ), + }, + }); export const getBatch = internalQuery({ - args: vBatchQueryArgs, - returns: vBatchResult(vBatch), - handler: async (ctx, { name }) => { + args: vQueryArgs, + returns: vQueryReturns, + handler: async (ctx, { cursor }) => { // Resume from after the last batch's commit timestamp. // We don't delete events, so the cursor allows us to avoid // handling scores multiple times. - const from = await cursorFor(ctx, name); const events = await ctx.db .query("scoreEvents") - .withIndex("insertedAt", (q) => q.gt("insertedAt", from)) + .withIndex("insertedAt", (q) => q.gt("insertedAt", cursor ?? 0n)) .take(BATCH_SIZE); if (events.length === 0) { return { kind: "idle" as const }; @@ -88,16 +88,16 @@ export const getBatch = internalQuery({ team: e.team, points: e.points, })), - // Rows come back in commit order, so the last one is how far we got. - cursor: lastCommitTs, }, + // Rows come back in commit order, so the last one is how far we got. + cursor: lastCommitTs, }; }, }); export const processBatch = internalMutation({ - args: vBatch, - handler: async (ctx, { events, cursor }) => { + args: vMutationArgs, + handler: async (ctx, { events }) => { // Fold the batch into one delta per team first, so we touch each aggregate // row once regardless of how many events it covers. const deltas = new Map(); @@ -115,8 +115,6 @@ export const processBatch = internalMutation({ await ctx.db.insert("teamTotals", { team, total: delta }); } } - // Only the loop writes the cursor, so this never conflicts with inserts. - await advanceCursor(ctx, WORKER, cursor); // Returning nothing re-runs immediately to drain the rest. }, }); @@ -127,8 +125,10 @@ export const getTotals = query({ handler: async (ctx) => { const totals = await ctx.db.query("teamTotals").take(100); // Scores still queued, waiting to be folded into the aggregates. Reading - // from the cursor keeps this off the tombstones too. - const from = await cursorFor(ctx, WORKER); + // from the worker's cursor keeps this off the tombstones too. + const from = ((await ctx.runQuery(components.batchWorker.lib.getCursor, { + name: WORKER, + })) ?? 0n) as bigint; const pending = ( await ctx.db .query("scoreEvents") diff --git a/example/convex/cursor.ts b/example/convex/cursor.ts deleted file mode 100644 index 1e70753..0000000 --- a/example/convex/cursor.ts +++ /dev/null @@ -1,66 +0,0 @@ -import type { MutationCtx, QueryCtx } from "./_generated/server.js"; - -// Resuming a queue scan from a cursor, instead of always reading from the front. -// -// The obvious way to drain a queue is `.take(BATCH_SIZE)` from the start of the -// table and delete the rows you processed. That works, but every delete leaves a -// tombstone in the index until it's vacuumed (a few minutes later), and a scan -// that always starts at the front has to walk over them. On a busy queue that -// turns a fixed-size batch into a read that keeps growing. -// -// A `v.commitTs()` field fixes that. The examples here call it `insertedAt` or -// `updatedAt`. Write `ctx.db.vars.commitTs` to it on every insert or patch and -// it resolves, when that mutation commits, to an int64 ordered by commit order. -// It is safe to use as a cursor because nothing can be inserted in the past. -// By comparison, `_creationTime` does not work because it's assigned based on -// when the mutation starts, so it may be committed in the past after reading -// a more recent event. -// The worker stores the timestamp it got to, and the next scan starts there, -// skipping all the tombstones from previous deletes. -// -// The examples in this app keep their cursors in one `cursors` table, keyed by -// worker name. Any single-writer spot works — if your worker already updates a -// document each batch (like `totals` in example.ts), a field on that document -// saves a write. Only the worker mutation should write it, to avoid write conflicts. -function cursorDoc(ctx: QueryCtx, name: string) { - return ctx.db - .query("cursors") - .withIndex("name", (q) => q.eq("name", name)) - .unique(); -} - -/** Where the next scan for `name` should start (inclusive). */ -export async function cursorFor(ctx: QueryCtx, name: string): Promise { - const cursor = await cursorDoc(ctx, name); - return cursor?.commitTs ?? 0n; -} - -/** - * Record where this batch stopped, so the next one resumes there. - * - * The cursor is *inclusive* — scans use `.gte(cursor)`, not `.gt(cursor)`. - * Everything a single mutation inserts shares one commit timestamp, so a batch - * can end in the middle of a tie, and `.gt(...)` would skip the rest of that - * mutation's rows. Re-reading from the tie costs at most one mutation's worth of - * already-processed rows (tombstones, if the worker deletes them). - */ -export async function advanceCursor( - ctx: MutationCtx, - name: string, - commitTs: bigint, -) { - const existing = await cursorDoc(ctx, name); - if (existing) { - await ctx.db.patch("cursors", existing._id, { commitTs }); - } else { - await ctx.db.insert("cursors", { name, commitTs }); - } -} - -/** Forget where `name` got to, so its next scan starts from the front. */ -export async function resetCursor(ctx: MutationCtx, name: string) { - const existing = await cursorDoc(ctx, name); - if (existing) { - await ctx.db.delete("cursors", existing._id); - } -} diff --git a/example/convex/e2e.test.ts b/example/convex/e2e.test.ts index fa8a845..cbc55fd 100644 --- a/example/convex/e2e.test.ts +++ b/example/convex/e2e.test.ts @@ -1,6 +1,6 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { initConvexTest } from "./setup.test"; -import { api } from "./_generated/api"; +import { api, components } from "./_generated/api"; describe("e2e harness worker", () => { beforeEach(() => { @@ -29,10 +29,7 @@ describe("e2e harness worker", () => { await t.mutation(api.e2e.reset, {}); const cursor = await t.run((ctx) => - ctx.db - .query("cursors") - .withIndex("name", (q) => q.eq("name", "e2e")) - .unique(), + ctx.runQuery(components.batchWorker.lib.getCursor, { name: "e2e" }), ); expect(cursor).toBeNull(); expect(await t.query(api.e2e.samples, {})).toEqual([]); diff --git a/example/convex/e2e.ts b/example/convex/e2e.ts index 59c256b..734e314 100644 --- a/example/convex/e2e.ts +++ b/example/convex/e2e.ts @@ -1,5 +1,5 @@ import { v } from "convex/values"; -import { ping, vBatchQueryArgs, vBatchResult } from "@convex-dev/batch-worker"; +import { defineBatchWorkerValidators, ping } from "@convex-dev/batch-worker"; import { components, internal } from "./_generated/api.js"; import { internalMutation, @@ -7,7 +7,6 @@ import { mutation, query, } from "./_generated/server.js"; -import { advanceCursor, cursorFor, resetCursor } from "./cursor.js"; // Instrumented worker used by e2e.mjs to measure performance. Kept separate // from the README example (different worker `name`) so it doesn't interfere. @@ -27,8 +26,10 @@ export const reset = mutation({ const docs = await ctx.db.query(table).collect(); for (const d of docs) await ctx.db.delete(table, d._id); } - // Reset this worker's cursor too, so the next scenario starts from scratch. - await resetCursor(ctx, WORKER); + // Clear this worker's cursor too, so the next scenario starts from scratch. + await ctx.runMutation(components.batchWorker.lib.setCursor, { + name: WORKER, + }); }, }); @@ -50,23 +51,24 @@ export const enqueue = mutation({ }, }); -export const getBatch = internalQuery({ - args: vBatchQueryArgs, - returns: vBatchResult( - v.object({ +const { vQueryArgs, vQueryReturns, vMutationArgs } = + defineBatchWorkerValidators({ + batch: { items: v.array( v.object({ id: v.id("e2eEvents"), creationTime: v.number() }), ), - cursor: v.int64(), - }), - ), - handler: async (ctx, { name }) => { + }, + }); + +export const getBatch = internalQuery({ + args: vQueryArgs, + returns: vQueryReturns, + handler: async (ctx, { cursor }) => { // Resume from the cursor so batch reads stay a fixed size no matter how // many rows (and tombstones) the harness has already churned through. - const from = await cursorFor(ctx, name); const events = await ctx.db .query("e2eEvents") - .withIndex("updatedAt", (q) => q.gte("updatedAt", from)) + .withIndex("updatedAt", (q) => q.gte("updatedAt", cursor ?? 0n)) .take(BATCH_SIZE); if (events.length === 0) { return { kind: "idle" as const }; @@ -78,20 +80,15 @@ export const getBatch = internalQuery({ id: e._id, creationTime: e._creationTime, })), - cursor: events.at(-1)!.updatedAt as bigint, }, + cursor: events.at(-1)!.updatedAt, }; }, }); export const processBatch = internalMutation({ - args: { - items: v.array( - v.object({ id: v.id("e2eEvents"), creationTime: v.number() }), - ), - cursor: v.int64(), - }, - handler: async (ctx, { items, cursor }) => { + args: vMutationArgs, + handler: async (ctx, { items }) => { const now = Date.now(); const latencies = items.map((b) => now - b.creationTime); await ctx.db.insert("e2eSamples", { @@ -103,7 +100,6 @@ export const processBatch = internalMutation({ for (const b of items) { await ctx.db.delete("e2eEvents", b.id); } - await advanceCursor(ctx, WORKER, cursor); }, }); @@ -115,7 +111,9 @@ export const samples = query({ export const pending = query({ args: {}, handler: async (ctx) => { - const from = await cursorFor(ctx, WORKER); + const from = ((await ctx.runQuery(components.batchWorker.lib.getCursor, { + name: WORKER, + })) ?? 0n) as bigint; return ( await ctx.db .query("e2eEvents") diff --git a/example/convex/example.test.ts b/example/convex/example.test.ts index e539d7c..465ecfb 100644 --- a/example/convex/example.test.ts +++ b/example/convex/example.test.ts @@ -1,6 +1,6 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { initConvexTest } from "./setup.test"; -import { api } from "./_generated/api"; +import { api, components } from "./_generated/api"; describe("example worker", () => { beforeEach(() => { @@ -107,16 +107,18 @@ describe("example worker", () => { await t.finishAllScheduledFunctions(vi.runAllTimers); const cursor = await t.run((ctx) => - ctx.db - .query("cursors") - .withIndex("name", (q) => q.eq("name", "events")) - .unique(), + ctx.runQuery(components.batchWorker.lib.getCursor, { name: "events" }), ); - expect(cursor?.commitTs).toBe(lastInsertedAt); + expect(cursor).toBe(lastInsertedAt); }); test("leaves events behind the cursor alone", async () => { const t = initConvexTest(); + // The cursor lives on the worker, so it has to exist before we can park it. + await t.mutation(api.example.addEvent, { value: 1 }); + await t.finishAllScheduledFunctions(vi.runAllTimers); + + // Insert an event without pinging, so the loop doesn't pick it up. await t.run(async (ctx) => { await ctx.db.insert("events", { value: 100, @@ -131,9 +133,9 @@ describe("example worker", () => { .query("events") .withIndex("insertedAt") .collect(); - await ctx.db.insert("cursors", { + await ctx.runMutation(components.batchWorker.lib.setCursor, { name: "events", - commitTs: (old.insertedAt as bigint) + 1n, + cursor: (old.insertedAt as bigint) + 1n, }); }); @@ -142,8 +144,8 @@ describe("example worker", () => { // Only the new event was summed — the scan really does start at the cursor. expect(await t.query(api.example.getTotals, {})).toEqual({ - total: 5, - count: 1, + total: 6, + count: 2, pending: 0, }); const left = await t.run((ctx) => ctx.db.query("events").collect()); diff --git a/example/convex/example.ts b/example/convex/example.ts index 7e946d0..1e0bc46 100644 --- a/example/convex/example.ts +++ b/example/convex/example.ts @@ -1,5 +1,5 @@ import { v } from "convex/values"; -import { ping, vBatchQueryArgs, vBatchResult } from "@convex-dev/batch-worker"; +import { defineBatchWorkerValidators, ping } from "@convex-dev/batch-worker"; import { components, internal } from "./_generated/api.js"; import { internalMutation, @@ -7,7 +7,6 @@ import { mutation, query, } from "./_generated/server.js"; -import { advanceCursor, cursorFor } from "./cursor.js"; // Use distinct `name`s if you want several independent queues backed by the // same component. @@ -24,7 +23,7 @@ export const addEvent = mutation({ handler: async (ctx, { value }) => { // `db.vars.commitTs` resolves to this mutation's commit timestamp. Every row // a mutation writes shares one value, and nothing can commit later with a - // smaller one — that's what makes it safe to cursor on. See cursor.ts. + // smaller one, which is what makes it safe to use as a cursor. await ctx.db.insert("events", { value, insertedAt: ctx.db.vars.commitTs }); await ping(ctx, components.batchWorker, { name: WORKER, @@ -34,49 +33,49 @@ export const addEvent = mutation({ }, }); +const vEvent = v.object({ id: v.id("events"), value: v.number() }); + +const { vQueryArgs, vQueryReturns, vMutationArgs, vMutationReturns } = + defineBatchWorkerValidators({ batch: { events: v.array(vEvent) } }); + /** * The work query: returns the next batch of work, or `idle` when the queue is - * empty. Its `batch` shape lines up with `processBatch`'s args. + * empty. */ -const vEvent = v.object({ id: v.id("events"), value: v.number() }); - export const getBatch = internalQuery({ - args: vBatchQueryArgs, - returns: vBatchResult({ - events: v.array(vEvent), - cursor: v.int64(), - }), - handler: async (ctx, { name }) => { - // Pick up where the last batch stopped, rather than from the front of the - // table — where every row we've already deleted leaves a tombstone to scan - // over. `gte` because the cursor is inclusive (see cursor.ts). - const from = await cursorFor(ctx, name); + args: vQueryArgs, + returns: vQueryReturns, + handler: async (ctx, { cursor }) => { + // Pick up where the last batch stopped, skipping the tombstone every row + // we deleted left at the front of the table. `gte` because the cursor is + // inclusive: everything one mutation inserts shares a commit timestamp, + // and a batch can end mid-tie. const events = await ctx.db .query("events") - .withIndex("insertedAt", (q) => q.gte("insertedAt", from)) + .withIndex("insertedAt", (q) => q.gte("insertedAt", cursor ?? 0n)) .take(BATCH_SIZE); if (events.length === 0) { return { kind: "idle" as const }; } return { kind: "work" as const, - batch: { - events: events.map((e) => ({ id: e._id, value: e.value })), - // Rows come back in commit order, so the last one is how far we got. - cursor: events.at(-1)!.insertedAt as bigint, - }, + batch: { events: events.map((e) => ({ id: e._id, value: e.value })) }, + // Rows come back in commit order, so the last one is how far we got. + // The component commits this with the batch and hands it back above. + cursor: events.at(-1)!.insertedAt, }; }, }); /** - * The worker mutation: processes a batch. It owns cleanup — advancing the cursor - * past what it processed, and deleting those rows so the table stays small. - * Returning nothing re-runs immediately to drain the rest. + * The worker mutation: processes a batch. It owns cleanup, deleting the rows it + * processed so the table stays small. Returning nothing re-runs immediately to + * drain the rest. */ export const processBatch = internalMutation({ - args: { events: v.array(vEvent), cursor: v.int64() }, - handler: async (ctx, { events, cursor }) => { + args: vMutationArgs, + returns: vMutationReturns, + handler: async (ctx, { events }) => { const sum = events.reduce((a, e) => a + e.value, 0); const totals = await ctx.db .query("totals") @@ -97,8 +96,6 @@ export const processBatch = internalMutation({ for (const { id } of events) { await ctx.db.delete("events", id); } - // Only the loop writes the cursor, so this never conflicts with inserts. - await advanceCursor(ctx, WORKER, cursor); }, }); @@ -110,8 +107,10 @@ export const getTotals = query({ .withIndex("key", (q) => q.eq("key", "all")) .unique(); // Events left in the queue waiting for the worker to pick them up. Reading - // from the cursor keeps this off the tombstones too. - const from = await cursorFor(ctx, WORKER); + // from the worker's cursor keeps this off the tombstones too. + const from = ((await ctx.runQuery(components.batchWorker.lib.getCursor, { + name: WORKER, + })) ?? 0n) as bigint; const pending = ( await ctx.db .query("events") diff --git a/example/convex/rateLimited.ts b/example/convex/rateLimited.ts index 6e9d26b..844bafb 100644 --- a/example/convex/rateLimited.ts +++ b/example/convex/rateLimited.ts @@ -1,6 +1,6 @@ import { v } from "convex/values"; import { SECOND, RateLimiter } from "@convex-dev/rate-limiter"; -import { ping, vBatchQueryArgs, vBatchResult } from "@convex-dev/batch-worker"; +import { defineBatchWorkerValidators, ping } from "@convex-dev/batch-worker"; import { components, internal } from "./_generated/api.js"; import { internalAction, @@ -9,7 +9,6 @@ import { mutation, query, } from "./_generated/server.js"; -import { advanceCursor, cursorFor } from "./cursor.js"; // Batching async LLM requests, paced by a token budget. // @@ -62,7 +61,7 @@ export const submitRequest = mutation({ inputTokens, state: "pending", // Lets the worker cursor through pending requests in commit order rather - // than rescanning from the front of the range. See cursor.ts. + // than rescanning from the front of the range. updatedAt: ctx.db.vars.commitTs, }); await ping(ctx, components.batchWorker, worker); @@ -77,15 +76,13 @@ const vLlmRequest = v.object({ inputTokens: v.number(), }); -const vBatch = v.object({ - requests: v.array(vLlmRequest), - cursor: v.int64(), -}); +const { vQueryArgs, vQueryReturns, vMutationArgs, vMutationReturns } = + defineBatchWorkerValidators({ batch: { requests: v.array(vLlmRequest) } }); export const getBatch = internalQuery({ - args: vBatchQueryArgs, - returns: vBatchResult(vBatch), - handler: async (ctx, { name }) => { + args: vQueryArgs, + returns: vQueryReturns, + handler: async (ctx, { cursor }) => { // TODO: also pick up "started" requests whose action never reported back // (e.g. `startedAt` older than some timeout) so a crashed batch isn't // stranded. That scan walks the (small) "started" range, so it doesn't use @@ -94,11 +91,10 @@ export const getBatch = internalQuery({ // Nothing is deleted here — requests are patched from "pending" to // "started" — but patching a row out of the pending range leaves a // tombstone there just the same, so the cursor still earns its keep. - const from = await cursorFor(ctx, name); const pending = await ctx.db .query("llmRequests") .withIndex("state_updatedAt", (q) => - q.eq("state", "pending").gte("updatedAt", from), + q.eq("state", "pending").gte("updatedAt", cursor ?? 0n), ) .take(BATCH_SIZE); if (pending.length === 0) { @@ -112,9 +108,9 @@ export const getBatch = internalQuery({ prompt: r.prompt, inputTokens: r.inputTokens, })), - // Rows come back in commit order, so the last one is how far we got. - cursor: pending.at(-1)!.updatedAt as bigint, }, + // Rows come back in commit order, so the last one is how far we got. + cursor: pending.at(-1)!.updatedAt, }; }, }); @@ -124,8 +120,9 @@ export const getBatch = internalQuery({ * started, and schedule the LLM call for when the budget clears. */ export const startBatch = internalMutation({ - args: vBatch, - handler: async (ctx, { requests, cursor }) => { + args: vMutationArgs, + returns: vMutationReturns, + handler: async (ctx, { requests }) => { const totalInputTokens = requests.reduce((a, r) => a + r.inputTokens, 0); // Reserve the whole batch's input tokens up front. `reserve: true` never @@ -148,7 +145,6 @@ export const startBatch = internalMutation({ updatedAt: ctx.db.vars.commitTs, }); } - await advanceCursor(ctx, WORKER, cursor); // Make the actual call when the reservation clears. Pass the prompts // through so the action doesn't have to re-read them. @@ -263,8 +259,10 @@ export const stats = query({ .withIndex("state_updatedAt", (q) => q.eq("state", state)) .take(500) ).length; - // The pending count reads from the cursor, like the work query does. - const from = await cursorFor(ctx, WORKER); + // The pending count reads from the worker's cursor, like the work query. + const from = ((await ctx.runQuery(components.batchWorker.lib.getCursor, { + name: WORKER, + })) ?? 0n) as bigint; const pending = ( await ctx.db .query("llmRequests") diff --git a/example/convex/schema.ts b/example/convex/schema.ts index 9359629..244a190 100644 --- a/example/convex/schema.ts +++ b/example/convex/schema.ts @@ -4,10 +4,10 @@ import { v } from "convex/values"; export default defineSchema({ // A simple work queue: each row is one event to be summed. // - // `updatedAt` is written as `ctx.db.vars.commitTs` and resolves, when the + // `insertedAt` is written as `ctx.db.vars.commitTs` and resolves, when the // writing mutation commits, to an int64 ordered by commit order. Indexed so - // the worker can resume where it left off instead of scanning from the front - // of the table every batch — see cursor.ts. + // the worker can resume from the cursor it last returned instead of scanning + // from the front of the table every batch. events: defineTable({ value: v.number(), insertedAt: v.commitTs(), @@ -19,14 +19,6 @@ export default defineSchema({ count: v.number(), }).index("key", ["key"]), - // Where each worker's next scan should start, one row per worker name. This is - // a plain int64 — the commit timestamp some *other* row resolved to, copied - // here — not a `v.commitTs()` field of its own. See cursor.ts. - cursors: defineTable({ - name: v.string(), - commitTs: v.int64(), - }).index("name", ["name"]), - // --- Rate-limited async LLM batches (see rateLimited.ts) --- llmRequests: defineTable({ prompt: v.string(), diff --git a/package-lock.json b/package-lock.json index e5a98b4..61cfbb6 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "@convex-dev/batch-worker", - "version": "0.2.0", + "version": "0.2.1-alpha.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "@convex-dev/batch-worker", - "version": "0.2.0", + "version": "0.2.1-alpha.0", "license": "Apache-2.0", "dependencies": { "convex-helpers": "^0.1.119" diff --git a/package.json b/package.json index 3d6ecbf..224f869 100644 --- a/package.json +++ b/package.json @@ -6,7 +6,7 @@ "bugs": { "url": "https://github.com/get-convex/batch-worker/issues" }, - "version": "0.2.0", + "version": "0.2.1-alpha.0", "license": "Apache-2.0", "keywords": [ "convex", @@ -64,7 +64,7 @@ } }, "peerDependencies": { - "convex": "^1.42.0", + "convex": "^1.43.0", "react": "^18.3.1 || ^19.0.0" }, "devDependencies": { diff --git a/src/client/index.test.ts b/src/client/index.test.ts index 9e0e8f5..61d2eea 100644 --- a/src/client/index.test.ts +++ b/src/client/index.test.ts @@ -12,14 +12,29 @@ import { mutationGeneric, queryGeneric, } from "convex/server"; -import { ping, vBatchQueryArgs, vBatchResult } from "./index.js"; +import { + defineBatchWorkerValidators, + ping, + vBatchQueryArgs, + vBatchResult, +} from "./index.js"; import { components, initConvexTest } from "./setup.test.js"; const schema = defineSchema({ items: defineTable({ value: v.number() }), + // For the cursor tests: `marks` is scanned from the cursor and never + // deleted, so only an advancing cursor keeps the loop from reprocessing. + marks: defineTable({ seq: v.int64() }).index("seq", ["seq"]), + // One row per batch, so a test can see how the cursor sliced the marks up. + batches: defineTable({ seqs: v.array(v.int64()) }), + // For the custom-cursor test: the cursor here is a string, not a timestamp. + letters: defineTable({ letter: v.string() }).index("letter", ["letter"]), + letterBatches: defineTable({ letters: v.array(v.string()) }), }); const WORKER = "items"; +const CURSOR_WORKER = "marks"; +const LETTER_WORKER = "letters"; export const getBatch = internalQueryGeneric({ args: vBatchQueryArgs, @@ -82,6 +97,172 @@ export const remaining = queryGeneric({ handler: async (ctx) => (await ctx.db.query("items").take(1000)).length, }); +// ── A worker driven purely by its cursor ─────────────────────────────────── + +const BATCH = 2; + +export const getMarks = internalQueryGeneric({ + args: vBatchQueryArgs, + returns: vBatchResult({ seqs: v.array(v.int64()) }), + handler: async (ctx, { cursor }) => { + const marks = await ctx.db + .query("marks") + .withIndex("seq", (q) => q.gte("seq", cursor ?? 0n)) + .take(BATCH); + if (marks.length === 0) { + return { kind: "idle" as const, cooldownMs: 100, pollIntervalMs: 10 }; + } + return { + kind: "work" as const, + batch: { seqs: marks.map((m) => m.seq as bigint) }, + // Nothing is deleted, so only this advancing cursor ends the drain. + cursor: (marks.at(-1)!.seq as bigint) + 1n, + }; + }, +}); + +export const processMarks = internalMutationGeneric({ + args: { seqs: v.array(v.int64()) }, + handler: async (ctx, { seqs }) => { + await ctx.db.insert("batches", { seqs }); + }, +}); + +// Same query, but the mutation insists the batch only got halfway. +export const processMarksPartially = internalMutationGeneric({ + args: { seqs: v.array(v.int64()) }, + handler: async (ctx, { seqs }) => { + await ctx.db.insert("batches", { seqs }); + return { cursor: seqs[0]! + 1n }; + }, +}); + +export const enqueueMark = mutationGeneric({ + args: { seq: v.int64(), partial: v.optional(v.boolean()) }, + handler: async (ctx, { seq, partial }) => { + await ctx.db.insert("marks", { seq }); + await ping(ctx, components.batchWorker, { + name: CURSOR_WORKER, + config: { debounceMs: 0 }, + workQuery: testApi.getMarks, + workerMutation: partial + ? testApi.processMarksPartially + : testApi.processMarks, + }); + }, +}); + +export const markCursor = queryGeneric({ + args: {}, + handler: async (ctx) => + ctx.runQuery(components.batchWorker.lib.getCursor, { name: CURSOR_WORKER }), +}); + +export const setMarkCursor = mutationGeneric({ + args: { cursor: v.optional(v.int64()) }, + handler: async (ctx, { cursor }) => + ctx.runMutation(components.batchWorker.lib.setCursor, { + name: CURSOR_WORKER, + cursor, + }), +}); + +export const batches = queryGeneric({ + args: {}, + handler: async (ctx) => + (await ctx.db.query("batches").take(1000)).map((b) => b.seqs), +}); + +// Seeding a cursor before the worker gets to process anything: `ping` creates +// the worker, `stop` in the same mutation keeps its loop from running. +export const seedMarks = mutationGeneric({ + args: { seqs: v.array(v.int64()), cursor: v.int64() }, + handler: async (ctx, { seqs, cursor }) => { + for (const seq of seqs) { + await ctx.db.insert("marks", { seq }); + } + await ping(ctx, components.batchWorker, { + name: CURSOR_WORKER, + config: { debounceMs: 0 }, + workQuery: testApi.getMarks, + workerMutation: testApi.processMarks, + }); + await ctx.runMutation(components.batchWorker.lib.stop, { + name: CURSOR_WORKER, + }); + await ctx.runMutation(components.batchWorker.lib.setCursor, { + name: CURSOR_WORKER, + cursor, + }); + }, +}); + +export const startMarks = mutationGeneric({ + args: {}, + handler: async (ctx) => + ctx.runMutation(components.batchWorker.lib.start, { name: CURSOR_WORKER }), +}); + +// ── A worker whose cursor isn't a commit timestamp ────────────────────────── + +const letterValidators = defineBatchWorkerValidators({ + batch: { letters: v.array(v.string()) }, + cursor: v.string(), +}); + +export const getLetters = internalQueryGeneric({ + args: letterValidators.vQueryArgs, + returns: letterValidators.vQueryReturns, + handler: async (ctx, { cursor }) => { + const rows = await ctx.db + .query("letters") + .withIndex("letter", (q) => q.gt("letter", cursor ?? "")) + .take(BATCH); + if (rows.length === 0) { + return { kind: "idle" as const, cooldownMs: 100, pollIntervalMs: 10 }; + } + return { + kind: "work" as const, + batch: { letters: rows.map((r) => r.letter) }, + cursor: rows.at(-1)!.letter, + }; + }, +}); + +export const processLetters = internalMutationGeneric({ + args: letterValidators.vMutationArgs, + returns: letterValidators.vMutationReturns, + handler: async (ctx, { letters }) => { + await ctx.db.insert("letterBatches", { letters }); + return null; + }, +}); + +export const enqueueLetter = mutationGeneric({ + args: { letter: v.string() }, + handler: async (ctx, { letter }) => { + await ctx.db.insert("letters", { letter }); + await ping(ctx, components.batchWorker, { + name: LETTER_WORKER, + config: { debounceMs: 0 }, + workQuery: testApi.getLetters, + workerMutation: testApi.processLetters, + }); + }, +}); + +export const letterCursor = queryGeneric({ + args: {}, + handler: async (ctx) => + ctx.runQuery(components.batchWorker.lib.getCursor, { name: LETTER_WORKER }), +}); + +export const letterBatches = queryGeneric({ + args: {}, + handler: async (ctx) => + (await ctx.db.query("letterBatches").take(1000)).map((b) => b.letters), +}); + const testApi = ( anyApi as unknown as ApiFromModules<{ "index.test": { @@ -92,6 +273,20 @@ const testApi = ( startWorker: typeof startWorker; stopWorker: typeof stopWorker; remaining: typeof remaining; + getMarks: typeof getMarks; + processMarks: typeof processMarks; + processMarksPartially: typeof processMarksPartially; + enqueueMark: typeof enqueueMark; + markCursor: typeof markCursor; + setMarkCursor: typeof setMarkCursor; + batches: typeof batches; + seedMarks: typeof seedMarks; + startMarks: typeof startMarks; + getLetters: typeof getLetters; + processLetters: typeof processLetters; + enqueueLetter: typeof enqueueLetter; + letterCursor: typeof letterCursor; + letterBatches: typeof letterBatches; }; }> )["index.test"]; @@ -129,3 +324,85 @@ describe("Worker client", () => { expect(await t.query(testApi.remaining, {})).toBe(0); }); }); + +describe("Cursor", () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + afterEach(() => { + vi.useRealTimers(); + }); + + test("round-trips from the query back into the next call's args", async () => { + const t = initConvexTest(schema); + for (const seq of [0n, 1n, 2n, 3n]) { + await t.mutation(testApi.enqueueMark, { seq }); + } + // Nothing deletes the marks, so the query only stops handing out the same + // batch once the cursor comes back to it as `args.cursor`. + await t.finishAllScheduledFunctions(vi.runAllTimers); + + expect(await t.query(testApi.batches, {})).toEqual([ + [0n, 1n], + [2n, 3n], + ]); + expect(await t.query(testApi.markCursor, {})).toBe(4n); + }); + + test("a cursor from the mutation overrides the query's", async () => { + const t = initConvexTest(schema); + for (const seq of [0n, 1n]) { + await t.mutation(testApi.enqueueMark, { seq, partial: true }); + } + await t.finishAllScheduledFunctions(vi.runAllTimers); + + // The query proposed 2n each time; the mutation only claimed the first + // mark of each batch, so the loop re-reads from there and inches forward. + expect(await t.query(testApi.batches, {})).toEqual([[0n, 1n], [1n]]); + expect(await t.query(testApi.markCursor, {})).toBe(2n); + }); + + test("setCursor overwrites it, and clears it when omitted", async () => { + const t = initConvexTest(schema); + await t.mutation(testApi.enqueueMark, { seq: 0n }); + await t.finishAllScheduledFunctions(vi.runAllTimers); + expect(await t.query(testApi.markCursor, {})).toBe(1n); + + await t.mutation(testApi.setMarkCursor, { cursor: 99n }); + expect(await t.query(testApi.markCursor, {})).toBe(99n); + + await t.mutation(testApi.setMarkCursor, {}); + expect(await t.query(testApi.markCursor, {})).toBe(null); + }); + + test("can be seeded before the worker processes anything", async () => { + const t = initConvexTest(schema); + await t.mutation(testApi.seedMarks, { seqs: [0n, 1n, 2n, 3n], cursor: 2n }); + await t.finishAllScheduledFunctions(vi.runAllTimers); + expect(await t.query(testApi.batches, {})).toEqual([]); + + // A ping doesn't resume a stopped worker, so the cursor stays seeded. + await t.mutation(testApi.enqueueMark, { seq: 4n }); + await t.finishAllScheduledFunctions(vi.runAllTimers); + expect(await t.query(testApi.batches, {})).toEqual([]); + + // `start` is what resumes it, and the scan begins at the seeded cursor. + await t.mutation(testApi.startMarks, {}); + await t.finishAllScheduledFunctions(vi.runAllTimers); + expect(await t.query(testApi.batches, {})).toEqual([[2n, 3n], [4n]]); + }); + + test("can be a type other than a commit timestamp", async () => { + const t = initConvexTest(schema); + for (const letter of ["a", "b", "c"]) { + await t.mutation(testApi.enqueueLetter, { letter }); + } + await t.finishAllScheduledFunctions(vi.runAllTimers); + + expect(await t.query(testApi.letterBatches, {})).toEqual([ + ["a", "b"], + ["c"], + ]); + expect(await t.query(testApi.letterCursor, {})).toBe("c"); + }); +}); diff --git a/src/client/index.ts b/src/client/index.ts index 66de188..d2a741f 100644 --- a/src/client/index.ts +++ b/src/client/index.ts @@ -11,14 +11,19 @@ import { type BatchQueryArgs, type BatchResult, type Config, + type DefaultCursor, + type WorkerResult, } from "../component/shared.js"; export { + defineBatchWorkerValidators, vBatchQueryArgs, vBatchResult, + vDefaultCursor, vWorkerResult, type BatchResult, type BatchQueryArgs, + type DefaultCursor, type WorkerResult, } from "../component/shared.js"; export type { @@ -32,11 +37,14 @@ export type { * already running). * * You provide: - * - a **work query** (args validated by {@link vBatchQueryArgs}, returns - * {@link vBatchResult}) that returns the next batch or `idle`, and + * - a **work query** that returns the next batch or `idle`, and * - a **worker mutation** that processes a batch and owns its cleanup. It may * return `{ debounceMs }` to throttle the loop. * + * Validate both with {@link defineBatchWorkerValidators}. The cursor type is + * taken from the work query's `cursor` arg, and both return types are checked + * against it. + * * @example * ```ts * export const enqueue = mutation({ @@ -52,7 +60,10 @@ export type { * }); * ``` */ -export async function ping( +export async function ping< + Batch extends DefaultFunctionArgs, + Cursor = DefaultCursor, +>( ctx: MutationCtx | ActionCtx, component: ComponentApi, args: { @@ -62,15 +73,15 @@ export async function ping( workQuery: FunctionReference< "query", "internal", - BatchQueryArgs, - BatchResult + BatchQueryArgs, + BatchResult> >; /** Processes a batch returned by the work query. */ workerMutation: FunctionReference< "mutation", "internal", Batch, - { debounceMs?: number } | null | void + WorkerResult> | void >; /** Loop configuration. */ config?: Partial; diff --git a/src/component/_generated/component.ts b/src/component/_generated/component.ts index 3cea202..da091e7 100644 --- a/src/component/_generated/component.ts +++ b/src/component/_generated/component.ts @@ -24,6 +24,13 @@ import type { FunctionReference } from "convex/server"; export type ComponentApi = { lib: { + getCursor: FunctionReference< + "query", + "internal", + { name: string }, + any, + Name + >; ping: FunctionReference< "mutation", "internal", @@ -36,6 +43,13 @@ export type ComponentApi = null, Name >; + setCursor: FunctionReference< + "mutation", + "internal", + { cursor?: any; name: string }, + null, + Name + >; start: FunctionReference< "mutation", "internal", diff --git a/src/component/functions.ts b/src/component/functions.ts index fdb3234..760f4b3 100644 --- a/src/component/functions.ts +++ b/src/component/functions.ts @@ -25,7 +25,10 @@ const withLogger = customCtx(() => ({ log: createLogger(env.LOG_LEVEL) })); export const query = customQuery(baseQuery, withLogger); export const internalQuery = customQuery(baseInternalQuery, withLogger); export const mutation = customMutation(baseMutation, withLogger); -export const internalMutation = customMutation(baseInternalMutation, withLogger); +export const internalMutation = customMutation( + baseInternalMutation, + withLogger, +); /** * Ctx types for this component's functions: the generated ctx augmented with diff --git a/src/component/kick.ts b/src/component/kick.ts index 40d1c1c..8356501 100644 --- a/src/component/kick.ts +++ b/src/component/kick.ts @@ -1,3 +1,4 @@ +import type { Value } from "convex/values"; import { internal } from "./_generated/api.js"; import type { Doc, Id } from "./_generated/dataModel.js"; import type { MutationCtx, QueryCtx } from "./functions.js"; @@ -165,14 +166,15 @@ export async function continueRunning( ctx: MutationCtx, worker: Doc<"workers">, delayMs: number, - lastWorkTs?: number, + opts?: { lastWorkTs?: number | undefined; cursor?: Value | undefined }, ): Promise { if (worker.status.kind !== "running") { await ctx.db.patch("workers", worker._id, { status: { kind: "running" } }); } await scheduleLoopRun(ctx, worker, { delayMs, - lastWorkTs, + lastWorkTs: opts?.lastWorkTs, + cursor: opts?.cursor, }); } @@ -210,7 +212,11 @@ export async function goIdle( async function scheduleLoopRun( ctx: MutationCtx, worker: Doc<"workers">, - opts: { delayMs: number; lastWorkTs?: number | undefined }, + opts: { + delayMs: number; + lastWorkTs?: number | undefined; + cursor?: Value | undefined; + }, ): Promise { const state = await getOrCreateWorkerState(ctx, worker); const generation = state.generation + 1n; @@ -219,10 +225,13 @@ async function scheduleLoopRun( internal.loop.loop, { name: worker.name, generation }, ); + // The cursor rides along on the patch this already does every iteration. + // `undefined` leaves the stored cursor as it was. await ctx.db.patch("workerState", state._id, { generation, runnerId, ...(opts.lastWorkTs !== undefined ? { lastWorkTs: opts.lastWorkTs } : {}), + ...(opts.cursor !== undefined ? { cursor: opts.cursor } : {}), }); // await ctx.db.patch("workers", worker._id, { status: worker.status }); diff --git a/src/component/lib.test.ts b/src/component/lib.test.ts index ef459aa..ca0bc1e 100644 --- a/src/component/lib.test.ts +++ b/src/component/lib.test.ts @@ -27,7 +27,9 @@ import { RUNNING_THRESHOLD_MS } from "./shared.js"; // `functions.ts`); convex-test's raw ctx has none, so add one here. Defaults to // a silent mock logger; pass one in to assert logging side-effects. const run = ( - t: { run: (fn: (ctx: Omit) => Promise) => Promise }, + t: { + run: (fn: (ctx: Omit) => Promise) => Promise; + }, fn: (ctx: MutationCtx) => Promise, log: Logger = createMockLogger(), ): Promise => t.run((ctx) => fn(Object.assign(ctx, { log }))); diff --git a/src/component/lib.ts b/src/component/lib.ts index 8dddf68..cfd1b16 100644 --- a/src/component/lib.ts +++ b/src/component/lib.ts @@ -1,6 +1,7 @@ import { v } from "convex/values"; import { mutation, query } from "./functions.js"; import { + getOrCreateWorkerState, getWorker, ping as pingHelper, start as startHelper, @@ -46,3 +47,45 @@ export const status = query({ return worker?.status; }, }); + +/** + * The cursor this worker last committed, or `null` if it doesn't have one. + * Read it when your own queries want to scan the same range the worker does + * (e.g. counting work still pending), so they skip the same tombstones. + */ +export const getCursor = query({ + args: { name: v.string() }, + returns: v.any(), + handler: async (ctx, args) => { + const worker = await getWorker(ctx, args.name); + if (!worker) return null; + const state = await ctx.db.get("workerState", worker.stateId); + return state?.cursor ?? null; + }, +}); + +/** + * Overwrite the worker's cursor. Omit `cursor` to clear it, so the next scan + * starts from the front again. + * + * Use this for migrations and recovery; in the steady state return a cursor + * from the work query. The loop writes this document every iteration, so a + * call made while the worker is busy is liable to hit a write conflict. + * + * Note: a worker must have been `ping`ed before the cursor can be set. + */ +export const setCursor = mutation({ + args: { name: v.string(), cursor: v.optional(v.any()) }, + returns: v.null(), + handler: async (ctx, args) => { + const worker = await getWorker(ctx, args.name); + if (!worker) { + throw new Error( + `[setCursor] no worker named "${args.name}". Ping it first.`, + ); + } + const state = await getOrCreateWorkerState(ctx, worker); + // Patching `undefined` removes the field. + await ctx.db.patch("workerState", state._id, { cursor: args.cursor }); + }, +}); diff --git a/src/component/loop.ts b/src/component/loop.ts index 1e62210..d2243bd 100644 --- a/src/component/loop.ts +++ b/src/component/loop.ts @@ -49,11 +49,15 @@ export const loop = internalMutation({ } const now = Date.now(); - const queryArgs: BatchQueryArgs = { name }; + // The cursor is sent once a batch has returned one. + const queryArgs: BatchQueryArgs = { + name, + ...(state.cursor !== undefined ? { cursor: state.cursor } : {}), + }; const queryRef = worker.workQuery as unknown as FunctionHandle< "query", - BatchQueryArgs, - BatchResult + BatchQueryArgs, + BatchResult >; // Stale snapshot read: no OCC dependency, so concurrent inserts while we @@ -69,11 +73,18 @@ export const loop = internalMutation({ const mutationRef = worker.workerMutation as unknown as FunctionHandle< "mutation", any, - WorkerResult + WorkerResult >; const ret = await ctx.runMutation(mutationRef, result.batch); const debounceMs = ret?.debounceMs ?? 0; - await continueRunning(ctx, worker, debounceMs, now); + // The mutation shares this transaction, so the cursor commits with the + // batch. A cursor from the mutation wins over the query's: it knows how + // far the work actually got. + const cursor = ret?.cursor !== undefined ? ret.cursor : result.cursor; + await continueRunning(ctx, worker, debounceMs, { + lastWorkTs: now, + cursor, + }); return; } @@ -118,8 +129,12 @@ export const loop = internalMutation({ /** A real (dependency-taking) read of the work query. */ async function confirmHasWork( ctx: MutationCtx, - queryRef: FunctionHandle<"query", BatchQueryArgs, BatchResult>, - queryArgs: BatchQueryArgs, + queryRef: FunctionHandle< + "query", + BatchQueryArgs, + BatchResult + >, + queryArgs: BatchQueryArgs, ): Promise { const confirm = await ctx.runQuery(queryRef, queryArgs); // Match the snapshot path's detection: a result with a `batch` is work. diff --git a/src/component/schema.ts b/src/component/schema.ts index a6ce7ef..e029909 100644 --- a/src/component/schema.ts +++ b/src/component/schema.ts @@ -34,5 +34,10 @@ export default defineSchema({ // `monitorLagMs` after the loop's next run and refreshed as it approaches. monitorId: v.optional(v.id("_scheduled_functions")), monitorRunAtMs: v.optional(v.number()), + // How far the work query said it got, passed back to it as `args.cursor`. + // Opaque here: it's whatever the app's query or worker mutation returned, + // so the app's own validators type it. Written on the same patch that + // schedules the next run. Absent until a batch returns one. + cursor: v.optional(v.any()), }), }); diff --git a/src/component/shared.ts b/src/component/shared.ts index 561d744..c22ae47 100644 --- a/src/component/shared.ts +++ b/src/component/shared.ts @@ -3,6 +3,7 @@ import { v, type Infer, type PropertyValidators, + type VCommitTs, type Validator, } from "convex/values"; @@ -66,32 +67,53 @@ export type Status = Infer; // ── The work query / worker mutation contract ────────────────────────────── /** - * The args your work query receives. Today just the worker's `name`, so a - * single query function can serve multiple named queues. Use this as your - * query's `args` validator for forward compatibility. + * The default cursor type: a commit timestamp, as written by + * `ctx.db.vars.commitTs` and read back as a `bigint`. */ -export const vBatchQueryArgs = v.object({ name: v.string() }); -export type BatchQueryArgs = Infer; +export const vDefaultCursor = v.commitTs(); +export type DefaultCursor = Infer; /** - * Builds the validator for what your work query returns: either a batch of - * work to process, or an explicit `idle` (optionally with a `timeoutMs` hint - * for when to check again — e.g. when the next item is scheduled). + * The args your work query receives: the worker's `name` (so a single query + * function can serve multiple named queues) and the `cursor` this worker last + * committed, if any. * - * @example - * export const getBatch = internalQuery({ - * args: vBatchQueryArgs, - * returns: vBatchResult(v.object({ ids: v.array(v.id("tasks")) })), - * handler: ... - * }); + * @deprecated Use {@link defineBatchWorkerValidators} to obtain all four + * argument and return validators, so the batch shape and the cursor type match + * between the query and the mutation. */ -export function vBatchResult< +export const vBatchQueryArgs = v.object({ + name: v.string(), + cursor: v.optional(vDefaultCursor), +}); +export type BatchQueryArgs = { + name: string; + cursor?: Cursor; +}; + +/** + * What your work query returns: either a batch of work to process, or an + * explicit `idle` (optionally with a `timeoutMs` hint for when to check again + * — e.g. when the next item is scheduled). + * + * Alongside a batch, return a `cursor` saying how far this batch got. The + * component commits it with the batch, and hands it back as `args.cursor` on + * the next call so the scan resumes there. It is committed only if the worker + * mutation commits. + * + * `idle` ends the loop until a ping wakes it, so it carries no cursor. To skip + * past a stretch you scanned and found nothing to do in, return `work` with an + * empty batch and the advanced cursor. + */ +function vBatchResultFor< B extends Validator | PropertyValidators, ->(batch: B) { + C extends Validator, +>(batch: B, cursor: C) { return v.union( v.object({ kind: v.optional(v.literal("work")), batch: asObjectValidator(batch), + cursor: v.optional(cursor), }), v.object({ kind: v.literal("idle"), @@ -114,14 +136,36 @@ export function vBatchResult< ); } +/** + * Builds the validator for what your work query returns: a batch with an + * optional commit-timestamp `cursor`, or `idle`. + * + * @deprecated Use {@link defineBatchWorkerValidators} to obtain all four + * argument and return validators, so the batch shape and the cursor type match + * between the query and the mutation. + */ +export function vBatchResult< + B extends Validator | PropertyValidators, +>(batch: B) { + return vBatchResultFor(batch, vDefaultCursor); +} + /** * What a work query returns: a `batch` of work to process, or `idle` with an * optional `timeoutMs` hint for when to look again. * * @typeParam Batch - the shape passed to your worker mutation. */ -export type BatchResult = - | { kind: "work"; batch: Batch } +export type BatchResult = + | { + kind?: "work"; + batch: Batch; + /** + * How far this batch got. Committed with the batch and passed to the + * next query call as `args.cursor`. Omit to leave the cursor untouched. + */ + cursor?: Cursor; + } | { kind: "idle"; /** @@ -143,14 +187,82 @@ export type BatchResult = * What a worker mutation may return to steer the loop. Returning nothing (or * null) re-runs immediately (drain as fast as possible). */ -export const vWorkerResult = v.union( - v.null(), - v.object({ - /** - * Don't run again — and ignore pings — for at least this long. Use to - * debounce / batch. - */ - debounceMs: v.optional(v.number()), - }), -); -export type WorkerResult = Infer; +function vWorkerResultFor>( + cursor: C, +) { + return v.union( + v.null(), + v.object({ + /** + * Don't run again — and ignore pings — for at least this long. Use to + * debounce / batch. + */ + debounceMs: v.optional(v.number()), + /** + * Overrides the cursor the work query returned. Return one when the batch + * made partial progress and you can derive how far it got, e.g. from a + * per-item timestamp carried in the batch. + */ + cursor: v.optional(cursor), + }), + ); +} + +/** + * What a worker mutation may return to steer the loop. + * + * @deprecated Use {@link defineBatchWorkerValidators} to obtain all four + * argument and return validators, so the batch shape and the cursor type match + * between the query and the mutation. + */ +export const vWorkerResult = vWorkerResultFor(vDefaultCursor); +export type WorkerResult = null | { + debounceMs?: number; + cursor?: Cursor; +}; + +/** + * Builds all four of a worker's argument and return validators from a single + * declaration, so the batch shape and the cursor type match between the work + * query and the worker mutation. + * + * `cursor` defaults to a commit timestamp ({@link vDefaultCursor}); pass one to + * key the worker on something else, like a `paginator` cursor. + * + * @example + * ```ts + * const { vQueryArgs, vQueryReturns, vMutationArgs, vMutationReturns } = + * defineBatchWorkerValidators({ + * batch: { ids: v.array(v.id("tasks")) }, + * cursor: v.string(), // optional; defaults to v.commitTs() + * }); + * + * export const getBatch = internalQuery({ + * args: vQueryArgs, + * returns: vQueryReturns, + * handler: async (ctx, { cursor }) => ..., + * }); + * + * export const processBatch = internalMutation({ + * args: vMutationArgs, + * returns: vMutationReturns, + * handler: async (ctx, { ids }) => ..., + * }); + * ``` + */ +export function defineBatchWorkerValidators< + B extends Validator | PropertyValidators, + C extends Validator = VCommitTs, +>(spec: { batch: B; cursor?: C }) { + const cursor = (spec.cursor ?? vDefaultCursor) as C; + return { + /** Your work query's `args`: `{ name, cursor? }`. */ + vQueryArgs: v.object({ name: v.string(), cursor: v.optional(cursor) }), + /** Your work query's `returns`: a batch with an optional cursor, or idle. */ + vQueryReturns: vBatchResultFor(spec.batch, cursor), + /** Your worker mutation's `args`: the batch the query returned. */ + vMutationArgs: asObjectValidator(spec.batch), + /** Your worker mutation's `returns`: null, or `{ debounceMs?, cursor? }`. */ + vMutationReturns: vWorkerResultFor(cursor), + }; +}