diff --git a/packages/task-graph/README.md b/packages/task-graph/README.md index 21d6c2b0b..f5e68a61f 100644 --- a/packages/task-graph/README.md +++ b/packages/task-graph/README.md @@ -832,7 +832,13 @@ If the registered `private` slot is present and the graph contains any task whos ```typescript import { CacheJanitor } from "@workglow/task-graph"; -const janitor = new CacheJanitor({ privateBacking }); +const janitor = new CacheJanitor({ + privateBacking, + // Snapshot of currently-live run IDs; the janitor excludes these from the + // sweep so a long-running in-flight run does not have its cache rows deleted. + // Pass `() => []` if no run is ever concurrent with the sweep. + liveRunIds: () => runner.activeRunIds(), +}); // Sweep run-private rows older than 24 hours. await janitor.sweepStaleRunPrivate(24 * 60 * 60 * 1000); ``` diff --git a/packages/task-graph/src/cache/CacheJanitor.ts b/packages/task-graph/src/cache/CacheJanitor.ts index 1591db37a..875920824 100644 --- a/packages/task-graph/src/cache/CacheJanitor.ts +++ b/packages/task-graph/src/cache/CacheJanitor.ts @@ -19,11 +19,11 @@ export interface CacheJanitorOptions { * in-flight long-running run cannot have its cache rows deleted out from * under it (a row's `createdAt` predates the run's age, not its current * activity, so a still-running run started days ago would otherwise be - * reaped). Default returns an empty list AND emits a one-time `console.warn` - * — defaulting is only safe when no run is ever active concurrent with the - * sweep. + * reaped). Callers that never sweep concurrent with a live run should pass + * `() => []` explicitly — the argument is required so the empty case is a + * conscious choice at the call site rather than a silent default. */ - liveRunIds?: () => Iterable; + liveRunIds: () => Iterable; } /** @@ -41,23 +41,10 @@ export interface CacheJanitorOptions { export class CacheJanitor { private readonly privateBacking: RunPrivateTaskOutputRepository; private readonly liveRunIds: () => Iterable; - private static defaultLiveRunIdsWarned = false; constructor({ privateBacking, liveRunIds }: CacheJanitorOptions) { this.privateBacking = privateBacking; - if (liveRunIds === undefined) { - if (!CacheJanitor.defaultLiveRunIdsWarned) { - CacheJanitor.defaultLiveRunIdsWarned = true; - // No logger dependency here (task-graph keeps its log surface narrow); - // a one-time console.warn is enough to surface the misconfiguration. - console.warn( - "CacheJanitor: no liveRunIds callback provided — sweeps may delete cache rows for in-flight runs. Pass `liveRunIds: () => activeRunIdSnapshot` to suppress." - ); - } - this.liveRunIds = () => []; - } else { - this.liveRunIds = liveRunIds; - } + this.liveRunIds = liveRunIds; } async sweepStaleRunPrivate(olderThanMs: number): Promise { diff --git a/packages/task-graph/src/storage/RunPrivateTaskOutputRepository.ts b/packages/task-graph/src/storage/RunPrivateTaskOutputRepository.ts index 5d2ede0f6..2187615ab 100644 --- a/packages/task-graph/src/storage/RunPrivateTaskOutputRepository.ts +++ b/packages/task-graph/src/storage/RunPrivateTaskOutputRepository.ts @@ -4,7 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { ITabularStorage } from "@workglow/storage"; +import type { ITabularStorage, PageCursor } from "@workglow/storage"; import { makeFingerprint } from "@workglow/util"; import type { TaskInput, TaskOutput } from "../task/TaskTypes"; import { @@ -99,10 +99,21 @@ export class RunPrivateTaskOutputRepository extends TaskOutputRepository { override async deleteRunOlderThan(runId: string, olderThanInMs: number): Promise { const cutoff = new Date(Date.now() - olderThanInMs).toISOString(); - await this.storage.deleteSearch({ runId, createdAt: { value: cutoff, operator: "<" } }); + await this.deleteRunOlderThanAt(runId, cutoff); this.emit("output_pruned"); } + /** + * Internal — timestamp-parameterized older-than delete used by both + * {@link deleteRunOlderThan} (which computes the cutoff from `olderThanInMs` + * on entry) and the all-runs sweep (which computes the cutoff once). Does + * NOT emit `output_pruned`; callers own the event so the sweep can emit + * exactly once at the end. + */ + private async deleteRunOlderThanAt(runId: string, cutoff: string): Promise { + await this.storage.deleteSearch({ runId, createdAt: { value: cutoff, operator: "<" } }); + } + override async sizeForRun(runId: string): Promise { // count() tallies matching rows without loading (and decoding) every blob. return await this.storage.count({ runId }); @@ -117,30 +128,41 @@ export class RunPrivateTaskOutputRepository extends TaskOutputRepository { * a run that started long ago but is still active must not have its cache rows * deleted out from under it. The tabular surface has no `NOT IN` operator * ({@link SearchOperator} is `=//>=` only), so the implementation - * enumerates distinct old `runId`s and calls {@link deleteRunOlderThan} per - * runId not in the exclude set. + * enumerates distinct old `runId`s via cursor-paginated `queryPage` — bounded + * page reads instead of materializing the full stale row set in memory — and + * calls {@link deleteRunOlderThan} per runId not in the exclude set. */ async clearOlderThan( olderThanInMs: number, excludeRunIds: ReadonlySet = new Set() ): Promise { const cutoff = new Date(Date.now() - olderThanInMs).toISOString(); - if (excludeRunIds.size === 0) { - await this.storage.deleteSearch({ createdAt: { value: cutoff, operator: "<" } }); - this.emit("output_pruned"); - return; - } - const oldRows = await this.storage.search({ - createdAt: { value: cutoff, operator: "<" }, - }); - const oldRunIds = new Set(); - for (const row of oldRows ?? []) { - const runId = (row as { runId?: unknown }).runId; - if (typeof runId === "string") oldRunIds.add(runId); - } - for (const runId of oldRunIds) { + const criteria = { createdAt: { value: cutoff, operator: "<" as const } }; + const seenRunIds = new Set(); + + // Default PK ordering (`runId` ASC — leading PK column) groups a run's rows + // contiguously, so distinct runIds accumulate quickly on typical data. The + // page size caps per-batch memory; `seenRunIds` is bounded by distinct + // runId count, not row count. + const pageLimit = 500; + let cursor: PageCursor | undefined; + // Termination follows the ITabularStorage contract: bail on empty page or + // undefined nextCursor (whichever comes first). + do { + const page = await this.storage.queryPage(criteria, { limit: pageLimit, cursor }); + for (const row of page.items) { + const runId = (row as { runId?: unknown }).runId; + if (typeof runId === "string") seenRunIds.add(runId); + } + if (page.items.length === 0) break; + cursor = page.nextCursor; + } while (cursor !== undefined); + + for (const runId of seenRunIds) { if (excludeRunIds.has(runId)) continue; - await this.deleteRunOlderThan(runId, olderThanInMs); + // Route through the internal helper so per-run delete does NOT emit + // `output_pruned`; the sweep emits once at the end below. + await this.deleteRunOlderThanAt(runId, cutoff); } this.emit("output_pruned"); } diff --git a/packages/test/src/test/task-graph-cache/CacheJanitor.test.ts b/packages/test/src/test/task-graph-cache/CacheJanitor.test.ts index aca8399b4..0ffebb902 100644 --- a/packages/test/src/test/task-graph-cache/CacheJanitor.test.ts +++ b/packages/test/src/test/task-graph-cache/CacheJanitor.test.ts @@ -30,7 +30,7 @@ describe("CacheJanitor", () => { ); expect(await backing.size()).toBe(2); - const janitor = new CacheJanitor({ privateBacking: backing }); + const janitor = new CacheJanitor({ privateBacking: backing, liveRunIds: () => [] }); await janitor.sweepStaleRunPrivate(7 * 24 * 3600_000); expect(await backing.size()).toBe(1); @@ -46,10 +46,110 @@ describe("CacheJanitor", () => { await backing.saveOutputForRun("rX", "T", { x: 1 }, { ok: "x" }, old); await backing.saveOutputForRun("rY", "T", { x: 1 }, { ok: "y" }, old); - const janitor = new CacheJanitor({ privateBacking: backing }); + const janitor = new CacheJanitor({ privateBacking: backing, liveRunIds: () => [] }); await janitor.sweepStaleRunPrivate(7 * 24 * 3600_000); // Every row is run-private; all are older than the cutoff, so all are reaped. expect(await backing.size()).toBe(0); }); + + it("clearOlderThan with a non-empty excludeRunIds does not throw", async () => { + // Regression: clearOlderThan called this.storage.search(...) which does not + // exist on ITabularStorage — the correct method is query(). This smoke test + // exercises the non-empty-excludeRunIds branch so the fix stays regressed. + const backing = new RunPrivateInMemoryTaskOutputRepository(); + await (backing as any).setupDatabase?.(); + + const old = new Date(Date.now() - 30 * 24 * 3600_000); + await backing.saveOutputForRun("live", "T", { x: 1 }, { ok: "live" }, old); + await backing.saveOutputForRun("stale", "T", { x: 1 }, { ok: "stale" }, old); + + await expect(backing.clearOlderThan(1, new Set(["live"]))).resolves.not.toThrow(); + }); + + it("omitting liveRunIds is a compile-time error", () => { + const backing = new RunPrivateInMemoryTaskOutputRepository(); + // @ts-expect-error - liveRunIds is a required option. + const janitor = new CacheJanitor({ privateBacking: backing }); + expect(janitor).toBeInstanceOf(CacheJanitor); + }); + + it("sweepStaleRunPrivate does not materialize the full stale row set at once", async () => { + // Regression: the previous implementation called `storage.query({...})` + // for ALL stale rows in a single call. The new implementation drives + // distinct-runId collection through cursor-paginated `queryPage` with a + // bounded `limit`, so SQL backends (which push both the WHERE and the + // keyset predicate into the SELECT) issue O(pageSize) round-trips. + // + // We seed 5000 stale rows across 3 runs and assert: + // (a) `queryPage` is invoked repeatedly (≥10 calls, matching the 500-row + // page size against 5000 rows), i.e. it truly paginates instead of + // one-shotting. + // (b) Every `queryPage` call carries a `request.limit`, so the memory + // profile is bounded per call. + // (c) The sweep completes and reaps all stale rows. + // + // (The `BaseTabularStorage` fallback engine used by InMemoryTabularStorage + // still fetches the underlying rows in one `query()` for composite-PK + // tables — that is a backend-level concern; SQL backends override + // `queryPage` with pushdown. What we're pinning here is the + // repository-level contract: it must ask for pages, not the whole table.) + const backing = new RunPrivateInMemoryTaskOutputRepository(); + await (backing as any).setupDatabase?.(); + + const old = new Date(Date.now() - 30 * 24 * 3600_000); + const runIds = ["r-a", "r-b", "r-c"]; + for (let i = 0; i < 5000; i++) { + const runId = runIds[i % runIds.length]; + await backing.saveOutputForRun(runId, "T", { i }, { ok: i }, old); + } + expect(await backing.size()).toBe(5000); + + const storage = backing.storage as any; + const originalQueryPage = storage.queryPage.bind(storage); + const queryPageCalls: Array<{ limit: number | undefined }> = []; + storage.queryPage = (criteria: unknown, request: { limit?: number } = {}) => { + queryPageCalls.push({ limit: request.limit }); + return originalQueryPage(criteria, request); + }; + + try { + const janitor = new CacheJanitor({ privateBacking: backing, liveRunIds: () => [] }); + await janitor.sweepStaleRunPrivate(7 * 24 * 3600_000); + } finally { + storage.queryPage = originalQueryPage; + } + + expect(queryPageCalls.length).toBeGreaterThanOrEqual(10); + // Every page fetch carries a bounded limit — no unbounded reads. + for (const call of queryPageCalls) { + expect(call.limit).toBeTypeOf("number"); + expect(call.limit).toBeGreaterThan(0); + } + expect(await backing.size()).toBe(0); + }); + + it("empty liveRunIds still respects the per-run indexed delete path", async () => { + // Seed 3 stale runs + 1 live-but-excluded run. All rows are older than the + // cutoff; the excluded live-run rows must survive and only the 3 stale + // runs' rows should be reaped. + const backing = new RunPrivateInMemoryTaskOutputRepository(); + await (backing as any).setupDatabase?.(); + + const old = new Date(Date.now() - 30 * 24 * 3600_000); + await backing.saveOutputForRun("stale-1", "T", { x: 1 }, { ok: 1 }, old); + await backing.saveOutputForRun("stale-2", "T", { x: 1 }, { ok: 2 }, old); + await backing.saveOutputForRun("stale-3", "T", { x: 1 }, { ok: 3 }, old); + await backing.saveOutputForRun("live", "T", { x: 1 }, { ok: "live" }, old); + expect(await backing.size()).toBe(4); + + const janitor = new CacheJanitor({ privateBacking: backing, liveRunIds: () => ["live"] }); + await janitor.sweepStaleRunPrivate(7 * 24 * 3600_000); + + expect(await backing.sizeForRun("stale-1")).toBe(0); + expect(await backing.sizeForRun("stale-2")).toBe(0); + expect(await backing.sizeForRun("stale-3")).toBe(0); + expect(await backing.sizeForRun("live")).toBe(1); + expect(await backing.getOutputForRun("live", "T", { x: 1 })).toEqual({ ok: "live" }); + }); }); diff --git a/packages/test/src/test/task-graph-output-cache/RunPrivateTaskOutputRepository.test.ts b/packages/test/src/test/task-graph-output-cache/RunPrivateTaskOutputRepository.test.ts new file mode 100644 index 000000000..2c5258105 --- /dev/null +++ b/packages/test/src/test/task-graph-output-cache/RunPrivateTaskOutputRepository.test.ts @@ -0,0 +1,38 @@ +/** + * @license + * Copyright 2026 Steven Roussey + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it } from "vitest"; +import { RunPrivateInMemoryTaskOutputRepository } from "../../binding/RunPrivateInMemoryTaskOutputRepository"; + +describe("RunPrivateTaskOutputRepository.clearOlderThan", () => { + it("preserves live runs even when they have stale rows", async () => { + // Regression: the sweep must exclude runIds in `excludeRunIds` even when + // every row on that run is older than the cutoff (a still-running job that + // started long ago would otherwise have its cache reaped mid-flight). + const backing = new RunPrivateInMemoryTaskOutputRepository(); + await (backing as any).setupDatabase?.(); + + const old = new Date(Date.now() - 30 * 24 * 3600_000); + // Stale run with several rows. + await backing.saveOutputForRun("stale", "T", { x: 1 }, { ok: 1 }, old); + await backing.saveOutputForRun("stale", "T", { x: 2 }, { ok: 2 }, old); + await backing.saveOutputForRun("stale", "T", { x: 3 }, { ok: 3 }, old); + // Live run whose rows also predate the cutoff. + await backing.saveOutputForRun("live", "T", { x: 1 }, { ok: "l1" }, old); + await backing.saveOutputForRun("live", "T", { x: 2 }, { ok: "l2" }, old); + + expect(await backing.sizeForRun("stale")).toBe(3); + expect(await backing.sizeForRun("live")).toBe(2); + + await backing.clearOlderThan(7 * 24 * 3600_000, new Set(["live"])); + + expect(await backing.sizeForRun("stale")).toBe(0); + expect(await backing.sizeForRun("live")).toBe(2); + // Every live-run row still readable end-to-end. + expect(await backing.getOutputForRun("live", "T", { x: 1 })).toEqual({ ok: "l1" }); + expect(await backing.getOutputForRun("live", "T", { x: 2 })).toEqual({ ok: "l2" }); + }); +}); diff --git a/packages/test/src/test/util/Container.test.ts b/packages/test/src/test/util/Container.test.ts index 7e4b3b101..6060df3c6 100644 --- a/packages/test/src/test/util/Container.test.ts +++ b/packages/test/src/test/util/Container.test.ts @@ -174,6 +174,87 @@ describe("Container", () => { expect(disposed).toBe(0); }); + it("register schedules eviction dispose and awaitReplacement resolves after it completes", async () => { + // Deferred pattern so we can observe: (1) disposer hasn't run yet + // synchronously after `register`, (2) resolving lets awaitReplacement + // settle, (3) `disposeCalled` is only true after that. + let resolveGate: (() => void) | undefined; + const gate = new Promise((resolve) => { + resolveGate = resolve; + }); + let disposeCalled = false; + container.register("svc", () => ({ + async dispose() { + await gate; + disposeCalled = true; + }, + })); + // Force instantiation so the singleton is cached. + container.get("svc"); + + container.register("svc", () => ({ value: "second" })); + // Disposer is still awaiting the gate; nothing has completed yet. + expect(disposeCalled).toBe(false); + + resolveGate!(); + await container.awaitReplacement("svc"); + expect(disposeCalled).toBe(true); + }); + + it("container.dispose() drains pending disposals from prior register replacements", async () => { + let resolveGate: (() => void) | undefined; + const gate = new Promise((resolve) => { + resolveGate = resolve; + }); + let disposeCalled = false; + container.register("svc", () => ({ + async dispose() { + await gate; + disposeCalled = true; + }, + })); + container.get("svc"); + + container.register("svc", () => ({ value: "second" })); + // Kick off dispose(); it must wait for the pending eviction disposer. + const disposePromise = container.dispose(); + expect(disposeCalled).toBe(false); + resolveGate!(); + await disposePromise; + expect(disposeCalled).toBe(true); + }); + + it("registerInstance eviction is drained by awaitReplacement", async () => { + let resolveGate: (() => void) | undefined; + const gate = new Promise((resolve) => { + resolveGate = resolve; + }); + let disposeCalled = false; + const first = { + async dispose() { + await gate; + disposeCalled = true; + }, + }; + container.registerInstance("svc", first); + container.registerInstance("svc", { value: "second" }); + expect(disposeCalled).toBe(false); + resolveGate!(); + await container.awaitReplacement("svc"); + expect(disposeCalled).toBe(true); + }); + + it("registerInstance clears any previously-registered factory for the token", () => { + container.register("svc", () => ({ value: "from-factory" })); + // Do not force-instantiate; the factory sits in `factories` with no + // cached entry in `services`. + container.registerInstance("svc", { value: "from-instance" }); + // The factory must be cleared so a subsequent remove() → get() cannot + // resurrect the factory-built object. + expect((container as any).factories.has("svc")).toBe(false); + expect(container.get<{ value: string }>("svc").value).toBe("from-instance"); + }); + it("survives a buggy disposer on the eviction path", async () => { const originalWarn = console.warn; console.warn = () => {}; diff --git a/packages/util/src/di/Container.ts b/packages/util/src/di/Container.ts index 112ae7749..47e30f6f0 100644 --- a/packages/util/src/di/Container.ts +++ b/packages/util/src/di/Container.ts @@ -18,6 +18,15 @@ export class Container { * disposing this container must NOT dispose them. */ private inheritedServices: Set = new Set(); + /** + * In-flight eviction disposals keyed by token. `register()` and + * `registerInstance()` invoke the previous singleton's disposer without + * blocking the registration; the returned promise lives here so + * {@link dispose} and {@link awaitReplacement} can drain it. Entries are + * identity-guarded so a subsequent replacement's promise doesn't get erased + * by an earlier disposer's `.finally()`. + */ + private pendingDisposals: Map> = new Map(); /** * Register a service factory. Replacing a factory disposes the previously @@ -36,7 +45,7 @@ export class Container { // stale cached instance and the re-registration would be silently dead. const previous = this.services.get(token); if (previous != null && !this.inheritedServices.has(token)) { - void this.disposeService(token, previous); + this.trackDisposal(token, this.disposeService(token, previous)); } this.services.delete(token); this.inheritedServices.delete(token); @@ -71,11 +80,15 @@ export class Container { registerInstance(token: string, instance: T): void { const previous = this.services.get(token); if (previous != null && previous !== instance && !this.inheritedServices.has(token)) { - void this.disposeService(token, previous); + this.trackDisposal(token, this.disposeService(token, previous)); } this.services.set(token, instance); this.singletons.add(token); - // An explicitly registered instance is owned by this container. + // An explicitly registered instance is owned by this container. A stale + // factory left behind by a prior register() would let a subsequent + // remove() → get() race resurrect the factory-built object instead of + // failing loudly; clear it. + this.factories.delete(token); this.inheritedServices.delete(token); } @@ -136,9 +149,19 @@ export class Container { /** * Dispose all instantiated singleton services and clear registrations. * Services implementing dispose(), Symbol.asyncDispose, or Symbol.dispose will be cleaned up. + * + * Also drains any eviction disposals scheduled by prior `register` / + * `registerInstance` replacements — otherwise the container would clear its + * state (and reject subsequent lookups) while asynchronous disposers were + * still holding resources open. */ async dispose(): Promise { const errors: unknown[] = []; + try { + await this.awaitRegistrations(); + } catch (err) { + errors.push(err); + } try { for (const [token, service] of this.services) { if (service == null) continue; @@ -156,12 +179,47 @@ export class Container { this.factories.clear(); this.singletons.clear(); this.inheritedServices.clear(); + this.pendingDisposals.clear(); } if (errors.length > 0) { throw new AggregateError(errors, "One or more services failed to dispose"); } } + /** + * Resolves once the in-flight eviction disposal for `token` (if any) has + * settled. Callers that immediately re-`register` and then need to observe + * side effects of the prior disposer (e.g. a released DB connection) can + * `await awaitReplacement(token)` between the two operations. + */ + async awaitReplacement(token: string): Promise { + const pending = this.pendingDisposals.get(token); + if (pending) await pending; + } + + /** Drains every in-flight eviction disposal. */ + async awaitRegistrations(): Promise { + // Snapshot: entries `.finally()` themselves out of the map, but a race + // between iteration and settlement could otherwise skip pending entries. + const pending = Array.from(this.pendingDisposals.values()); + await Promise.allSettled(pending); + } + + /** + * Track an in-flight eviction disposal so callers (and {@link dispose}) can + * drain it. Identity-guarded: if a later replacement schedules another + * disposal for the same token, this promise's `.finally()` doesn't erase the + * newer entry. + */ + private trackDisposal(token: string, promise: Promise): void { + this.pendingDisposals.set(token, promise); + void promise.finally(() => { + if (this.pendingDisposals.get(token) === promise) { + this.pendingDisposals.delete(token); + } + }); + } + /** * Eviction-path disposer used by {@link register} and {@link registerInstance} * when a cached singleton is replaced. Errors are caught here so a buggy