fix: eliminate rebalancer waiter timeout race

[ReubenBond]

The rebalancer diagnostic regression test could time out under macOS CI load even though the awaited event was emitted synchronously. The observer completed an inner task, but returned an async WaitAsync wrapper whose continuation could lose to the timeout timer when the thread pool was delayed.

Complete event waiters directly while holding the observer's waiter lock, and serialize timeout removal through that same lock. This makes event delivery and timeout completion atomic without increasing timeout values. The regression coverage now verifies that prior events do not satisfy new waiters, new emissions complete waiters synchronously, and a timed-out waiter does not interfere with subsequent events.

Microsoft Reviewers: Open in CodeFlow

[Copilot]

Pull request overview

This PR updates the RebalancerDiagnosticObserver test helper to eliminate a timeout race when waiting for activation rebalancer diagnostic events under CI load, and strengthens regression coverage to validate the corrected semantics.

Changes:

• Replace Task.WaitAsync(timeout)-based waiting with explicit per-waiter timers and complete waiters atomically while holding the waiter lock.
• Ensure timeout removal and waiter completion are serialized via the same lock and that completed waiters stop their timers.
• Expand regression tests to assert that prior events do not satisfy new waiters, that new emissions complete waiters immediately, and that a timed-out waiter does not impact subsequent waits.

Show a summary per file

File	Description
test/TestInfrastructure/TestExtensions/Diagnostics/RebalancerDiagnosticObserver.cs	Reworks waiter lifecycle/timeout handling to avoid continuation-vs-timeout races.
test/Orleans.Core.Tests/Diagnostics/DiagnosticInfrastructureRegressionTests.cs	Updates and adds tests to validate the fixed waiter/timeout behavior and synchronous completion.

Copilot's findings

Suppressed comments (2)

test/Orleans.Core.Tests/Diagnostics/DiagnosticInfrastructureRegressionTests.cs:98

• This test now relies on the helper's 60s default timeout. Keeping an explicit (but generous) timeout makes failures surface faster and avoids long hangs if the event is never observed.

        var waitTask = observer.WaitForSessionStopAsync();

test/Orleans.Core.Tests/Diagnostics/DiagnosticInfrastructureRegressionTests.cs:117

• This waiter uses the helper's 60s default timeout. Consider using an explicit timeout here too so the test doesn't hang for up to a minute if the later event is not delivered.

        var waitTask = observer.WaitForSessionStopAsync();

• Files reviewed: 2/2 changed files
• Comments generated: 2

[Copilot]

Copilot's findings

Suppressed comments (4)

test/TestInfrastructure/TestExtensions/Diagnostics/RebalancerDiagnosticObserver.cs:254

