Add a scheduler that ensures single batch of changes from live query due to a transaction that touches multiple source collections

[samwillis]

stacked on #625

Live Query Scheduler Overview

Overview

This change introduces a scoped, dependency-aware scheduler that guarantees every live-query builder runs at most once per transaction, even when multiple source collections or derived queries fire in the same mutate call. The scheduler groups work by transaction id, dedupes entries by the CollectionConfigBuilder instance, tracks dependencies between builders, and flushes synchronously when Transaction.mutate exits. The net effect: optimistic updates from a single transaction coalesce into a single graph run for each live query, so downstream frameworks see one change batch per transaction.

How scheduling works

• Each CollectionSubscriber immediately forwards changes to the D2 input and calls CollectionConfigBuilder.scheduleGraphRun.
• scheduleGraphRun records which upstream builders the current builder depends on (based on the collections it subscribed to) and hands a job to the shared scheduler with:
  • contextId: the transaction id, grouping all work triggered by that transaction.
  • jobId: the builder instance; repeated schedules for the same builder merge into a single entry.
  • dependencies: the set of upstream builders that must finish before this builder can run.
• The scheduler maintains, per transaction, a queue, entry map, dependency map, and “completed” set. When flushing, it only executes a job once all dependencies are marked completed, deferring the job otherwise. It loops until no work remains and throws if it detects an impossible cycle.

When the graph actually runs

• Inside Transaction.mutate, we call registerTransaction before the user’s callback and unregisterTransaction in a finally block afterward.
• registerTransaction clears any stale entries from previous failed scopes.
• unregisterTransaction calls scheduler.flush(tx.id). flush now loops while a context map exists, so if a job enqueues additional work during its own execution (e.g., the join scheduling itself after its parents run), the scheduler immediately picks up the new entry before leaving the transaction scope.
• If there is no active transaction, scheduleGraphRun runs maybeRunGraph immediately (backward compatible behaviour).

Nested live queries (the “diamond” cases)

• Builders register themselves when a live-query collection is created. Whenever a builder subscribes to another live query, it records that dependency.
• During a transaction, updates enqueue jobs for the parent builders. When those builders finish, they mark themselves as completed, which in turn allows child jobs (such as joins) to run exactly once at the end of the flush.
• A hybrid variant (join between liveQueryA and raw collectionB) behaves the same way: even if collectionB fires first, the join job is deferred until liveQueryA has finished.
• Because dependencies are tracked explicitly, there are no deadlocks. If the scheduler ever loops without making progress, it throws a clear “dependency cycle” error.

Order-by loaders and batching

• The maybeRunGraph loop keeps calling graph.run() while pendingWork() is true. If an order-by loader requests more rows, it pushes those changes into the same D2 input, triggering pendingWork() again.
• During this loop CollectionConfigBuilder.isGraphRunning is true, so any nested call to scheduleGraphRun is ignored; the loader’s extra rows are consumed by the ongoing loop.
• As a result, even hungry top-K queries emit a single batch per synchronous transaction run. Multiple batches only occur when a loader deliberately yields results asynchronously (for example, the incremental sync features we’re adding)—in that case the UI should expect staged updates.

Tests

scheduler.test.ts now covers:

1. Basic single run per transaction (two collections mutated inside one mutate).
2. Nested transactions.
3. Rollback cleanup.
4. Multiple subscribers receiving exactly one batch.
5. Loader deduping.
6. Diamond dependency (live query joining two parents) – verifies parents run first and the join runs once (tracked via getRunCount).
7. Hybrid diamond (join of a live query with a base collection) – same single-run guarantee, even when the raw collection fires first.

Each of the diamond tests mutates once and then performs another transaction with updates, demonstrating that we still emit exactly one additional batch—no duplicates, no missed runs—and that getRunCount() only increments once per transaction.

Summary

