fix(#568): close the WaitForShardState race against an already-published state

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.
+
+    /// <summary>
+    /// 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.
+    /// </summary>
+    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");
+    }
 }

src/JasperFx.Events/Daemon/ShardStateTracker.cs

@@ -180,7 +180,9 @@ public IReadOnlyList<ShardState> CurrentStates()
         => _states.Enumerate().Select(x => x.Value).ToList();
 
     /// <summary>
-    ///     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.
     /// </summary>
     /// <param name="expected"></param>
     /// <param name="timeout"></param>
@@ -233,7 +235,9 @@ public Task<ShardState> WaitForShardState(ShardName name, long sequence, TimeSpa
     }
 
     /// <summary>
-    ///     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.
     /// </summary>
     /// <param name="condition"></param>
     /// <param name="description"></param>

src/JasperFx.Events/Daemon/ShardStatusWatcher.cs

@@ -10,41 +10,64 @@ internal class ShardStatusWatcher : IObserver<ShardState>
 {
     private readonly TaskCompletionSource<ShardState> _completion;
     private readonly Func<ShardState, bool> _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<ShardState>();
-
-
-        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<ShardState, bool> condition, ShardStateTracker tracker,
         TimeSpan timeout)
+        : this(tracker, condition, timeout, $"{description} was not detected in the time allowed")
     {
-        _condition = condition;
-        _completion = new TaskCompletionSource<ShardState>();
+        Debug.WriteLine("Subscribed to watch shard state: " + description);
+    }
 
+    private ShardStatusWatcher(ShardStateTracker tracker, Func<ShardState, bool> condition, TimeSpan timeout,
+        string timeoutMessage)
+    {
+        _condition = condition;
+        _completion = new TaskCompletionSource<ShardState>(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<ShardState> 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;
     }
-}
+}