Always dispose outside of lock. Use Interlocked.Exchange to clear

Author: ReubenBond

src/Libraries/Microsoft.Extensions.ServiceDiscovery/ServiceEndpointWatcher.cs

@@ -135,17 +135,17 @@ private async Task RefreshAsyncInternal()
         CacheStatus newCacheState;
         try
         {
+            // Capture the existing change token registration, if any, so we can dispose it.
+            // Disposing inside the lock can cause a deadlock since Dispose waits for any in-flight callback to complete,
+            // but the callback may be waiting to acquire _lock.
             IDisposable? changeTokenRegistration;
             lock (_lock)
             {
-                // Capture the existing change token registration, if any, so we can dispose it outside the lock.
-                // Disposing inside the lock can cause a deadlock since Dispose waits for any in-flight callback to complete,
-                // but the callback may be waiting to acquire _lock.
                 changeTokenRegistration = _changeTokenRegistration;
                 _changeTokenRegistration = null;
             }
 
-            // Dispose the existing change token registration, if any, outside the lock to avoid deadlock.
+            // Dispose the existing change token registration, if any.
             changeTokenRegistration?.Dispose();
 
             Log.ResolvingEndpoints(_logger, ServiceName);
@@ -159,6 +159,10 @@ private async Task RefreshAsyncInternal()
             var endpoints = builder.Build();
             newCacheState = CacheStatus.Valid;
 
+            // Capture the existing timer, if any, so we can dispose it outside the lock.
+            // Disposing inside the lock can cause a deadlock.
+            ITimer? pollingTimer = null;
+
             lock (_lock)
             {
                 // Check if we need to poll for updates or if we can register for change notification callbacks.
@@ -167,8 +171,7 @@ private async Task RefreshAsyncInternal()
                     // Initiate a background refresh when the change token fires.
                     _changeTokenRegistration = endpoints.ChangeToken.RegisterChangeCallback(static state => _ = ((ServiceEndpointWatcher)state!).RefreshAsync(force: false), this);
 
-                    // Dispose the existing timer, if any, since we are reliant on change tokens for updates.
-                    _pollingTimer?.Dispose();
+                    pollingTimer = _pollingTimer;
                     _pollingTimer = null;
                 }
                 else
@@ -180,6 +183,8 @@ private async Task RefreshAsyncInternal()
                 newEndpoints = endpoints;
                 newCacheState = CacheStatus.Valid;
             }
+
+            pollingTimer?.Dispose();
         }
         catch (Exception exception)
         {
@@ -236,16 +241,14 @@ private async Task RefreshAsyncInternal()
 
     private void SchedulePollingTimer()
     {
+        ITimer? timerToDispose = null;
         lock (_lock)
         {
             if (_disposalCancellation.IsCancellationRequested)
             {
-                _pollingTimer?.Dispose();
-                _pollingTimer = null;
-                return;
+                timerToDispose = Interlocked.Exchange(ref _pollingTimer, null);
             }
-
-            if (_pollingTimer is null)
+            else if (_pollingTimer is null)
             {
                 _pollingTimer = _timeProvider.CreateTimer(s_pollingAction, this, _options.RefreshPeriod, TimeSpan.Zero);
             }
@@ -254,6 +257,8 @@ private void SchedulePollingTimer()
                 _pollingTimer.Change(_options.RefreshPeriod, TimeSpan.Zero);
             }
         }
+
+        timerToDispose?.Dispose();
     }
 
     /// <inheritdoc/>
@@ -268,16 +273,15 @@ public async ValueTask DisposeAsync()
             _logger.LogError(exception, "Error cancelling disposal cancellation token.");
         }
 
+        // Capture the existing change token registration and polling timer so we can dispose them.
+        // Disposing _changeTokenRegistration inside the lock can cause a deadlock since Dispose waits for any in-flight
+        // callback to complete, but the callback may be waiting to acquire _lock.
         IDisposable? changeTokenRegistration;
         ITimer? pollingTimer;
         lock (_lock)
         {
-            // Capture the existing change token registration and polling timer so we can dispose them outside the lock.
-            // Disposing _changeTokenRegistration inside the lock can cause a deadlock since Dispose waits for any in-flight
-            // callback to complete, but the callback may be waiting to acquire _lock.
             changeTokenRegistration = _changeTokenRegistration;
             _changeTokenRegistration = null;
-
             pollingTimer = _pollingTimer;
             _pollingTimer = null;
         }