• Scheduler coalesces work by transaction + builder.
• Transaction.mutate clears before and flushes after every scope, so everything runs synchronously in the optimistic phase.
• flush loops until no jobs remain, ensuring child live queries (joins) run exactly once per transaction after their parents.
• Comprehensive tests cover simple, nested, rollback, multi-subscriber, loader, diamond, and hybrid patterns; asynchronous loaders will still produce multiple batches by design.

[changeset-bot]

🦋 Changeset detected

Latest commit: c46a15a

The changes in this PR will be included in the next version bump.

This PR includes changesets to release 12 packages

Name	Type
@tanstack/db	Patch
@tanstack/angular-db	Patch
@tanstack/electric-db-collection	Patch
@tanstack/query-db-collection	Patch
@tanstack/react-db	Patch
@tanstack/rxdb-db-collection	Patch
@tanstack/solid-db	Patch
@tanstack/svelte-db	Patch
@tanstack/trailbase-db-collection	Patch
@tanstack/vue-db	Patch
todos	Patch
@tanstack/db-example-react-todo	Patch

Not sure what this means? Click here to learn what changesets are.

Click here if you're a maintainer who wants to add another changeset to this PR

[pkg-pr-new]

More templates

• todos
• tanstack-start-db-projects-example
• @tanstack/db-example-react-todo
• @tanstack/db-example-solid-todo

@tanstack/angular-db

npm i https://pkg.pr.new/@tanstack/angular-db@628

@tanstack/db

npm i https://pkg.pr.new/@tanstack/db@628

@tanstack/db-ivm

npm i https://pkg.pr.new/@tanstack/db-ivm@628

@tanstack/electric-db-collection

npm i https://pkg.pr.new/@tanstack/electric-db-collection@628

@tanstack/query-db-collection

npm i https://pkg.pr.new/@tanstack/query-db-collection@628

@tanstack/react-db

npm i https://pkg.pr.new/@tanstack/react-db@628

@tanstack/rxdb-db-collection

npm i https://pkg.pr.new/@tanstack/rxdb-db-collection@628

@tanstack/solid-db

npm i https://pkg.pr.new/@tanstack/solid-db@628

@tanstack/svelte-db

npm i https://pkg.pr.new/@tanstack/svelte-db@628

@tanstack/trailbase-db-collection

npm i https://pkg.pr.new/@tanstack/trailbase-db-collection@628

@tanstack/vue-db

npm i https://pkg.pr.new/@tanstack/vue-db@628

commit: 4d5ff08

[github-actions]

Size Change: +2.97 kB (+3.95%)

Total Size: 78.2 kB

Filename	Size	Change
./packages/db/dist/esm/query/live-query-collection.js	416 B	+76 B (+22.35%)	🚨
./packages/db/dist/esm/query/live/collection-config-builder.js	4.43 kB	+1.18 kB (+36.44%)	🚨
./packages/db/dist/esm/query/live/collection-subscriber.js	1.83 kB	+77 B (+4.39%)
./packages/db/dist/esm/transactions.js	3.08 kB	+44 B (+1.45%)
./packages/db/dist/esm/query/live/collection-registry.js	349 B	+349 B (new file)	🆕
./packages/db/dist/esm/scheduler.js	1.24 kB	+1.24 kB (new file)	🆕

ℹ️ View Unchanged

