Fix GrpcChannel handle leak in AzureManaged backend

src/Client/AzureManaged/DurableTaskSchedulerClientExtensions.cs

@@ -1,7 +1,10 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+using System.Collections.Concurrent;
+using System.Linq;
 using Azure.Core;
+using Grpc.Net.Client;
 using Microsoft.DurableTask.Client.Grpc;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.DependencyInjection.Extensions;
@@ -99,11 +102,23 @@ static void ConfigureSchedulerOptions(
     /// <summary>
     /// Configuration class that sets up gRPC channels for client options
     /// using the provided Durable Task Scheduler options.
+    /// Channels are cached per configuration key and disposed when the service provider is disposed.
     /// </summary>
-    /// <param name="schedulerOptions">Monitor for accessing the current scheduler options configuration.</param>
-    class ConfigureGrpcChannel(IOptionsMonitor<DurableTaskSchedulerClientOptions> schedulerOptions) :
-        IConfigureNamedOptions<GrpcDurableTaskClientOptions>
+    sealed class ConfigureGrpcChannel : IConfigureNamedOptions<GrpcDurableTaskClientOptions>, IAsyncDisposable
     {
+        readonly IOptionsMonitor<DurableTaskSchedulerClientOptions> schedulerOptions;
+        readonly ConcurrentDictionary<string, Lazy<GrpcChannel>> channels = new();
+        int disposed;
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="ConfigureGrpcChannel"/> class.
+        /// </summary>
+        /// <param name="schedulerOptions">Monitor for accessing the current scheduler options configuration.</param>
+        public ConfigureGrpcChannel(IOptionsMonitor<DurableTaskSchedulerClientOptions> schedulerOptions)
+        {
+            this.schedulerOptions = schedulerOptions;
+        }
+
         /// <summary>
         /// Configures the default named options instance.
         /// </summary>
@@ -117,8 +132,72 @@ class ConfigureGrpcChannel(IOptionsMonitor<DurableTaskSchedulerClientOptions> sc
         /// <param name="options">The options instance to configure.</param>
         public void Configure(string? name, GrpcDurableTaskClientOptions options)
         {
-            DurableTaskSchedulerClientOptions source = schedulerOptions.Get(name ?? Options.DefaultName);
-            options.Channel = source.CreateChannel();
+#if NET7_0_OR_GREATER
+            ObjectDisposedException.ThrowIf(this.disposed == 1, this);
+#else
+            if (this.disposed == 1)
+            {
+                throw new ObjectDisposedException(nameof(ConfigureGrpcChannel));
+            }
+#endif
+
+            string optionsName = name ?? Options.DefaultName;
+            DurableTaskSchedulerClientOptions source = this.schedulerOptions.Get(optionsName);
+
+            // Create a cache key based on the options name, endpoint, and task hub.
+            // This ensures channels are reused for the same configuration
+            // but separate channels are created for different configurations.
+            string cacheKey = $"{optionsName}:{source.EndpointAddress}:{source.TaskHubName}";
+            options.Channel = this.channels.GetOrAdd(
+                cacheKey,
+                _ => new Lazy<GrpcChannel>(source.CreateChannel)).Value;
+        }
+
+        /// <inheritdoc/>
+        public async ValueTask DisposeAsync()
+        {
+            if (Interlocked.Exchange(ref this.disposed, 1) == 1)
+            {
+                return;
+            }
+
+            List<Exception>? exceptions = null;
+            foreach (Lazy<GrpcChannel> channel in this.channels.Values.Where(lazy => lazy.IsValueCreated))
+            {
+                try
+                {
+                    await DisposeChannelAsync(channel.Value).ConfigureAwait(false);
+                }
+                catch (Exception ex)
+                {
+                    exceptions ??= new List<Exception>();
+                    exceptions.Add(ex);
+                }
+            }
+
+            this.channels.Clear();
+            GC.SuppressFinalize(this);
+
+            if (exceptions is { Count: > 0 })
+            {
+                throw new AggregateException(exceptions);
+            }
+        }
+
+        static async Task DisposeChannelAsync(GrpcChannel channel)
+        {
+            // ShutdownAsync is the graceful way to close a gRPC channel.
+            using (channel)
+            {
+                try
+                {
+                    await channel.ShutdownAsync().ConfigureAwait(false);
+                }
+                catch (Exception ex) when (ex is OperationCanceledException or ObjectDisposedException)
+                {
+                    // Ignore expected shutdown/disposal errors
+                }
+            }
         }
     }
 }

src/Worker/AzureManaged/DurableTaskSchedulerWorkerExtensions.cs

@@ -1,7 +1,10 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+using System.Collections.Concurrent;
+using System.Linq;
 using Azure.Core;
+using Grpc.Net.Client;
 using Microsoft.DurableTask.Worker.Grpc;
 using Microsoft.DurableTask.Worker.Grpc.Internal;
 using Microsoft.Extensions.DependencyInjection;
@@ -101,11 +104,23 @@ static void ConfigureSchedulerOptions(
     /// <summary>
     /// Configuration class that sets up gRPC channels for worker options
     /// using the provided Durable Task Scheduler options.
+    /// Channels are cached per configuration key and disposed when the service provider is disposed.
     /// </summary>
-    /// <param name="schedulerOptions">Monitor for accessing the current scheduler options configuration.</param>
-    class ConfigureGrpcChannel(IOptionsMonitor<DurableTaskSchedulerWorkerOptions> schedulerOptions) :
-        IConfigureNamedOptions<GrpcDurableTaskWorkerOptions>
+    sealed class ConfigureGrpcChannel : IConfigureNamedOptions<GrpcDurableTaskWorkerOptions>, IAsyncDisposable
     {
+        readonly IOptionsMonitor<DurableTaskSchedulerWorkerOptions> schedulerOptions;
+        readonly ConcurrentDictionary<string, Lazy<GrpcChannel>> channels = new();
+        int disposed;
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="ConfigureGrpcChannel"/> class.
+        /// </summary>
+        /// <param name="schedulerOptions">Monitor for accessing the current scheduler options configuration.</param>
+        public ConfigureGrpcChannel(IOptionsMonitor<DurableTaskSchedulerWorkerOptions> schedulerOptions)
+        {
+            this.schedulerOptions = schedulerOptions;
+        }
+
         /// <summary>
         /// Configures the default named options instance.
         /// </summary>
@@ -119,9 +134,74 @@ class ConfigureGrpcChannel(IOptionsMonitor<DurableTaskSchedulerWorkerOptions> sc
         /// <param name="options">The options instance to configure.</param>
         public void Configure(string? name, GrpcDurableTaskWorkerOptions options)
         {
-            DurableTaskSchedulerWorkerOptions source = schedulerOptions.Get(name ?? Options.DefaultName);
-            options.Channel = source.CreateChannel();
+#if NET7_0_OR_GREATER
+            ObjectDisposedException.ThrowIf(this.disposed == 1, this);
+#else
+            if (this.disposed == 1)
+            {
+                throw new ObjectDisposedException(nameof(ConfigureGrpcChannel));
+            }
+#endif
+
+            string optionsName = name ?? Options.DefaultName;
+            DurableTaskSchedulerWorkerOptions source = this.schedulerOptions.Get(optionsName);
+
+            // Create a cache key based on the options name, endpoint, and task hub.
+            // This ensures channels are reused for the same configuration
+            // but separate channels are created for different configurations.
+            // Use a delimiter character (\u001F) that will not appear in typical endpoint URIs.
+            string cacheKey = $"{optionsName}\u001F{source.EndpointAddress}\u001F{source.TaskHubName}";
+            options.Channel = this.channels.GetOrAdd(
+                cacheKey,
+                _ => new Lazy<GrpcChannel>(source.CreateChannel)).Value;
             options.ConfigureForAzureManaged();
         }
+
+        /// <inheritdoc/>
+        public async ValueTask DisposeAsync()
+        {
+            if (Interlocked.Exchange(ref this.disposed, 1) == 1)
+            {
+                return;
+            }
+
+            List<Exception>? exceptions = null;
+            foreach (Lazy<GrpcChannel> channel in this.channels.Values.Where(lazy => lazy.IsValueCreated))
+            {
+                try
+                {
+                    await DisposeChannelAsync(channel.Value).ConfigureAwait(false);
+                }
+                catch (Exception ex)
+                {
+                    exceptions ??= new List<Exception>();
+                    exceptions.Add(ex);
+                }
+            }
+
+            this.channels.Clear();
+            GC.SuppressFinalize(this);
+
+            if (exceptions is { Count: > 0 })
+            {
+                throw new AggregateException(exceptions);
+            }
+        }
+
+        static async Task DisposeChannelAsync(GrpcChannel channel)
+        {
+            // ShutdownAsync is the graceful way to close a gRPC channel.
+            using (channel)
+            {
+                try
+                {
+                    await channel.ShutdownAsync().ConfigureAwait(false);
+                }
+                catch (Exception ex) when (ex is OperationCanceledException or ObjectDisposedException)
+                {
+                    // Ignore expected shutdown/disposal errors
+                }
+            }
+        }
     }
 }

test/Client/AzureManaged.Tests/DurableTaskSchedulerClientExtensionsTests.cs

@@ -5,6 +5,7 @@
 using Azure.Identity;
 using FluentAssertions;
 using Grpc.Core;
+using Grpc.Net.Client;
 using Microsoft.DurableTask.Client.Grpc;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Options;
@@ -280,4 +281,120 @@ public void UseDurableTaskScheduler_WithConnectionStringAndRetryOptions_ShouldCo
             clientOptions.RetryOptions.RetryableStatusCodes.Should().Contain(StatusCode.Unknown);
         }
     }
