From 3dd5055b8d554ed8d775d2f9fd9bb0077c4cb8cf Mon Sep 17 00:00:00 2001 From: ABCxFF <79597906+abcxff@users.noreply.github.com> Date: Thu, 3 Sep 2026 12:18:36 -0400 Subject: [PATCH] fix(workflow-engine): clear pruned loop history via batched, bounded deletes Port of the workflows-repo fix. Add batchDelete to the EngineDriver interface and all implementations, and route loop-history pruning through runDeletes(), which coalesces keys into transaction-sized (MAX_KV_BATCH_ENTRIES) batchDelete chunks run in bounded rounds (MAX_CONCURRENT_DELETES=64). Previously deleteEntriesWithPrefix / flush fanned out one unbounded Promise.all of single-statement deletes per key; cutting a loop history of >128 entries exceeded the actor SQLite transaction coordinator's 128-permit admission cap (non-blocking try_acquire -> transaction_queue_full). --- .../packages/rivetkit/src/workflow/driver.ts | 13 ++ .../packages/workflow-engine/src/driver.ts | 6 + .../packages/workflow-engine/src/storage.ts | 58 ++++++--- .../packages/workflow-engine/src/testing.ts | 7 ++ .../tests/prune-fanout.test.ts | 111 ++++++++++++++++++ 5 files changed, 181 insertions(+), 14 deletions(-) create mode 100644 rivetkit-typescript/packages/workflow-engine/tests/prune-fanout.test.ts diff --git a/rivetkit-typescript/packages/rivetkit/src/workflow/driver.ts b/rivetkit-typescript/packages/rivetkit/src/workflow/driver.ts index bce8c52032..feb3b2cc9d 100644 --- a/rivetkit-typescript/packages/rivetkit/src/workflow/driver.ts +++ b/rivetkit-typescript/packages/rivetkit/src/workflow/driver.ts @@ -116,6 +116,11 @@ class WorkflowStorage { ); } + async batchDelete(keys: Uint8Array[]): Promise { + if (keys.length === 0) return; + await this.deleteRawKeys(keys.map((key) => makeWorkflowKey(key))); + } + async deletePrefix(prefix: Uint8Array): Promise { const start = makeWorkflowKey(prefix); const end = computeUpperBound(start); @@ -344,6 +349,10 @@ export class ActorWorkflowDriver implements EngineDriver { await this.#runCtx.internalKeepAwake(this.#storage.delete(key)); } + async batchDelete(keys: Uint8Array[]): Promise { + await this.#runCtx.internalKeepAwake(this.#storage.batchDelete(keys)); + } + async deletePrefix(prefix: Uint8Array): Promise { await this.#runCtx.internalKeepAwake( this.#storage.deletePrefix(prefix), @@ -438,6 +447,10 @@ export class ActorWorkflowControlDriver implements EngineDriver { await this.#storage.delete(key); } + async batchDelete(keys: Uint8Array[]): Promise { + await this.#storage.batchDelete(keys); + } + async deletePrefix(prefix: Uint8Array): Promise { await this.#storage.deletePrefix(prefix); } diff --git a/rivetkit-typescript/packages/workflow-engine/src/driver.ts b/rivetkit-typescript/packages/workflow-engine/src/driver.ts index 808d3c0377..c9905a87ba 100644 --- a/rivetkit-typescript/packages/workflow-engine/src/driver.ts +++ b/rivetkit-typescript/packages/workflow-engine/src/driver.ts @@ -54,6 +54,12 @@ export interface EngineDriver { */ delete(key: Uint8Array): Promise; + /** + * Batch delete multiple keys in a single operation. + * Should be atomic if possible. + */ + batchDelete(keys: Uint8Array[]): Promise; + /** * Delete all keys with a given prefix. */ diff --git a/rivetkit-typescript/packages/workflow-engine/src/storage.ts b/rivetkit-typescript/packages/workflow-engine/src/storage.ts index 3073c55680..2ccad628af 100644 --- a/rivetkit-typescript/packages/workflow-engine/src/storage.ts +++ b/rivetkit-typescript/packages/workflow-engine/src/storage.ts @@ -43,6 +43,9 @@ import type { export const MAX_KV_BATCH_ENTRIES = 128; export const MAX_KV_BATCH_PAYLOAD_BYTES = 976 * 1024; +/** Max delete ops (one transaction/permit each) run at once, under the 128-permit cap. */ +export const MAX_CONCURRENT_DELETES = 64; + /** * Create an empty storage instance. */ @@ -330,18 +333,8 @@ export async function flush( // Apply pending deletions after the batch write. These are collected // by collectLoopPruning so pruning happens alongside the state write. if (pendingDeletions) { - const deleteOps: Promise[] = []; - for (const prefix of pendingDeletions.prefixes) { - deleteOps.push(driver.deletePrefix(prefix)); - } - for (const range of pendingDeletions.ranges) { - deleteOps.push(driver.deleteRange(range.start, range.end)); - } - for (const key of pendingDeletions.keys) { - deleteOps.push(driver.delete(key)); - } - if (deleteOps.length > 0) { - await Promise.all(deleteOps); + const didChange = await runDeletes(driver, pendingDeletions); + if (didChange) { historyUpdated = true; } } @@ -397,6 +390,44 @@ function splitBatchWrites(writes: KVWrite[]): KVWrite[][] { return chunks; } +/** + * Split delete keys into batches within one KV transaction (KV_TX_MAX_ROWS). + */ +function splitBatchDeletes(keys: Uint8Array[]): Uint8Array[][] { + const chunks: Uint8Array[][] = []; + for (let i = 0; i < keys.length; i += MAX_KV_BATCH_ENTRIES) { + chunks.push(keys.slice(i, i + MAX_KV_BATCH_ENTRIES)); + } + return chunks; +} + +/** + * Apply deletions concurrently in bounded rounds; returns whether anything was deleted. + */ +async function runDeletes( + driver: EngineDriver, + deletions: PendingDeletions, +): Promise { + const ops = [ + ...deletions.prefixes.map((prefix) => () => driver.deletePrefix(prefix)), + ...deletions.ranges.map( + (range) => () => driver.deleteRange(range.start, range.end), + ), + ...splitBatchDeletes(deletions.keys).map( + (chunk) => () => driver.batchDelete(chunk), + ), + ]; + if (ops.length === 0) { + return false; + } + for (let i = 0; i < ops.length; i += MAX_CONCURRENT_DELETES) { + await Promise.all( + ops.slice(i, i + MAX_CONCURRENT_DELETES).map((op) => op()), + ); + } + return true; +} + /** * Delete entries with a given location prefix (used for loop forgetting). * Also cleans up associated metadata from both memory and driver. @@ -410,8 +441,7 @@ export async function deleteEntriesWithPrefix( const deletions = collectDeletionsForPrefix(storage, prefixLocation); // Apply deletions to driver - await driver.deletePrefix(deletions.prefixes[0]!); - await Promise.all(deletions.keys.map((key) => driver.delete(key))); + await runDeletes(driver, deletions); if (deletions.keys.length > 0 && onHistoryUpdated) { onHistoryUpdated(); diff --git a/rivetkit-typescript/packages/workflow-engine/src/testing.ts b/rivetkit-typescript/packages/workflow-engine/src/testing.ts index f57984cf16..259db41177 100644 --- a/rivetkit-typescript/packages/workflow-engine/src/testing.ts +++ b/rivetkit-typescript/packages/workflow-engine/src/testing.ts @@ -173,6 +173,13 @@ export class InMemoryDriver implements EngineDriver { this.kv.delete(keyToHex(key)); } + async batchDelete(keys: Uint8Array[]): Promise { + await sleep(this.latency); + for (const key of keys) { + this.kv.delete(keyToHex(key)); + } + } + async deletePrefix(prefix: Uint8Array): Promise { await sleep(this.latency); for (const [hexKey, entry] of this.kv) { diff --git a/rivetkit-typescript/packages/workflow-engine/tests/prune-fanout.test.ts b/rivetkit-typescript/packages/workflow-engine/tests/prune-fanout.test.ts new file mode 100644 index 0000000000..971add2e77 --- /dev/null +++ b/rivetkit-typescript/packages/workflow-engine/tests/prune-fanout.test.ts @@ -0,0 +1,111 @@ +import { describe, expect, it } from "vitest"; +import { + deleteEntriesWithPrefix, + MAX_CONCURRENT_DELETES, + MAX_KV_BATCH_ENTRIES, +} from "../src/storage.js"; +import { + appendName, + createEntry, + createStorage, + emptyLocation, + InMemoryDriver, + setEntry, +} from "../src/testing.js"; + +describe("Workflow Engine Storage delete fan-out", () => { + // Records batchDelete sizes to assert keys are coalesced, not deleted one-by-one. + class BatchDeleteRecordingDriver extends InMemoryDriver { + batchSizes: number[] = []; + singleDeletes = 0; + + override async batchDelete(keys: Uint8Array[]): Promise { + this.batchSizes.push(keys.length); + await super.batchDelete(keys); + } + + override async delete(key: Uint8Array): Promise { + this.singleDeletes++; + await super.delete(key); + } + } + + it("clears a large history prefix in transaction-sized delete batches", async () => { + const driver = new BatchDeleteRecordingDriver(); + driver.latency = 1; + const storage = createStorage(); + const loopLocation = appendName(storage, emptyLocation(), "loop"); + + // Span several batches so chunking is exercised. + const entryCount = MAX_KV_BATCH_ENTRIES * 3 + 7; + for (let i = 0; i < entryCount; i++) { + const location = appendName(storage, loopLocation, `iter-${i}`); + const entry = createEntry(location, { + type: "step", + data: { output: i }, + }); + setEntry(storage, location, entry); + } + + await deleteEntriesWithPrefix(storage, driver, loopLocation); + + // All keys deleted via transaction-sized batches, no per-key fan-out. + expect(driver.singleDeletes).toBe(0); + expect(driver.batchSizes).toHaveLength( + Math.ceil(entryCount / MAX_KV_BATCH_ENTRIES), + ); + for (const size of driver.batchSizes) { + expect(size).toBeLessThanOrEqual(MAX_KV_BATCH_ENTRIES); + } + expect(driver.batchSizes.reduce((a, b) => a + b, 0)).toBe(entryCount); + expect(storage.history.entries.size).toBe(0); + }); + + // Tracks concurrent delete ops so the test can assert the fan-out stays bounded. + class ConcurrencyTrackingDriver extends InMemoryDriver { + inFlight = 0; + peakInFlight = 0; + + async #track(op: Promise): Promise { + this.inFlight++; + this.peakInFlight = Math.max(this.peakInFlight, this.inFlight); + try { + return await op; + } finally { + this.inFlight--; + } + } + + override batchDelete(keys: Uint8Array[]): Promise { + return this.#track(super.batchDelete(keys)); + } + + override deletePrefix(prefix: Uint8Array): Promise { + return this.#track(super.deletePrefix(prefix)); + } + } + + it("bounds concurrent delete ops for a prune larger than the cap", async () => { + const driver = new ConcurrencyTrackingDriver(); + driver.latency = 1; + const storage = createStorage(); + const loopLocation = appendName(storage, emptyLocation(), "loop"); + + // Enough keys to yield more batches than MAX_CONCURRENT_DELETES. + const entryCount = MAX_CONCURRENT_DELETES * MAX_KV_BATCH_ENTRIES + 1; + for (let i = 0; i < entryCount; i++) { + const location = appendName(storage, loopLocation, `iter-${i}`); + const entry = createEntry(location, { + type: "step", + data: { output: i }, + }); + setEntry(storage, location, entry); + } + + await deleteEntriesWithPrefix(storage, driver, loopLocation); + + expect(driver.peakInFlight).toBeLessThanOrEqual(MAX_CONCURRENT_DELETES); + expect(driver.peakInFlight).toBe(MAX_CONCURRENT_DELETES); + expect(storage.history.entries.size).toBe(0); + }); +});