From 36a29cb86749ed90692018d1832c711f0f638e67 Mon Sep 17 00:00:00 2001 From: "Jeremy D. Miller" Date: Sun, 26 Jul 2026 06:46:56 -0500 Subject: [PATCH] fix(#568): close the WaitForShardState race against an already-published state MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ShardStateTracker publication is asynchronous — PublishAsync posts to a Block and returns, and the consumer thread walks the listener list as it stood when it started. A ShardStatusWatcher that subscribes after its state was already delivered never sees it, because it only ever inspected states arriving at OnNext and never re-read the tracker's state map. If nothing else was published for that shard, the wait burned its full timeout (1 minute by default) and then claimed the shard never reached a sequence it had reached before the wait was even set up. The window widens exactly when the machine is busy, so it read as a flaky test rather than a bug. ShardStatusWatcher now re-checks tracker.CurrentStates() against its condition immediately after subscribing, so the state is either delivered to OnNext or found in the snapshot — never lost between the two. TrySetResult makes the double hit harmless. That also closes WaitForShardCondition's wider gap for free: it never consulted current state at all, so a condition already satisfied by every known state still waited on the next publication. A user condition that throws on some unrelated shard's state has always just meant "no match" (ShardStateTracker.publish swallows and logs whatever a listener throws), so the snapshot re-check preserves that instead of throwing out of WaitForShardCondition. Also tightened the watcher's own bookkeeping: TrySetException on the timeout and error paths, and dispose the timeout CancellationTokenSource once the wait resolves rather than leaving a timer per wait. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01QYfvCxHCUHo9MriHQocuoD --- .../Daemon/ShardStateTrackerTests.cs | 130 ++++++++++++++++++ .../Daemon/ShardStateTracker.cs | 8 +- .../Daemon/ShardStatusWatcher.cs | 82 +++++++---- 3 files changed, 195 insertions(+), 25 deletions(-) diff --git a/src/EventTests/Daemon/ShardStateTrackerTests.cs b/src/EventTests/Daemon/ShardStateTrackerTests.cs index 125519e4..addf77bd 100644 --- a/src/EventTests/Daemon/ShardStateTrackerTests.cs +++ b/src/EventTests/Daemon/ShardStateTrackerTests.cs @@ -1,3 +1,4 @@ +using JasperFx.Core; using JasperFx.Core.Reflection; using JasperFx.Events.Daemon; using JasperFx.Events.Projections; @@ -289,4 +290,133 @@ public async Task publish_leaves_the_database_null_when_the_tracker_has_none() observer.States.ShouldAllBe(x => x.DatabaseIdentifier == null); } + + // jasperfx#568 — "publish, then wait" is the natural shape for user code and test suites, but + // publication is asynchronous: PublishAsync posts to a Block and returns, and the consumer thread + // walks the listener list as it stood when it started. A watcher that subscribes after the state + // has already been delivered never sees it, and the wait used to burn its full timeout (1 minute by + // default) claiming a shard never reached a sequence it reached before the wait even started. + // + // Each of these publishes and then waits for the delivery to land in the tracker's snapshot BEFORE + // waiting, which is the deterministic form of the race — the state is provably already published, + // and nothing further will be published for that shard. + + /// + /// Block until the tracker's own listener has recorded a state for the shard at or past the sequence, + /// which is exactly the "state was published before the wait was set up" precondition. + /// + private async Task drainTo(string shardName, long sequence) + { + using var timeout = new CancellationTokenSource(10.Seconds()); + while (!timeout.IsCancellationRequested) + { + var current = theTracker.CurrentState(shardName); + if (current != null && current.Sequence >= sequence) return; + + await Task.Delay(10.Milliseconds()); + } + + throw new TimeoutException($"The tracker never recorded {shardName} at sequence {sequence}"); + } + + [Fact] + public async Task shard_status_watcher_completes_from_the_already_published_state() + { + // Straight at the watcher, since WaitForShardState's own pre-check would otherwise mask it: + // subscribing after publication must still complete. + await theTracker.PublishAsync(new ShardState("Trip:All", 45) { AgentStatus = "Running" }); + await drainTo("Trip:All", 45); + + var watcher = new ShardStatusWatcher(theTracker, new ShardState("Trip:All", 45), 5.Seconds()); + + var state = await watcher.Task; + state.ShardName.ShouldBe("Trip:All"); + state.Sequence.ShouldBe(45); + } + + [Fact] + public async Task wait_for_shard_state_does_not_hang_on_an_already_published_state() + { + await theTracker.PublishAsync(new ShardState("Trip:All", 45) { AgentStatus = "Running" }); + await drainTo("Trip:All", 45); + + var state = await theTracker.WaitForShardState("Trip:All", 45, 5.Seconds()); + + state.Sequence.ShouldBe(45); + } + + [Fact] + public async Task wait_for_shard_condition_does_not_hang_on_an_already_published_state() + { + // WaitForShardCondition never checked the current state at all, so a condition already satisfied + // by every known state waited for the *next* publication regardless of timing. + await theTracker.PublishAsync(new ShardState("Trip:All", 45) { AgentStatus = "Paused" }); + await drainTo("Trip:All", 45); + + var state = await theTracker.WaitForShardCondition( + x => x.ShardName == "Trip:All" && x.AgentStatus == "Paused", + "Trip:All is paused", 5.Seconds()); + + state.AgentStatus.ShouldBe("Paused"); + } + + [Fact] + public async Task wait_for_high_water_mark_does_not_hang_on_an_already_published_mark() + { + await theTracker.MarkHighWaterAsync(1500); + await drainTo(ShardState.HighWaterMark, 1500); + + var state = await theTracker.WaitForHighWaterMark(1200, 5.Seconds()); + + state.Sequence.ShouldBeGreaterThanOrEqualTo(1200); + } + + [Fact] + public async Task wait_for_shard_state_still_waits_for_a_state_that_has_not_happened_yet() + { + // The re-check must not turn every wait into "complete against whatever is there now". + await theTracker.PublishAsync(new ShardState("Trip:All", 10)); + await drainTo("Trip:All", 10); + + var waiter = theTracker.WaitForShardState("Trip:All", 45, 10.Seconds()); + waiter.IsCompleted.ShouldBeFalse(); + + await theTracker.PublishAsync(new ShardState("Trip:All", 45)); + + (await waiter).Sequence.ShouldBe(45); + } + + [Fact] + public async Task wait_for_shard_condition_still_waits_for_a_condition_not_yet_met() + { + await theTracker.PublishAsync(new ShardState("Trip:All", 10) { AgentStatus = "Running" }); + await drainTo("Trip:All", 10); + + var waiter = theTracker.WaitForShardCondition( + x => x.ShardName == "Trip:All" && x.AgentStatus == "Paused", "Trip:All is paused", 10.Seconds()); + waiter.IsCompleted.ShouldBeFalse(); + + await theTracker.PublishAsync(new ShardState("Trip:All", 11) { AgentStatus = "Paused" }); + + (await waiter).AgentStatus.ShouldBe("Paused"); + } + + [Fact] + public async Task a_condition_that_blows_up_on_an_unrelated_state_is_still_just_a_miss() + { + // The tracker has always swallowed and logged whatever a listener throws, so a condition that + // is only safe for its own shard used to mean "no match yet". Re-checking the snapshot must not + // turn that into an exception thrown straight out of WaitForShardCondition. + await theTracker.PublishAsync(new ShardState("Other:All", 10)); + await drainTo("Other:All", 10); + + var waiter = theTracker.WaitForShardCondition( + x => x.AgentStatus!.Length > 0 && x.ShardName == "Trip:All", "Trip:All reports a status", + 10.Seconds()); + waiter.IsCompleted.ShouldBeFalse(); + + await theTracker.PublishAsync(new ShardState("Trip:All", 11) { AgentStatus = "Running" }); + + (await waiter).ShardName.ShouldBe("Trip:All"); + } } diff --git a/src/JasperFx.Events/Daemon/ShardStateTracker.cs b/src/JasperFx.Events/Daemon/ShardStateTracker.cs index fb50d5b3..ba477d98 100644 --- a/src/JasperFx.Events/Daemon/ShardStateTracker.cs +++ b/src/JasperFx.Events/Daemon/ShardStateTracker.cs @@ -180,7 +180,9 @@ public IReadOnlyList CurrentStates() => _states.Enumerate().Select(x => x.Value).ToList(); /// - /// Use to "wait" for an expected projection shard state + /// Use to "wait" for an expected projection shard state. Safe to call after the state has already + /// been published — the returned task completes immediately from the current state snapshot rather + /// than waiting on the next publication. See jasperfx#568. /// /// /// @@ -233,7 +235,9 @@ public Task WaitForShardState(ShardName name, long sequence, TimeSpa } /// - /// Use to "wait" for an expected projection shard condition + /// Use to "wait" for an expected projection shard condition. The condition is evaluated against + /// every state this tracker has already seen before waiting on new publications, so a condition + /// that is already satisfied completes immediately. See jasperfx#568. /// /// /// diff --git a/src/JasperFx.Events/Daemon/ShardStatusWatcher.cs b/src/JasperFx.Events/Daemon/ShardStatusWatcher.cs index 6dd3552a..e5a9c518 100644 --- a/src/JasperFx.Events/Daemon/ShardStatusWatcher.cs +++ b/src/JasperFx.Events/Daemon/ShardStatusWatcher.cs @@ -10,41 +10,64 @@ internal class ShardStatusWatcher : IObserver { private readonly TaskCompletionSource _completion; private readonly Func _condition; + private readonly CancellationTokenSource _timeout; private readonly IDisposable _unsubscribe; public ShardStatusWatcher(ShardStateTracker tracker, ShardState expected, TimeSpan timeout) + : this(tracker, x => x.ShardName == expected.ShardName && x.Sequence >= expected.Sequence, timeout, + $"Shard {expected.ShardName} did not reach sequence number {expected.Sequence} in the time allowed") { - _condition = x => x.ShardName == expected.ShardName && x.Sequence >= expected.Sequence; - _completion = new TaskCompletionSource(); - - - var timeout1 = new CancellationTokenSource(timeout); - timeout1.Token.Register(() => - { - _completion.TrySetException(new TimeoutException( - $"Shard {expected.ShardName} did not reach sequence number {expected.Sequence} in the time allowed")); - }); - - _unsubscribe = tracker.Subscribe(this); } public ShardStatusWatcher(string description, Func condition, ShardStateTracker tracker, TimeSpan timeout) + : this(tracker, condition, timeout, $"{description} was not detected in the time allowed") { - _condition = condition; - _completion = new TaskCompletionSource(); + Debug.WriteLine("Subscribed to watch shard state: " + description); + } + private ShardStatusWatcher(ShardStateTracker tracker, Func condition, TimeSpan timeout, + string timeoutMessage) + { + _condition = condition; + _completion = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); - var timeout1 = new CancellationTokenSource(timeout); - timeout1.Token.Register(() => + _timeout = new CancellationTokenSource(timeout); + _timeout.Token.Register(() => { - _completion.TrySetException(new TimeoutException( - $"{description} was not detected in the time allowed")); + if (_completion.TrySetException(new TimeoutException(timeoutMessage))) + { + // ReSharper disable once ConstantConditionalAccessQualifier -- the token can fire before + // the constructor has finished assigning _unsubscribe if the timeout is tiny. + _unsubscribe?.Dispose(); + } }); _unsubscribe = tracker.Subscribe(this); - Debug.WriteLine("Subscribed to watch shard state: " + description); + // jasperfx#568: publication is asynchronous — PublishAsync posts to a Block and returns, and the + // consumer thread walks the listener list that existed when it started. A state can therefore be + // delivered (and recorded in the tracker's state map) in the window between a caller checking the + // map and this watcher subscribing, and nothing else may ever be published for that shard. Since + // OnNext only ever sees states delivered AFTER this point, re-check what the tracker already knows + // now that we're subscribed. TrySetResult makes a double hit harmless — whichever path wins first, + // the other is a no-op. This is also the only current-state check WaitForShardCondition gets. + foreach (var state in tracker.CurrentStates()) + { + try + { + if (check(state)) + { + break; + } + } + catch (Exception) + { + // A user supplied condition that blows up on some unrelated shard's state has always + // just meant "no match" -- ShardStateTracker.publish() swallows and logs whatever a + // listener throws. Keep that contract here instead of throwing out of WaitFor*. + } + } } public Task Task => _completion.Task; @@ -55,15 +78,28 @@ public void OnCompleted() public void OnError(Exception error) { - _completion.SetException(error); + if (_completion.TrySetException(error)) + { + _unsubscribe.Dispose(); + _timeout.Dispose(); + } } public void OnNext(ShardState value) { - if (_condition(value)) + check(value); + } + + private bool check(ShardState state) + { + if (!_condition(state)) return false; + + if (_completion.TrySetResult(state)) { - _completion.SetResult(value); _unsubscribe.Dispose(); + _timeout.Dispose(); } + + return true; } -} \ No newline at end of file +}