• WaitForEventAsync always creates a timer-based timeout even when timeout is TimeSpan.Zero/negative, which means the returned task only faults once a ThreadPool timer callback runs. For an immediate timeout (and to keep tests deterministic under threadpool starvation), return a faulted task directly when timeout <= TimeSpan.Zero.

        lock (_waitersLock)
        {
            var waiter = new EventWaiter<TEvent>(predicate);
            _waiters.Add(waiter);
            waiter.StartTimeout(timeout, () => TimeoutWaiter(waiter, timeoutMessage));

test/Orleans.Core.Tests/Diagnostics/DiagnosticInfrastructureRegressionTests.cs:79

• This test now uses the default 30s timeout (WaitForCycleAsync()), which can make failures hang for a long time. Keeping an explicit short timeout preserves fast failure while still validating synchronous completion via IsCompletedSuccessfully.

        var waitTask = observer.WaitForCycleAsync();

test/Orleans.Core.Tests/Diagnostics/DiagnosticInfrastructureRegressionTests.cs:98

• This test now uses the default 60s timeout (WaitForSessionStopAsync()), which can make failures hang for a long time. Prefer an explicit short timeout to keep failures quick and deterministic.

        var waitTask = observer.WaitForSessionStopAsync();

test/TestInfrastructure/TestExtensions/Diagnostics/RebalancerDiagnosticObserver.cs:236

• WaitUntilAsync currently always registers a waiter + timer even when timeout is TimeSpan.Zero/negative. That relies on a ThreadPool timer callback to fault the task, which can be delayed under load (the exact flakiness this helper is trying to avoid). Consider treating non-positive timeouts as an immediate timeout by returning a faulted task without enqueuing a waiter.

This issue also appears on line 250 of the same file.

            if (predicate())
            {
                return Task.CompletedTask;
            }

• Files reviewed: 2/2 changed files
• Comments generated: 0 new

[Copilot]

Copilot's findings

Suppressed comments (5)

test/Orleans.Core.Tests/Diagnostics/DiagnosticInfrastructureRegressionTests.cs:120

• Consider using an explicit short timeout here too (instead of the observer's 60s default) so that if the waiter doesn't complete due to a regression, the test fails quickly instead of potentially hanging for a long time.

        var waitTask = observer.WaitForSessionStopAsync();
        ActivationRebalancerEvents.EmitSessionStop(siloAddress, "latest", 1);
        Assert.True(waitTask.IsCompletedSuccessfully);

test/Orleans.Core.Tests/Diagnostics/DiagnosticInfrastructureRegressionTests.cs:83

• These tests now rely on the observer's default 30s/60s timeouts. If the regression reappears, the test can hang for much longer than necessary. Other tests in this file use explicit short timeouts (e.g., 100ms/1s above), so it would be better to keep an explicit timeout here as well to fail fast on regressions.

        var waitTask = observer.WaitForCycleAsync();
        Assert.False(waitTask.IsCompleted);

        ActivationRebalancerEvents.EmitCycleStop(siloAddress, 2, 2, 0.2, TimeSpan.FromMilliseconds(1), false);
        Assert.True(waitTask.IsCompletedSuccessfully);

test/Orleans.Core.Tests/Diagnostics/DiagnosticInfrastructureRegressionTests.cs:102

• This test now uses the observer's default 60s timeout. To keep regressions failing quickly (consistent with other tests in this file which use explicit timeouts), pass an explicit short timeout here too.

        var waitTask = observer.WaitForSessionStopAsync();
        Assert.False(waitTask.IsCompleted);

        ActivationRebalancerEvents.EmitSessionStop(siloAddress, "latest", 2);
        Assert.True(waitTask.IsCompletedSuccessfully);

test/TestInfrastructure/TestExtensions/Diagnostics/RebalancerDiagnosticObserver.cs:322

• Now that each waiter can own a Timer, disposing the observer while there are outstanding waiters can leave timers running and callbacks firing after disposal (potentially affecting later tests and keeping the observer rooted longer than necessary). It would be safer to fail/clear any outstanding waiters and dispose their timers when RebalancerDiagnosticObserver is disposed.

    public void Dispose()
    {
        _subscription?.Dispose();
    }

test/TestInfrastructure/TestExtensions/Diagnostics/RebalancerDiagnosticObserver.cs:256

• WaitFor...Async(TimeSpan.Zero) now relies on a thread-pool Timer firing to fault the waiter. Under thread-pool starvation (the same kind of CI load this helper is meant to harden against), the timer callback can be delayed, causing a supposed immediate timeout to take arbitrarily long (or hang the test until the framework timeout). Consider handling TimeSpan.Zero synchronously while holding _waitersLock so the returned task is faulted immediately without requiring a timer callback.

    private Task WaitUntilAsync(Func<bool> predicate, TimeSpan timeout, Func<string> timeoutMessage)
    {
        lock (_waitersLock)
        {
            if (predicate())
            {
                return Task.CompletedTask;
            }

            var waiter = new ConditionWaiter(predicate);
            _waiters.Add(waiter);
            waiter.StartTimeout(timeout, () => TimeoutWaiter(waiter, timeoutMessage));
            return waiter.Task;
        }
    }

    private Task<TEvent> WaitForEventAsync<TEvent>(
        Func<TEvent, bool> predicate,
        TimeSpan timeout,
        Func<string> timeoutMessage)
        where TEvent : ActivationRebalancerEvents.RebalancerEvent
    {
        lock (_waitersLock)
        {
            var waiter = new EventWaiter<TEvent>(predicate);
            _waiters.Add(waiter);
            waiter.StartTimeout(timeout, () => TimeoutWaiter(waiter, timeoutMessage));
            return waiter.Task;
        }

• Files reviewed: 2/2 changed files
• Comments generated: 0 new