From a7b1dc1956aa0abbe26c7bd0dd7420cdf2677600 Mon Sep 17 00:00:00 2001
From: Reuben Bond <rebond@microsoft.com>
Date: Fri, 31 Jul 2026 14:21:16 -0700
Subject: [PATCH] fix: clean up rebalancer diagnostic waiters

Complete zero-duration waits without a timer callback and fail outstanding waiters when their observer is disposed. Keep timer ownership on the waiter and clean up scheduling failures.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 0c7d7f36-28fd-45f3-aeaf-218a5979a0d7
---
 ...DiagnosticInfrastructureRegressionTests.cs | 20 ++++++-
 .../RebalancerDiagnosticObserver.cs           | 57 +++++++++++++++++--
 2 files changed, 70 insertions(+), 7 deletions(-)

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<TimeoutException>(() => observer.WaitForSessionStopAsync(TimeSpan.Zero));
+        var timedOutWaitTask = observer.WaitForSessionStopAsync(TimeSpan.Zero);
+        Assert.True(timedOutWaitTask.IsCompleted);
+        await Assert.ThrowsAsync<TimeoutException>(() => timedOutWaitTask);
+
+        var timedOutCountWaitTask = observer.WaitForSessionStopCountAsync(1, TimeSpan.Zero);
+        Assert.True(timedOutCountWaitTask.IsCompleted);
+        await Assert.ThrowsAsync<TimeoutException>(() => 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<ObjectDisposedException>(() => 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<Activa
     private readonly object _waitersLock = new();
     private readonly List<Waiter> _waiters = [];
     private IDisposable? _subscription;
+    private bool _disposed;
 
     /// <summary>
     /// Gets all captured cycle start events.
@@ -228,14 +229,21 @@ private Task WaitUntilAsync(Func<bool> predicate, TimeSpan timeout, Func<string>
     {
         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<TEvent> WaitForEventAsync<TEvent>(
     {
         lock (_waitersLock)
         {
+            ObjectDisposedException.ThrowIf(_disposed, this);
+
+            if (timeout == TimeSpan.Zero)
+            {
+                return Task.FromException<TEvent>(new TimeoutException(timeoutMessage()));
+            }
+
             var waiter = new EventWaiter<TEvent>(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()