Harden gRPC worker and client against silent disconnects

src/Client/AzureManaged/DurableTaskSchedulerClientExtensions.cs

@@ -4,9 +4,11 @@
 using System.Collections.Concurrent;
 using System.Diagnostics;
 using System.Linq;
+using System.Threading;
 using Azure.Core;
 using Grpc.Net.Client;
 using Microsoft.DurableTask.Client.Grpc;
+using Microsoft.DurableTask.Client.Grpc.Internal;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.DependencyInjection.Extensions;
 using Microsoft.Extensions.Options;
@@ -156,7 +158,92 @@ public void Configure(string? name, GrpcDurableTaskClientOptions options)
             string cacheKey = $"{optionsName}\u001F{source.EndpointAddress}\u001F{source.TaskHubName}\u001F{source.ResourceId}\u001F{credentialType}\u001F{source.AllowInsecureCredentials}\u001F{retryOptionsKey}";
             options.Channel = this.channels.GetOrAdd(
                 cacheKey,
-                _ => new Lazy<GrpcChannel>(source.CreateChannel)).Value;
+                _ => new Lazy<GrpcChannel>(source.CreateChannel, LazyThreadSafetyMode.PublicationOnly)).Value;
+            options.SetChannelRecreator((oldChannel, ct) => this.RecreateChannelAsync(cacheKey, source, oldChannel, ct));
+        }
+
+        /// <summary>
+        /// Atomically swaps the cached channel for the given key with a freshly created one and schedules
+        /// graceful disposal of the old channel after a grace period so any in-flight RPCs from peer
+        /// clients can drain. Returns the currently cached channel if a peer client has already recreated it.
+        /// </summary>
+        Task<GrpcChannel> RecreateChannelAsync(
+            string cacheKey,
+            DurableTaskSchedulerClientOptions source,
+            GrpcChannel oldChannel,
+            CancellationToken cancellation)
+        {
+            cancellation.ThrowIfCancellationRequested();
+
+            if (this.disposed == 1)
+            {
+                throw new ObjectDisposedException(nameof(ConfigureGrpcChannel));
+            }
+
+            if (!this.channels.TryGetValue(cacheKey, out Lazy<GrpcChannel>? currentLazy))
+            {
+                // PublicationOnly avoids permanently caching a transient CreateChannel exception.
+                Lazy<GrpcChannel> created = new(source.CreateChannel, LazyThreadSafetyMode.PublicationOnly);
+                if (this.channels.TryAdd(cacheKey, created))
+                {
+                    return Task.FromResult(created.Value);
+                }
+
+                this.channels.TryGetValue(cacheKey, out currentLazy);
+            }
+
+            if (currentLazy is null)
+            {
+                throw new InvalidOperationException("Failed to obtain a cached gRPC channel after recreation attempt.");
+            }
+
+            if (currentLazy.IsValueCreated && !ReferenceEquals(currentLazy.Value, oldChannel))
+            {
+                // A peer client already swapped in a new channel; reuse it.
+                return Task.FromResult(currentLazy.Value);
+            }
+
+            // Materialize the new channel BEFORE swapping the dictionary so a CreateChannel failure
+            // leaves the existing entry intact. If we swapped a not-yet-materialized Lazy and then
+            // CreateChannel threw, the dictionary would point to a permanently-failing Lazy and the
+            // old channel would have already been queued for disposal — an unrecoverable state.
+            GrpcChannel newChannel = source.CreateChannel();
+            Lazy<GrpcChannel> newLazy = new(newChannel);
+            if (!this.channels.TryUpdate(cacheKey, newLazy, currentLazy))
+            {
+                // Lost the race; whoever won has the freshest entry. Dispose the channel we just
+                // created so it doesn't leak.
+                this.channels.TryGetValue(cacheKey, out Lazy<GrpcChannel>? winner);
+                _ = ScheduleDeferredDisposeAsync(newChannel);
+                return Task.FromResult(winner?.Value ?? newChannel);
+            }
+
+            if (currentLazy.IsValueCreated)
+            {
+                _ = ScheduleDeferredDisposeAsync(currentLazy.Value);
+            }
+
+            return Task.FromResult(newChannel);
+        }
+
+        static async Task ScheduleDeferredDisposeAsync(GrpcChannel channel)
+        {
+            try
+            {
+                await Task.Delay(TimeSpan.FromSeconds(30)).ConfigureAwait(false);
+                await DisposeChannelAsync(channel).ConfigureAwait(false);
+            }
+            catch (Exception ex) when (ex is not OutOfMemoryException
+                                        and not StackOverflowException
+                                        and not ThreadAbortException)
+            {
+                if (ex is not OperationCanceledException and not ObjectDisposedException)
+                {
+                    Trace.TraceError(
+                        "Unexpected exception while deferred-disposing gRPC channel in DurableTaskSchedulerClientExtensions.ScheduleDeferredDisposeAsync: {0}",
+                        ex);
+                }
+            }
         }
 
         /// <inheritdoc/>