diff --git a/test/Orleans.Core.Tests/Diagnostics/DiagnosticInfrastructureRegressionTests.cs b/test/Orleans.Core.Tests/Diagnostics/DiagnosticInfrastructureRegressionTests.cs index ee59a4e1ccd..253a103330f 100644 --- a/test/Orleans.Core.Tests/Diagnostics/DiagnosticInfrastructureRegressionTests.cs +++ b/test/Orleans.Core.Tests/Diagnostics/DiagnosticInfrastructureRegressionTests.cs @@ -112,7 +112,13 @@ public async Task RebalancerDiagnosticObserver_WaitAfterTimeout_CanObserveLaterE using var observer = RebalancerDiagnosticObserver.Create(); var siloAddress = SiloAddress.New(new IPEndPoint(IPAddress.Loopback, 12002), 3); - await Assert.ThrowsAsync(() => observer.WaitForSessionStopAsync(TimeSpan.Zero)); + var timedOutWaitTask = observer.WaitForSessionStopAsync(TimeSpan.Zero); + Assert.True(timedOutWaitTask.IsCompleted); + await Assert.ThrowsAsync(() => timedOutWaitTask); + + var timedOutCountWaitTask = observer.WaitForSessionStopCountAsync(1, TimeSpan.Zero); + Assert.True(timedOutCountWaitTask.IsCompleted); + await Assert.ThrowsAsync(() => timedOutCountWaitTask); var waitTask = observer.WaitForSessionStopAsync(); ActivationRebalancerEvents.EmitSessionStop(siloAddress, "latest", 1); @@ -122,6 +128,18 @@ public async Task RebalancerDiagnosticObserver_WaitAfterTimeout_CanObserveLaterE Assert.Equal("latest", result.Reason); } + [Fact, TestCategory("BVT")] + public async Task RebalancerDiagnosticObserver_Dispose_CompletesOutstandingWaiters() + { + var observer = RebalancerDiagnosticObserver.Create(); + var waitTask = observer.WaitForSessionStopAsync(); + + observer.Dispose(); + + Assert.True(waitTask.IsCompleted); + await Assert.ThrowsAsync(() => waitTask); + } + [Fact, TestCategory("BVT")] public void InMemoryLoggerProvider_FormatsStoredThreadId() { diff --git a/test/TestInfrastructure/TestExtensions/Diagnostics/RebalancerDiagnosticObserver.cs b/test/TestInfrastructure/TestExtensions/Diagnostics/RebalancerDiagnosticObserver.cs index 181aff45cd2..521a502eafd 100644 --- a/test/TestInfrastructure/TestExtensions/Diagnostics/RebalancerDiagnosticObserver.cs +++ b/test/TestInfrastructure/TestExtensions/Diagnostics/RebalancerDiagnosticObserver.cs @@ -21,6 +21,7 @@ public sealed class RebalancerDiagnosticObserver : IDisposable, IObserver _waiters = []; private IDisposable? _subscription; + private bool _disposed; /// /// Gets all captured cycle start events. @@ -228,14 +229,21 @@ private Task WaitUntilAsync(Func predicate, TimeSpan timeout, Func { lock (_waitersLock) { + ObjectDisposedException.ThrowIf(_disposed, this); + if (predicate()) { return Task.CompletedTask; } + if (timeout == TimeSpan.Zero) + { + return Task.FromException(new TimeoutException(timeoutMessage())); + } + var waiter = new ConditionWaiter(predicate); - _waiters.Add(waiter); waiter.StartTimeout(timeout, () => TimeoutWaiter(waiter, timeoutMessage)); + _waiters.Add(waiter); return waiter.Task; } } @@ -248,9 +256,16 @@ private Task WaitForEventAsync( { lock (_waitersLock) { + ObjectDisposedException.ThrowIf(_disposed, this); + + if (timeout == TimeSpan.Zero) + { + return Task.FromException(new TimeoutException(timeoutMessage())); + } + var waiter = new EventWaiter(predicate); - _waiters.Add(waiter); waiter.StartTimeout(timeout, () => TimeoutWaiter(waiter, timeoutMessage)); + _waiters.Add(waiter); return waiter.Task; } } @@ -317,7 +332,24 @@ private void SignalWaiters(ActivationRebalancerEvents.RebalancerEvent value) public void Dispose() { - _subscription?.Dispose(); + Interlocked.Exchange(ref _subscription, null)?.Dispose(); + + lock (_waitersLock) + { + if (_disposed) + { + return; + } + + _disposed = true; + foreach (var waiter in _waiters) + { + waiter.TrySetException(new ObjectDisposedException(nameof(RebalancerDiagnosticObserver))); + waiter.StopTimeout(); + } + + _waiters.Clear(); + } } private abstract class Waiter @@ -330,13 +362,26 @@ private abstract class Waiter public void StartTimeout(TimeSpan timeout, Action callback) { - var timer = new System.Threading.Timer( + if (timeout == System.Threading.Timeout.InfiniteTimeSpan) + { + return; + } + + _timeoutTimer = new System.Threading.Timer( static state => ((Action)state!).Invoke(), callback, System.Threading.Timeout.InfiniteTimeSpan, System.Threading.Timeout.InfiniteTimeSpan); - Interlocked.Exchange(ref _timeoutTimer, timer)?.Dispose(); - timer.Change(timeout, System.Threading.Timeout.InfiniteTimeSpan); + + try + { + _timeoutTimer.Change(timeout, System.Threading.Timeout.InfiniteTimeSpan); + } + catch + { + StopTimeout(); + throw; + } } public void StopTimeout()