Filename	Size
./packages/db/dist/esm/collection/change-events.js	958 B
./packages/db/dist/esm/collection/changes.js	1.01 kB
./packages/db/dist/esm/collection/events.js	683 B
./packages/db/dist/esm/collection/index.js	3.14 kB
./packages/db/dist/esm/collection/indexes.js	1.16 kB
./packages/db/dist/esm/collection/lifecycle.js	1.8 kB
./packages/db/dist/esm/collection/mutations.js	2.59 kB
./packages/db/dist/esm/collection/state.js	3.81 kB
./packages/db/dist/esm/collection/subscription.js	1.69 kB
./packages/db/dist/esm/collection/sync.js	1.32 kB
./packages/db/dist/esm/deferred.js	230 B
./packages/db/dist/esm/errors.js	3.46 kB
./packages/db/dist/esm/index.js	1.6 kB
./packages/db/dist/esm/indexes/auto-index.js	745 B
./packages/db/dist/esm/indexes/base-index.js	605 B
./packages/db/dist/esm/indexes/btree-index.js	1.82 kB
./packages/db/dist/esm/indexes/lazy-index.js	1.25 kB
./packages/db/dist/esm/local-only.js	827 B
./packages/db/dist/esm/local-storage.js	2.02 kB
./packages/db/dist/esm/optimistic-action.js	294 B
./packages/db/dist/esm/proxy.js	3.87 kB
./packages/db/dist/esm/query/builder/functions.js	615 B
./packages/db/dist/esm/query/builder/index.js	3.93 kB
./packages/db/dist/esm/query/builder/ref-proxy.js	938 B
./packages/db/dist/esm/query/compiler/evaluators.js	1.56 kB
./packages/db/dist/esm/query/compiler/expressions.js	631 B
./packages/db/dist/esm/query/compiler/group-by.js	2.11 kB
./packages/db/dist/esm/query/compiler/index.js	2.19 kB
./packages/db/dist/esm/query/compiler/joins.js	2.67 kB
./packages/db/dist/esm/query/compiler/order-by.js	1.27 kB
./packages/db/dist/esm/query/compiler/select.js	1.28 kB
./packages/db/dist/esm/query/ir.js	785 B
./packages/db/dist/esm/query/optimizer.js	3.1 kB
./packages/db/dist/esm/SortedMap.js	1.24 kB
./packages/db/dist/esm/utils.js	943 B
./packages/db/dist/esm/utils/browser-polyfills.js	365 B
./packages/db/dist/esm/utils/btree.js	6.02 kB
./packages/db/dist/esm/utils/comparison.js	754 B
./packages/db/dist/esm/utils/index-optimization.js	1.62 kB

_{compressed-size-action::db-package-size}

[github-actions]

Size Change: 0 B

Total Size: 1.44 kB

ℹ️ View Unchanged

Filename	Size
./packages/react-db/dist/esm/index.js	152 B
./packages/react-db/dist/esm/useLiveQuery.js	1.29 kB

_{compressed-size-action::react-db-package-size}

[kevin-dp]

I had a close look at this PR and left some comments. My main concern is the time complexity of the scheduler's flush method.

[samwillis]

@kevin-dp I've largely refactored this PR based on your feedback, removing the closures and hopefully simplifying it. Summary for cursor:

Overview

This refactoring addressed PR feedback about excessive closure usage in the live query scheduler and builder, while also fixing a type error, preventing memory leaks, and improving overall code clarity, testability, and performance.

Problems Addressed

1. Type Error

• Problem: liveQueryJoin.utils.getRunCount is not a function when running tests
• Root Cause: createLiveQueryCollection was overwriting LiveQueryCollectionUtils when custom utils were provided via config.utils
• Fix: Changed to merge utils using spread syntax: options.utils = { ...options.utils, ...config.utils }

2. Excessive Closure Usage

• Problem: PR review feedback indicated too many closures, particularly in the scheduler and how it interacted with the builder
• Previous Architecture:
  • Scheduler managed state via createEntry/updateEntry callbacks that closed over config, syncState, and callbacks
  • Complex indirection: builder created closures → passed to scheduler → scheduler called closures to manage state
  • State was scattered between scheduler's generic state management and builder's specific needs

3. Memory Leaks

• Problem 1: pendingGraphRuns map entries could leak when sync ended before transactions flushed
• Problem 2: Scheduler's onClear listener held strong reference to builder, preventing garbage collection
• Impact: Stale callbacks and builder instances accumulated in long-lived apps that create/destroy live queries

Solution: Separation of Concerns

New Architecture Principles

1. Scheduler only orchestrates execution order

   • Simplified to just managing job dependencies and execution sequencing
   • No longer manages caller-specific state
   • Jobs provide a simple run: () => void function
   • Deduplication happens by replacing the run function when the same job is scheduled multiple times

