fix: acquire step batch configs on the caller's transaction - #4617
Merged
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Contributor
Benchmark resultsCompared against |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR mitigates a potential deadlock/pool starvation scenario by ensuring step batch config lookups participate in the caller’s existing transaction (when available), and adds a shared LRU cache to reduce repeated DB lookups for step batch config presence during scheduling.
Changes:
- Updated
GetStepBatchConfigsto accept an optional*OptimisticTxand to query via the provided transaction when present. - Added an expirable LRU cache (
stepIdBatchConfigCache) for step batch-config presence, including caching negative results. - Updated scheduling call sites and the test fake repository to match the new repository interface.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| pkg/scheduling/v1/queuer.go | Passes the active optimistic transaction (or nil) into GetStepBatchConfigs to ensure DB access happens on the caller’s connection. |
| pkg/scheduling/v1/batch_scheduler_test.go | Updates the fake repository method signature to match the new interface contract. |
| pkg/repository/shared.go | Adds a shared expirable LRU cache for step batch-config presence and initializes it in newSharedRepository. |
| pkg/repository/scheduler.go | Updates the QueueRepository interface to accept an optional transaction for GetStepBatchConfigs. |
| pkg/repository/scheduler_queue.go | Implements transaction-aware querying for step batch configs and introduces cache-backed lookups (including negative caching). |
juliusgeo
reviewed
Aug 4, 2026
juliusgeo
reviewed
Aug 4, 2026
GetStepBatchConfigs always queried on the pool, so the optimistic assign path — which holds a transaction at that point, as its sibling lookups show — acquired a second pool connection while holding one. Under pool saturation, concurrent schedulers each holding a connection and waiting for another deadlock until their timeouts unwind them. The method now takes the caller's OptimisticTx with the same nil-means-pool convention as GetDesiredLabels and GetStepSlotRequests. Batch configs are also now cached per step id with negative caching, mirroring the slot-requests cache and relying on the same invariant: step configuration is immutable per step id, since config changes create a new workflow version and therefore new step ids. Steady-state assignment batches skip the lookup entirely.
abelanger5
force-pushed
the
belanger/step-batch-deadlock
branch
from
August 4, 2026 20:08
5939033 to
0f99b7c
Compare
juliusgeo
approved these changes
Aug 4, 2026
…eadlock # Conflicts: # pkg/repository/scheduler.go
abelanger5
commented
Aug 4, 2026
| var expectedAllKeys int64 | ||
| for _, k := range config.EventKeys { | ||
| expectedAllKeys += int64(pushedByKey[k]) * executionsPerPush(k, config.EventFanout, config.DagSteps) | ||
| expectedAllKeys += pushedByKey[k] * executionsPerPush(k, config.EventFanout, config.DagSteps) |
Contributor
Author
There was a problem hiding this comment.
drive-by linting issue
abelanger5
added a commit
that referenced
this pull request
Aug 5, 2026
… merge The merge of main brought #4617 (GetStepBatchConfigs now runs on the caller's transaction) into pkg/scheduling/v1, but git knows nothing about the v1alpha copies, so the merge completed without conflicts and left v1alpha calling the old repository signature — every consumer of the package failed to compile. Mirror the merged queuer.go and batch_scheduler_test.go into v1alpha; this is exactly the drift the scheduler-sync workflow flags, and it does on the broken merge commit. Also fix a pre-existing flaky test the reshuffled suite exposed: TestLeaseManager_SendWorkerIds/SendQueues raced a goroutine calling a non-blocking send (select/default) against the test's own receive on an unbuffered channel — if the send ran first the message was dropped and the receive blocked until the 10m suite timeout. Buffer the channel and send synchronously; mirrored across both packages.
grutt
added a commit
that referenced
this pull request
Aug 5, 2026
* test: benchmarks * perf: pool based scheduler * fix race * cleanup bench * fix: worker race * perf(scheduler): replace lock-based scheduler with per-tenant event loop All scheduling state is now owned by a single run-loop goroutine; every read and write goes through an op sent to the loop. This removes actionsMu, unackedMu, workersMu, assignedCountMu, replenishMu, the per-slot RWMutex, and the atomic unused counter, along with the lock ordering rules between them. - Replenish DB reads run outside the loop, so assignment is never blocked on I/O. The rebuild reconciles reads against assignments that acked mid-read via a per-pool ackedDuringReplenish counter instead of holding a write lock across three DB round trips. - Pools keep an exact freelist of slot indexes, making selection O(requested units) instead of an O(pool) scan with per-slot locking, and eliminating unused-counter drift by construction. Expiry is pool-level: every slot in a pool comes from the same replenish read. - Rate-limit checks and callbacks run outside the loop so it never blocks on another subsystem's locks. - The assignment ring offset lives on the action, so round-robin placement persists across batches. - Drop dead code: slotMeta, getRankedSlots, rankedValidSlots, slot ack bookkeeping. Adds a concurrent assign+replenish+snapshot benchmark. vs the previous branch head (n=6): concurrent bench -64% time / -89% allocs, TryAssignBatch -55% to -98.8%, replenish geomean roughly -30%, package geomean -46%. * feat(scheduler): gate v1alpha event-loop scheduler behind a per-shard flag Address PR feedback and make the new scheduler opt-in: - The actor/event-loop scheduler moves to pkg/scheduling/v1alpha; pkg/scheduling/v1 reverts to the existing lock-based scheduler. - A single flag selects the implementation for the whole shard: SERVER_V1ALPHA_SCHEDULER_ENABLED (runtime.v1AlphaSchedulerEnabled, default false). The loader constructs one pool or the other; nothing else in the engine changes behavior. - pkg/scheduling hosts the shared engine-facing surface: the Pool interface, the QueueResults/ConcurrencyResults/AssignedItemWithTask DTOs (aliased from both packages so channel and result types are identical either way), and the optimistic-scheduling sentinel errors (shared so errors.Is matches across implementations). Both pools carry compile-time interface assertions. Review feedback on the v1alpha scheduler itself: - document that slotPool is confined to the run loop and not concurrency-safe on its own - snapshot ticker reduced from 10-90ms to 1-1.5s now that every tick reports (previously most ticks were skipped on lock contention) - replace context.Background() do() calls with an explicit mustDo: acks, nacks and worker updates must not be cancellable once their trigger has committed, otherwise loop state diverges from the database - snapshot/utilization reporting split out of scheduler.go into snapshot.go * fix(scheduler): rename v1alpha integration test package and imports The integration tests use an external test package (package v1_test) and are hidden behind the integration build tag, so the package rename sweep missed them and untagged builds couldn't catch it: the v1alpha directory contained two packages and every consumer of the config loader failed to compile in CI. Rename to v1alpha_test and point the scheduling import at v1alpha so they exercise the new pool. * ci: add the v1alpha scheduler to the load testing matrix The load and load-deadlock jobs now run the full cross product of optimistic-scheduling and v1alpha-scheduler. The testing harness maps TESTING_MATRIX_V1ALPHA_SCHEDULER onto SERVER_V1ALPHA_SCHEDULER_ENABLED, mirroring the optimistic scheduling flag. * ci: enforce sync between the v1 and v1alpha scheduler packages While the v1alpha scheduler is rolling out, the two packages must stay in lockstep. A new scheduler-sync workflow enforces two rules via hack/ci/check-scheduler-sync.sh: 1. Files shared between the packages must be byte-for-byte identical after normalizing the package clause and the scheduling/v1[alpha] self-import path. The rewritten scheduler core (event loop, pool-owned slots) and its tests are the explicit divergent list. 2. On pull requests, a change to a divergent core file in one package must be accompanied by a change to the other package, so scheduler fixes get ported (or explicitly recorded as not applicable) instead of silently drifting. Also mirrors the v1alpha #nosec G115 annotations in batch_scheduler.go back into v1 so the file stays in the shared set. * fix(scheduler): address PR review on the rate limiter and docs - shouldRefill had the comparison inverted: it returned true while the refill deadline was still in the future — exactly when flushToDatabase's early-return guard makes the flush a no-op — and false once the deadline passed, so use() served stale limits for up to a second until the background flush tick. Refill when the deadline has been reached; the guard in flushToDatabase plus MAX_RATE_LIMIT_UPDATE_FREQUENCY still bound database traffic. Adds a regression test. - stop the loopFlush ticker on shutdown - fix a rate limiter test assertion that checked the wrong key type against unflushed and could never fail - drop SERVER_V1ALPHA_SCHEDULER_ENABLED from the self-hosting docs: it is an internal rollout control, not a supported option All changes mirrored across scheduling/v1 and scheduling/v1alpha. * fix(ci): keep the go-benchmarks job within its budget The go-benchmarks job was timing out on this PR and would have failed on the base-branch run right after: - TryAssignBatch replenished off-timer on every iteration; when the op got ~80x faster, b.N grew to thousands of iterations and the off-timer replenish dominated wall time. It also never acked, so at high b.N the fixture ran out of capacity and later iterations measured the noSlots path instead of assignment. The bench now acks every batch and replenishes only when the action's reachable capacity runs low. - the concurrent bench now acks assignments like a queuer flush, so unacked slots don't accumulate across iterations and skew the background replenish churn - the workflow skips packages that don't exist on the base branch (pkg/scheduling/v1alpha is new) and gets a larger timeout to fit two scheduler packages' suites at count=6 * fix(scheduler): park assignments that race an in-flight replenish On task completion the engine notifies the queuers and then triggers an on-demand replenish. In the v1 scheduler the woken queuer blocked on the actions write lock behind that replenish and always woke to fresh capacity — an accidental rendezvous. The v1alpha run loop never blocks, so the assignment raced ahead, saw the not-yet-applied pools, reported noSlots, and paid a full one-second queue poll before retrying. On a single-slot worker this roughly halved drain throughput. Assignment batches that miss while a replenish cycle is in flight now park on the run loop and retry once when the cycle ends, bounded by a 100ms timeout so a stalled replenish (slow database reads) cannot hold assignment results hostage the way the v1 lock could. In the single-slot embedded e2e this takes the drain rate from 1.8 to 24 tasks/s (v1: 3.3). Also reset ackedDuringReplenish just before the availability read: acks landing earlier in the cycle are already visible to the read, so counting them only inflated the conservative double-subtract window. Measured collision rate in the e2e was zero; the counter is a narrow guard, not a hot path. * fix(scheduler): mirror main's batch-config changes into v1alpha after merge The merge of main brought #4617 (GetStepBatchConfigs now runs on the caller's transaction) into pkg/scheduling/v1, but git knows nothing about the v1alpha copies, so the merge completed without conflicts and left v1alpha calling the old repository signature — every consumer of the package failed to compile. Mirror the merged queuer.go and batch_scheduler_test.go into v1alpha; this is exactly the drift the scheduler-sync workflow flags, and it does on the broken merge commit. Also fix a pre-existing flaky test the reshuffled suite exposed: TestLeaseManager_SendWorkerIds/SendQueues raced a goroutine calling a non-blocking send (select/default) against the test's own receive on an unbuffered channel — if the send ran first the message was dropped and the receive blocked until the 10m suite timeout. Buffer the channel and send synchronously; mirrored across both packages. --------- Co-authored-by: Gabe Ruttner <gabriel.ruttner@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Mitigates a deadlock risk I noticed while staring at some traces, where
GetStepBatchConfigsrequired a new connection from the pool within an existing tx, which immediately raises alarm bells. Also adds a cache to mirror behavior elsewhere.Fixes # (issue)
Type of change
What's Changed
Checklist
Changes have been:
🤖 AI Disclosure