+
+    [Fact]
+    public void UseDurableTaskScheduler_SameConfiguration_ReusesSameChannel()
+    {
+        // Arrange
+        ServiceCollection services = new ServiceCollection();
+        Mock<IDurableTaskClientBuilder> mockBuilder = new Mock<IDurableTaskClientBuilder>();
+        mockBuilder.Setup(b => b.Services).Returns(services);
+        DefaultAzureCredential credential = new DefaultAzureCredential();
+
+        // Act
+        mockBuilder.Object.UseDurableTaskScheduler(ValidEndpoint, ValidTaskHub, credential);
+        using ServiceProvider provider = services.BuildServiceProvider();
+
+        // Resolve options multiple times to trigger channel configuration
+        IOptionsFactory<GrpcDurableTaskClientOptions> optionsFactory = provider.GetRequiredService<IOptionsFactory<GrpcDurableTaskClientOptions>>();
+        GrpcDurableTaskClientOptions options1 = optionsFactory.Create(Options.DefaultName);
+        GrpcDurableTaskClientOptions options2 = optionsFactory.Create(Options.DefaultName);
+
+        // Assert
+        options1.Channel.Should().NotBeNull();
+        options2.Channel.Should().NotBeNull();
+        options1.Channel.Should().BeSameAs(options2.Channel, "same configuration should reuse the same channel");
+    }
+
+    [Fact]
+    public void UseDurableTaskScheduler_DifferentNamedOptions_UsesSeparateChannels()
+    {
+        // Arrange
+        ServiceCollection services = new ServiceCollection();
+        Mock<IDurableTaskClientBuilder> mockBuilder1 = new Mock<IDurableTaskClientBuilder>();
+        Mock<IDurableTaskClientBuilder> mockBuilder2 = new Mock<IDurableTaskClientBuilder>();
+        mockBuilder1.Setup(b => b.Services).Returns(services);
+        mockBuilder1.Setup(b => b.Name).Returns("client1");
+        mockBuilder2.Setup(b => b.Services).Returns(services);
+        mockBuilder2.Setup(b => b.Name).Returns("client2");
+        DefaultAzureCredential credential = new DefaultAzureCredential();
+
+        // Act - configure two different named clients with different endpoints
+        mockBuilder1.Object.UseDurableTaskScheduler("endpoint1.westus3.durabletask.io", ValidTaskHub, credential);
+        mockBuilder2.Object.UseDurableTaskScheduler("endpoint2.westus3.durabletask.io", ValidTaskHub, credential);
+        ServiceProvider provider = services.BuildServiceProvider();
+
+        // Resolve options for both named clients
+        IOptionsMonitor<GrpcDurableTaskClientOptions> optionsMonitor = provider.GetRequiredService<IOptionsMonitor<GrpcDurableTaskClientOptions>>();
+        GrpcDurableTaskClientOptions options1 = optionsMonitor.Get("client1");
+        GrpcDurableTaskClientOptions options2 = optionsMonitor.Get("client2");
+
+        // Assert
+        options1.Channel.Should().NotBeNull();
+        options2.Channel.Should().NotBeNull();
+        options1.Channel.Should().NotBeSameAs(options2.Channel, "different configurations should use different channels");
+    }
+
+    [Fact]
+    public async Task UseDurableTaskScheduler_ServiceProviderDispose_DisposesChannels()
+    {
+        // Arrange
+        ServiceCollection services = new ServiceCollection();
+        Mock<IDurableTaskClientBuilder> mockBuilder = new Mock<IDurableTaskClientBuilder>();
+        mockBuilder.Setup(b => b.Services).Returns(services);
+        DefaultAzureCredential credential = new DefaultAzureCredential();
+
+        // Act
+        mockBuilder.Object.UseDurableTaskScheduler(ValidEndpoint, ValidTaskHub, credential);
+        ServiceProvider provider = services.BuildServiceProvider();
+
+        // Resolve options to trigger channel creation
+        IOptionsMonitor<GrpcDurableTaskClientOptions> optionsMonitor = provider.GetRequiredService<IOptionsMonitor<GrpcDurableTaskClientOptions>>();
+        GrpcDurableTaskClientOptions options = optionsMonitor.Get(Options.DefaultName);
+        options.Channel.Should().NotBeNull();
+        GrpcChannel channel = options.Channel!;
+
+        // Dispose the service provider - this should dispose the ConfigureGrpcChannel which disposes channels
+        await provider.DisposeAsync();
+
+        // Assert - verify the channel was disposed by checking it throws ObjectDisposedException
+        Action action = () => channel.CreateCallInvoker();
+        action.Should().Throw<ObjectDisposedException>("channel should be disposed after provider disposal");
+
+        // Also verify that creating a new provider and getting options still works
+        ServiceCollection services2 = new ServiceCollection();
+        Mock<IDurableTaskClientBuilder> mockBuilder2 = new Mock<IDurableTaskClientBuilder>();
+        mockBuilder2.Setup(b => b.Services).Returns(services2);
+        mockBuilder2.Object.UseDurableTaskScheduler(ValidEndpoint, ValidTaskHub, credential);
+        await using ServiceProvider provider2 = services2.BuildServiceProvider();
+
+        IOptionsMonitor<GrpcDurableTaskClientOptions> newOptionsMonitor = provider2.GetRequiredService<IOptionsMonitor<GrpcDurableTaskClientOptions>>();
+        GrpcDurableTaskClientOptions newOptions = newOptionsMonitor.Get(Options.DefaultName);
+        newOptions.Channel.Should().NotBeNull();
+        newOptions.Channel.Should().NotBeSameAs(channel, "new provider should create a new channel");
+    }
+
+    [Fact]
+    public async Task UseDurableTaskScheduler_ConfigureAfterDispose_ThrowsObjectDisposedException()
+    {
+        // Arrange
+        ServiceCollection services = new ServiceCollection();
+        Mock<IDurableTaskClientBuilder> mockBuilder = new Mock<IDurableTaskClientBuilder>();
+        mockBuilder.Setup(b => b.Services).Returns(services);
+        DefaultAzureCredential credential = new DefaultAzureCredential();
+
+        // Act
+        mockBuilder.Object.UseDurableTaskScheduler(ValidEndpoint, ValidTaskHub, credential);
+        ServiceProvider provider = services.BuildServiceProvider();
+
+        // Resolve options monitor before disposal
+        IOptionsMonitor<GrpcDurableTaskClientOptions> optionsMonitor = provider.GetRequiredService<IOptionsMonitor<GrpcDurableTaskClientOptions>>();
+
+        // Dispose the service provider
+        await provider.DisposeAsync();
+
+        // Assert - attempting to get options after disposal should throw
+        Action action = () => optionsMonitor.Get(Options.DefaultName);
+        action.Should().Throw<ObjectDisposedException>("configuring options after disposal should throw");
+    }
 }