2. Builder manages its own state

   • Introduced currentSyncConfig and currentSyncState as public instance properties
   • Manages pendingGraphRuns map internally (per scheduler context/transaction)
   • Clear lifecycle: properties set when sync starts, cleared when sync stops
   • Made public for testability (builder is internal, not public API)

3. Lifecycle Coordination

   • Builder registers a clearPendingGraphRun listener via scheduler.onClear() when sync starts
   • When scheduler clears a context (e.g., transaction rollback), builder cleans up its pending state
   • Listener is unregistered when sync stops to prevent memory leaks
   • Graceful degradation: if sync ends while transaction is in-flight, execution returns early without error

Key Changes

Scheduler (packages/db/src/scheduler.ts)

Before:

interface ScheduleOptions<TState> {
  contextId?: SchedulerContextId
  jobId: unknown
  dependencies?: Iterable<unknown>
  createEntry: () => SchedulerEntry<TState>
  updateEntry: (entry: SchedulerEntry<TState>) => void
  run: (state: TState) => void
}

After:

interface ScheduleOptions {
  contextId?: SchedulerContextId
  jobId: unknown
  dependencies?: Iterable<unknown>
  run: () => void
}

Changes:

• Removed createEntry, updateEntry, and generic TState parameter
• Simplified state management to just storing run functions
• Added onClear() listener mechanism for cleanup coordination
• State is now: { jobs: Map<unknown, () => void>, ... } instead of { entries: Map<unknown, SchedulerEntry<TState>>, ... }

Performance Optimizations:

• Dependency set creation: new Set(dependencies); depSet.delete(jobId) instead of manual filtering
• Changed from [...deps].every() to for...of loop to avoid intermediate array allocation
• Simplified hasPendingJobs to single return statement: return !!context && context.jobs.size > 0

CollectionConfigBuilder (packages/db/src/query/live/collection-config-builder.ts)

Key Additions:

export class CollectionConfigBuilder<TContext, TResult> {
  // Current sync session state - set when sync starts, cleared when stops
  // Public for testing (builder is internal, not public API)
  public currentSyncConfig: Parameters<SyncConfig<TResult>['sync']>[0] | undefined
  public currentSyncState: FullSyncState | undefined

  // Builder manages its own pending state per scheduler context
  private readonly pendingGraphRuns = new Map<SchedulerContextId, PendingGraphRun>()
  
  // Scheduler listener lifecycle management (prevents memory leaks)
  private unsubscribeFromSchedulerClears?: () => void
}

Flow Changes:

1. scheduleGraphRun():

   • Uses instance properties currentSyncConfig and currentSyncState
   • Manages pendingGraphRuns state directly (get or create, accumulate callbacks)
   • Passes simple run function to scheduler: run: () => this.executeGraphRun(contextId, pendingToPass)
   • No closures over config/syncState passed to scheduler

2. executeGraphRun():

   • Removes pending state from map FIRST (before checking sync state) to prevent memory leaks
   • Returns early without error if sync has ended (graceful degradation)
   • Re-schedules during execution create fresh state (map entry already removed)
   • Increments run count and executes graph with combined loader callbacks

3. clearPendingGraphRun():

   • New public method called by scheduler's onClear listener
   • Ensures pendingGraphRuns map stays in sync when scheduler clears a context
   • Prevents stale callbacks from surviving transaction rollbacks

4. maybeRunGraph():

   • Signature simplified from maybeRunGraph(config, syncState, callback) to maybeRunGraph(callback)
   • Uses instance properties currentSyncConfig and currentSyncState directly

5. syncFn() lifecycle:

   • Start:
     • Sets currentSyncConfig and currentSyncState instance properties
     • Registers onClear listener with scheduler
   • Stop:
     • Unregisters onClear listener (prevents builder memory leak)
     • Clears currentSyncConfig and currentSyncState
     • Clears all pendingGraphRuns (prevents transaction entry leaks)
     • Resets caches, clears lazy source state, unsubscribes and clears subscriptions

