xds: reuse GrpcXdsTransport and underlying gRPC channel to the same xDS server by ref-counting

This PR implements reusing the gRPC xDS transport (and underlying gRPC
channel) to the same xDS server by ref-counting, which is already
implemented in gRPC C++
([link](https://github.com/grpc/grpc/blob/5a3a5d53145b94895610825e783a8896a61a3c73/src/core/xds/grpc/xds_transport_grpc.cc#L399-L414))
and gRPC Go
([link](https://github.com/grpc/grpc-go/blob/81c7924ec9f5f4a01c18b82c9d67691c1cd93bd5/internal/xds/clients/grpctransport/grpc_transport.go#L78-L120)).
This optimization is expected to reduce memory footprint of the xDS
management server and xDS enabled clients as channel establishment and
lifecycle management of the connection is expensive.

* Implemented a map to store `GrpcXdsTransport` instances keyed by the
`Bootstrapper.ServerInfo` and each `GrpcXdsTransport` has a ref count.
Note, the map cannot be simply keyed by the xDS server address as the
client could have different channel credentials to the same xDS server,
which should be counted as different transport instances.
* When `GrpcXdsTransportFactory.create()` is called, the existing
transport is reused if it already exists in the map and increment its
ref count, otherwise create a new transport, store it in the map, and
increment its ref count.
* When `GrpcXdsTransport.shutdown()` is called, its ref count is
decremented and the underlying gRPC channel is shut down when its ref
count reaches zero.
* Note this ref-counting of the `GrpcXdsTransport` is different and
orthogonal to the ref-counting of the xDS client keyed by the xDS server
target name to allow for xDS-based fallback per [gRFC
A71](https://github.com/grpc/proposal/blob/master/A71-xds-fallback.md).

Prod risk level: Low
* Reusing the underlying gRPC channel to the xDS server would not affect
the gRPC xDS (ADS/LRS) streams which would be multiplexed on the same
channel, however, this means new xDS (ADS/LRS) streams and RPCs may fail
due to hitting the limit of `MAX_CONCURRENT_STREAMS`.

Tested:
* Verified end-to-end with a xDS enabled gRPC Java client communicating
to multiple different gRPC backend servers behind *different targets*
using the xDS management server for name resolution and endpoint
discovery. Verified gRPC xDS transport creation, ref-counting, reuse,
shutdown, deletion from map when ref count is zero all worked as
expected.

Implementation details / context:
* Used `java.util.concurrent.ConcurrentHashMap` APIs `compute` and
`computeIfPresent` where the entire method invocation is performed
atomically to achieve a concurrent and thread-safe solution which
follows Java best practices.

Alternatives considered: 
* Write own synchronization logic with synchronized block and locks.
After discussion internally, it was preferred to use existing
concurrency libraries which is less error-prone and should offer better
performance.

Author: danielzhaotongliu

xds/src/main/java/io/grpc/xds/GrpcXdsTransportFactory.java

@@ -31,39 +31,74 @@
 import io.grpc.Status;
 import io.grpc.xds.client.Bootstrapper;
 import io.grpc.xds.client.XdsTransportFactory;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
 
+/**
+ * A factory for creating gRPC-based transports for xDS communication.
+ *
+ * <p>WARNING: This class reuses channels when possible, based on the provided {@link
+ * Bootstrapper.ServerInfo} with important considerations. The {@link Bootstrapper.ServerInfo}
+ * includes {@link ChannelCredentials}, which is compared by reference equality. This means every
+ * {@link Bootstrapper.BootstrapInfo} would have non-equal copies of {@link
+ * Bootstrapper.ServerInfo}, even if they all represent the same xDS server configuration. For gRPC
+ * name resolution with the {@code xds} and {@code google-c2p} scheme, this transport sharing works
+ * as expected as it internally reuses a single {@link Bootstrapper.BootstrapInfo} instance.
+ * Otherwise, new transports would be created for each {@link Bootstrapper.ServerInfo} despite them
+ * possibly representing the same xDS server configuration and defeating the purpose of transport
+ * sharing.
+ */
 final class GrpcXdsTransportFactory implements XdsTransportFactory {
 
   private final CallCredentials callCredentials;
+  // The map of xDS server info to its corresponding gRPC xDS transport.
+  // This enables reusing and sharing the same underlying gRPC channel.
+  //
+  // NOTE: ConcurrentHashMap is used as a per-entry lock and all reads and writes must be a mutation
+  // via the ConcurrentHashMap APIs to acquire the per-entry lock in order to ensure thread safety
+  // for reference counting of each GrpcXdsTransport instance.
+  private static final Map<Bootstrapper.ServerInfo, GrpcXdsTransport> xdsServerInfoToTransportMap =
+      new ConcurrentHashMap<>();
 
   GrpcXdsTransportFactory(CallCredentials callCredentials) {
     this.callCredentials = callCredentials;
   }
 
   @Override
   public XdsTransport create(Bootstrapper.ServerInfo serverInfo) {
-    return new GrpcXdsTransport(serverInfo, callCredentials);
+    return xdsServerInfoToTransportMap.compute(
+        serverInfo,
+        (info, transport) -> {
+          if (transport == null) {
+            transport = new GrpcXdsTransport(serverInfo, callCredentials);
+          }
+          ++transport.refCount;
+          return transport;
+        });
   }
 
   @VisibleForTesting
   public XdsTransport createForTest(ManagedChannel channel) {
-    return new GrpcXdsTransport(channel, callCredentials);
+    return new GrpcXdsTransport(channel, callCredentials, null);
   }
 
   @VisibleForTesting
   static class GrpcXdsTransport implements XdsTransport {
 
     private final ManagedChannel channel;
     private final CallCredentials callCredentials;
+    private final Bootstrapper.ServerInfo serverInfo;
+    // Must only be accessed via the ConcurrentHashMap APIs which act as the locking methods.
+    private int refCount = 0;
 
     public GrpcXdsTransport(Bootstrapper.ServerInfo serverInfo) {
       this(serverInfo, null);
     }
 
     @VisibleForTesting
     public GrpcXdsTransport(ManagedChannel channel) {
-      this(channel, null);
+      this(channel, null, null);
     }
 
     public GrpcXdsTransport(Bootstrapper.ServerInfo serverInfo, CallCredentials callCredentials) {
@@ -73,12 +108,17 @@ public GrpcXdsTransport(Bootstrapper.ServerInfo serverInfo, CallCredentials call
           .keepAliveTime(5, TimeUnit.MINUTES)
           .build();
       this.callCredentials = callCredentials;
+      this.serverInfo = serverInfo;
     }
 
     @VisibleForTesting
-    public GrpcXdsTransport(ManagedChannel channel, CallCredentials callCredentials) {
+    public GrpcXdsTransport(
+        ManagedChannel channel,
+        CallCredentials callCredentials,
+        Bootstrapper.ServerInfo serverInfo) {
       this.channel = checkNotNull(channel, "channel");
       this.callCredentials = callCredentials;
+      this.serverInfo = serverInfo;
     }
 
     @Override
@@ -98,7 +138,19 @@ public <ReqT, RespT> StreamingCall<ReqT, RespT> createStreamingCall(
 
     @Override
     public void shutdown() {
-      channel.shutdown();
+      if (serverInfo == null) {
+        channel.shutdown();
+        return;
+      }
+      xdsServerInfoToTransportMap.computeIfPresent(
+          serverInfo,
+          (info, transport) -> {
+            if (--transport.refCount == 0) { // Prefix decrement and return the updated value.
+              transport.channel.shutdown();
+              return null; // Remove mapping.
+            }
+            return transport;
+          });
     }
 
     private class XdsStreamingCall<ReqT, RespT> implements

xds/src/test/java/io/grpc/xds/GrpcXdsTransportFactoryTest.java

@@ -30,20 +30,24 @@
 import io.grpc.Server;
 import io.grpc.Status;
 import io.grpc.stub.StreamObserver;
+import io.grpc.testing.GrpcCleanupRule;
 import io.grpc.xds.client.Bootstrapper;
 import io.grpc.xds.client.XdsTransportFactory;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
 @RunWith(JUnit4.class)
 public class GrpcXdsTransportFactoryTest {
 
+  @Rule public final GrpcCleanupRule grpcCleanupRule = new GrpcCleanupRule();
+
   private Server server;
 
   @Before
@@ -118,6 +122,59 @@ public void callApis() throws Exception {
     xdsTransport.shutdown();
   }
 
+  @Test
+  public void refCountedXdsTransport_sameXdsServerAddress_returnsExistingTransport() {
+    Bootstrapper.ServerInfo xdsServerInfo =
+        Bootstrapper.ServerInfo.create(
+            "localhost:" + server.getPort(), InsecureChannelCredentials.create());
+    GrpcXdsTransportFactory xdsTransportFactory = new GrpcXdsTransportFactory(null);
+    // Calling create() for the first time creates a new GrpcXdsTransport instance.
+    // The ref count was previously 0 and now is 1.
+    XdsTransportFactory.XdsTransport transport1 = xdsTransportFactory.create(xdsServerInfo);
+    // Calling create() for the second time to the same xDS server address returns the same
+    // GrpcXdsTransport instance. The ref count was previously 1 and now is 2.
+    XdsTransportFactory.XdsTransport transport2 = xdsTransportFactory.create(xdsServerInfo);
+    assertThat(transport1).isSameInstanceAs(transport2);
+    // Calling shutdown() for the first time does not shut down the GrpcXdsTransport instance.
+    // The ref count was previously 2 and now is 1.
+    transport1.shutdown();
+    // Calling shutdown() for the second time shuts down the GrpcXdsTransport instance.
+    // The ref count was previously 1 and now is 0.
+    transport2.shutdown();
+  }
+
+  @Test
+  public void refCountedXdsTransport_differentXdsServerAddress_returnsDifferentTransport()
+      throws Exception {
+    // Create and start a second xDS server on a different port.
+    Server server2 =
+        grpcCleanupRule.register(
+            Grpc.newServerBuilderForPort(0, InsecureServerCredentials.create())
+                .addService(echoAdsService())
+                .build()
+                .start());
+    Bootstrapper.ServerInfo xdsServerInfo1 =
+        Bootstrapper.ServerInfo.create(
+            "localhost:" + server.getPort(), InsecureChannelCredentials.create());
+    Bootstrapper.ServerInfo xdsServerInfo2 =
+        Bootstrapper.ServerInfo.create(
+            "localhost:" + server2.getPort(), InsecureChannelCredentials.create());
+    GrpcXdsTransportFactory xdsTransportFactory = new GrpcXdsTransportFactory(null);
+    // Calling create() to the first xDS server creates a new GrpcXdsTransport instance.
+    // The ref count was previously 0 and now is 1.
+    XdsTransportFactory.XdsTransport transport1 = xdsTransportFactory.create(xdsServerInfo1);
+    // Calling create() to the second xDS server creates a different GrpcXdsTransport instance.
+    // The ref count was previously 0 and now is 1.
+    XdsTransportFactory.XdsTransport transport2 = xdsTransportFactory.create(xdsServerInfo2);
+    assertThat(transport1).isNotSameInstanceAs(transport2);
+    // Calling shutdown() shuts down the GrpcXdsTransport instance for the first xDS server.
+    // The ref count was previously 1 and now is 0.
+    transport1.shutdown();
+    // Calling shutdown() shuts down the GrpcXdsTransport instance for the second xDS server.
+    // The ref count was previously 1 and now is 0.
+    transport2.shutdown();
+  }
+
   private static class FakeEventHandler implements
       XdsTransportFactory.EventHandler<DiscoveryResponse> {
     private final BlockingQueue<DiscoveryResponse> respQ = new LinkedBlockingQueue<>();