From 158ce9cbbda48e1da1a8962d82fcdbef1767c7ea Mon Sep 17 00:00:00 2001
From: Amanda Tarafa Mas <atarafamas@google.com>
Date: Fri, 4 Oct 2024 11:51:25 -0700
Subject: [PATCH] feat: Support TPC connection sharing across gRPC channels

Closes #13655
---
 .../Google.Cloud.PubSub.V1/ClientBuilderChannelHelper.cs  | 8 +++++---
 .../Google.Cloud.PubSub.V1/PublisherClientBuilder.cs      | 8 +++++++-
 .../Google.Cloud.PubSub.V1/SubscriberClientBuilder.cs     | 8 +++++++-
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/apis/Google.Cloud.PubSub.V1/Google.Cloud.PubSub.V1/ClientBuilderChannelHelper.cs b/apis/Google.Cloud.PubSub.V1/Google.Cloud.PubSub.V1/ClientBuilderChannelHelper.cs
index c5f3751201cf..94edbc1c7091 100644
--- a/apis/Google.Cloud.PubSub.V1/Google.Cloud.PubSub.V1/ClientBuilderChannelHelper.cs
+++ b/apis/Google.Cloud.PubSub.V1/Google.Cloud.PubSub.V1/ClientBuilderChannelHelper.cs
@@ -27,9 +27,11 @@ internal static class ClientBuilderChannelHelper
         // Set max metadata size to 4 MB i.e., 4194304 bytes.
         .WithCustomOption("grpc.max_metadata_size", 4194304);
 
-        internal static GrpcChannelOptions GetUnlimitedSendReceiveChannelOptions() =>
-            // Use a random arg to prevent sub-channel re-use in gRPC, so each channel uses its own connection.
-            s_unlimitedSendReceiveChannelOptions.WithCustomOption("sub-channel-separator", Guid.NewGuid().ToString());
+        internal static GrpcChannelOptions GetUnlimitedSendReceiveChannelOptions(bool sharedConnections) =>
+            sharedConnections ?
+                s_unlimitedSendReceiveChannelOptions :
+                // Use a random arg to prevent sub-channel re-use in gRPC, so each channel uses its own connection.
+                s_unlimitedSendReceiveChannelOptions.WithCustomOption("sub-channel-separator", Guid.NewGuid().ToString());
 
     internal static Task DisposeChannelAsync(ChannelBase channel)
     {
diff --git a/apis/Google.Cloud.PubSub.V1/Google.Cloud.PubSub.V1/PublisherClientBuilder.cs b/apis/Google.Cloud.PubSub.V1/Google.Cloud.PubSub.V1/PublisherClientBuilder.cs
index 59dd5d2bdc8d..2f47240c889c 100644
--- a/apis/Google.Cloud.PubSub.V1/Google.Cloud.PubSub.V1/PublisherClientBuilder.cs
+++ b/apis/Google.Cloud.PubSub.V1/Google.Cloud.PubSub.V1/PublisherClientBuilder.cs
@@ -56,6 +56,12 @@ public int? ClientCount
         }
     }
 
+    /// <summary>
+    /// Whether the <see cref="PublisherServiceApiClient"/>s may share gRPC connections or not.
+    /// Defaults to false, meaning, each client has its own gRPC connection.
+    /// </summary>
+    public bool ClientsShareConnections { get; set; }
+
     /// <summary>
     /// Additional settings for batching, message ordering etc. Default settings will be used if this is null.
     /// </summary>
@@ -108,7 +114,7 @@ private async Task<PublisherClient> BuildAsyncImpl(CancellationToken cancellatio
         var shutdowns = new Func<Task>[clientCount];
         for (int i = 0; i < clientCount; i++)
         {
-            var grpcChannelOptions = ClientBuilderChannelHelper.GetUnlimitedSendReceiveChannelOptions();
+            var grpcChannelOptions = ClientBuilderChannelHelper.GetUnlimitedSendReceiveChannelOptions(ClientsShareConnections);
 
             var builder = new PublisherServiceApiClientBuilder(this, grpcChannelOptions);
             clients[i] = isAsync
diff --git a/apis/Google.Cloud.PubSub.V1/Google.Cloud.PubSub.V1/SubscriberClientBuilder.cs b/apis/Google.Cloud.PubSub.V1/Google.Cloud.PubSub.V1/SubscriberClientBuilder.cs
index ca7ebe15d143..142dff99b1c0 100644
--- a/apis/Google.Cloud.PubSub.V1/Google.Cloud.PubSub.V1/SubscriberClientBuilder.cs
+++ b/apis/Google.Cloud.PubSub.V1/Google.Cloud.PubSub.V1/SubscriberClientBuilder.cs
@@ -56,6 +56,12 @@ public int? ClientCount
         }
     }
 
+    /// <summary>
+    /// Whether the <see cref="PublisherServiceApiClient"/>s may share gRPC connections or not.
+    /// Defaults to false, meaning, each client has its own gRPC connection.
+    /// </summary>
+    public bool ClientsShareConnections { get; set; }
+
     /// <summary>
     /// Additional settings for batching, message ordering etc. Default settings will be used if this is null.
     /// </summary>
@@ -108,7 +114,7 @@ private async Task<SubscriberClient> BuildAsyncImpl(CancellationToken cancellati
         var shutdowns = new Func<Task>[clientCount];
         for (int i = 0; i < clientCount; i++)
         {
-            var grpcChannelOptions = ClientBuilderChannelHelper.GetUnlimitedSendReceiveChannelOptions();
+            var grpcChannelOptions = ClientBuilderChannelHelper.GetUnlimitedSendReceiveChannelOptions(ClientsShareConnections);
 
             var builder = new SubscriberServiceApiClientBuilder(this, grpcChannelOptions);
             clients[i] = isAsync