CollectionSubscriber (packages/db/src/query/live/collection-subscriber.ts)

Before:

constructor(
  private alias: string,
  private collectionId: string,
  private collection: Collection,
  private config: Parameters<SyncConfig<TResult>['sync']>[0],
  private syncState: FullSyncState,
  private collectionConfigBuilder: CollectionConfigBuilder<TContext, TResult>
)

After:

constructor(
  private alias: string,
  private collectionId: string,
  private collection: Collection,
  private collectionConfigBuilder: CollectionConfigBuilder<TContext, TResult>
)

Changes:

• Removed config (was never used) and syncState parameters
• Accesses sync state via collectionConfigBuilder.currentSyncState instead
• Consistent with the new architecture: state lives on the builder

Code Quality Improvements:

• Renamed "bound loader" to "loadMoreCallback" for clarity
• Simplified callback caching using ??= operator instead of manual ?? pattern
• Improved code comments to explain callback purpose rather than implementation details

Benefits

1. Eliminated Closure Anti-patterns

• Scheduler no longer accepts or manages closures over state
• Builder state management is explicit through instance properties
• Clearer data flow: state lives where it's used, not captured in closures

2. Memory Leak Prevention

• Fixed leak where pendingGraphRuns entries weren't cleaned up when sync ended before transaction flush
• Fixed leak where scheduler's listener Set held strong reference to builder after teardown
• Defense-in-depth cleanup: per-transaction cleanup, full map clear on sync stop, listener unregistration
• Prevents accumulation of stale callbacks and builder instances across subscription cycles

3. Performance Optimizations

• Eliminated intermediate array allocations in scheduler's dependency checking
• More efficient dependency set construction using Set constructor
• Simplified conditional logic throughout (fewer branches, clearer intent)

4. Improved Testability

• Instance properties currentSyncConfig, currentSyncState, scheduleGraphRun, and clearPendingGraphRun are public
• Tests can set state directly without any casts
• Internal class boundaries respected (builder is not public API, so public members are fine)

5. Better Separation of Concerns

• Scheduler: Pure orchestration of execution order and dependencies
• Builder: Owns all live query state and lifecycle
• Clear interfaces and responsibilities

6. Lifecycle Safety

• Explicit state lifecycle tied to sync session (set on start, cleared on stop)
• Graceful handling of edge cases (sync ends during transaction, transaction rollback)
• Listener registration/unregistration follows sync session lifecycle

7. Simpler Mental Model

• State is stored where it's used (instance properties) rather than captured in closures
• Execution flow is more linear and easier to trace
• Less indirection between scheduler and builder

Testing

• All existing tests pass (1095 passed, 3 skipped)
• Updated scheduler.test.ts to work with new public properties
• Mock signatures updated to reflect simplified maybeRunGraph signature
• No behavioral changes to live query functionality

Migration Notes

This is an internal refactoring with no public API changes:

• Scheduler class is internal
• CollectionConfigBuilder is internal
• CollectionSubscriber is internal
• createLiveQueryCollection public API unchanged
• Live query behavior unchanged from user perspective
• Existing code using the public API requires no modifications

Summary

The refactoring successfully addresses PR feedback by eliminating problematic closure patterns while improving code clarity, testability, and lifecycle management. The new architecture draws clear boundaries: the scheduler orchestrates execution order, while the builder manages its own state through the lifecycle of a sync session.

Key Achievements:

• ✅ Eliminated closure anti-patterns (primary goal)
• ✅ Fixed multiple memory leaks in transaction/sync lifecycle
• ✅ Improved performance through allocation reduction
• ✅ Enhanced testability with explicit state management
• ✅ Better code clarity with improved naming and simplified logic
• ✅ No breaking changes to public API
• ✅ All tests passing

[kevin-dp]

Great, i like the simplified implementation!