Add ShutdownChannelsAsync api

2 files changed, 36 insertions, 0 deletions

diff --git a/src/csharp/Grpc.Core/Channel.cs b/src/csharp/Grpc.Core/Channel.cs
index b58a6a7381..e0fc7180da 100644
--- a/src/csharp/Grpc.Core/Channel.cs
+++ b/src/csharp/Grpc.Core/Channel.cs
@@ -88,6 +88,7 @@ namespace Grpc.Core
                     this.handle = ChannelSafeHandle.CreateInsecure(target, nativeChannelArgs);
                 }
             }
+            GrpcEnvironment.RegisterChannel(this);
         }
 
         /// <summary>
@@ -209,6 +210,7 @@ namespace Grpc.Core
                 GrpcPreconditions.CheckState(!shutdownRequested);
                 shutdownRequested = true;
             }
+            GrpcEnvironment.UnregisterChannel(this);
 
             shutdownTokenSource.Cancel();
 
diff --git a/src/csharp/Grpc.Core/GrpcEnvironment.cs b/src/csharp/Grpc.Core/GrpcEnvironment.cs
index 6e56b6e8e3..c25022a5d4 100644
--- a/src/csharp/Grpc.Core/GrpcEnvironment.cs
+++ b/src/csharp/Grpc.Core/GrpcEnvironment.cs
@@ -54,12 +54,15 @@ namespace Grpc.Core
         static int refCount;
         static int? customThreadPoolSize;
         static int? customCompletionQueueCount;
+        static readonly HashSet<Channel> registeredChannels = new HashSet<Channel>();
 
         static ILogger logger = new ConsoleLogger();
 
+        readonly object myLock = new object();
         readonly GrpcThreadPool threadPool;
         readonly DebugStats debugStats = new DebugStats();
         readonly AtomicCounter cqPickerCounter = new AtomicCounter();
+
         bool isClosed;
 
         /// <summary>
@@ -110,6 +113,37 @@ namespace Grpc.Core
             }
         }
 
+        internal static void RegisterChannel(Channel channel)
+        {
+            lock (staticLock)
+            {
+                GrpcPreconditions.CheckNotNull(channel);
+                registeredChannels.Add(channel);
+            }
+        }
+
+        internal static void UnregisterChannel(Channel channel)
+        {
+            lock (staticLock)
+            {
+                GrpcPreconditions.CheckNotNull(channel);
+                GrpcPreconditions.CheckArgument(registeredChannels.Remove(channel), "Channel not found in the registered channels set.");
+            }
+        }
+
+        /// <summary>
+        /// Requests shutdown of all channels created by the current process.
+        /// </summary>
+        public static Task ShutdownChannelsAsync()
+        {
+            HashSet<Channel> snapshot = null;
+            lock (staticLock)
+            {
+                snapshot = new HashSet<Channel>(registeredChannels);
+            }
+            return Task.WhenAll(snapshot.Select((channel) => channel.ShutdownAsync()));
+        }
+
         /// <summary>
         /// Gets application-wide logger used by gRPC.
         /// </summary>