Remote: gRPC load balancing. (Part 5)

Refactor ReferenceCountedChannel to use DynamicConnectionPool when creating new calls. This change allows existing remote execution/cache client dynamically create new connections on demand.

This change includes rxjava3 to final jar so the size of install_base is increased (~6M for macOS).

PiperOrigin-RevId: 359687769

Author: Googler

src/main/java/com/google/devtools/build/lib/remote/BUILD

@@ -42,7 +42,6 @@ java_library(
         ":ExecutionStatusException",
         ":ReferenceCountedChannel",
         ":Retrier",
-        "//src/main/java/com/google/devtools/build/lib:build-request-options",
         "//src/main/java/com/google/devtools/build/lib:runtime",
         "//src/main/java/com/google/devtools/build/lib/actions",
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
@@ -75,6 +74,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/remote/common",
         "//src/main/java/com/google/devtools/build/lib/remote/disk",
         "//src/main/java/com/google/devtools/build/lib/remote/downloader",
+        "//src/main/java/com/google/devtools/build/lib/remote/grpc",
         "//src/main/java/com/google/devtools/build/lib/remote/http",
         "//src/main/java/com/google/devtools/build/lib/remote/logging",
         "//src/main/java/com/google/devtools/build/lib/remote/merkletree",
@@ -98,6 +98,7 @@ java_library(
         "//third_party:guava",
         "//third_party:jsr305",
         "//third_party:netty",
+        "//third_party:rxjava3",
         "//third_party/grpc:grpc-jar",
         "//third_party/protobuf:protobuf_java",
         "//third_party/protobuf:protobuf_java_util",
@@ -128,10 +129,9 @@ java_library(
     name = "ReferenceCountedChannel",
     srcs = [
         "ReferenceCountedChannel.java",
-        "ReferenceCountedChannelPool.java",
     ],
     deps = [
-        "//third_party:guava",
+        "//src/main/java/com/google/devtools/build/lib/remote/grpc",
         "//third_party:netty",
         "//third_party/grpc:grpc-jar",
     ],

src/main/java/com/google/devtools/build/lib/remote/GoogleChannelConnectionFactory.java

@@ -0,0 +1,60 @@
+// Copyright 2021 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.remote;
+
+import com.google.devtools.build.lib.authandtls.AuthAndTLSOptions;
+import com.google.devtools.build.lib.remote.grpc.ChannelConnectionFactory;
+import io.grpc.ClientInterceptor;
+import io.reactivex.rxjava3.core.Single;
+import java.util.List;
+
+/**
+ * A {@link ChannelConnectionFactory} which creates {@link ChannelConnection} using {@link
+ * ChannelFactory}.
+ */
+public class GoogleChannelConnectionFactory implements ChannelConnectionFactory {
+  private final ChannelFactory channelFactory;
+  private final String target;
+  private final String proxy;
+  private final AuthAndTLSOptions options;
+  private final List<ClientInterceptor> interceptors;
+  private final int maxConcurrency;
+
+  public GoogleChannelConnectionFactory(
+      ChannelFactory channelFactory,
+      String target,
+      String proxy,
+      AuthAndTLSOptions options,
+      List<ClientInterceptor> interceptors,
+      int maxConcurrency) {
+    this.channelFactory = channelFactory;
+    this.target = target;
+    this.proxy = proxy;
+    this.options = options;
+    this.interceptors = interceptors;
+    this.maxConcurrency = maxConcurrency;
+  }
+
+  @Override
+  public Single<ChannelConnection> create() {
+    return Single.fromCallable(
+        () ->
+            new ChannelConnection(channelFactory.newChannel(target, proxy, options, interceptors)));
+  }
+
+  @Override
+  public int maxConcurrency() {
+    return maxConcurrency;
+  }
+}

src/main/java/com/google/devtools/build/lib/remote/ReferenceCountedChannel.java

@@ -13,84 +13,101 @@
 // limitations under the License.
 package com.google.devtools.build.lib.remote;
 
+import com.google.devtools.build.lib.remote.grpc.ChannelConnectionFactory;
+import com.google.devtools.build.lib.remote.grpc.ChannelConnectionFactory.ChannelConnection;
+import com.google.devtools.build.lib.remote.grpc.DynamicConnectionPool;
+import com.google.devtools.build.lib.remote.grpc.SharedConnectionFactory.SharedConnection;
 import io.grpc.CallOptions;
+import io.grpc.Channel;
 import io.grpc.ClientCall;
-import io.grpc.ManagedChannel;
+import io.grpc.ForwardingClientCall;
+import io.grpc.ForwardingClientCallListener;
+import io.grpc.Metadata;
 import io.grpc.MethodDescriptor;
+import io.grpc.Status;
 import io.netty.util.AbstractReferenceCounted;
 import io.netty.util.ReferenceCounted;
-import java.util.concurrent.TimeUnit;
+import java.io.IOException;
 
 /**
- * A wrapper around a {@link io.grpc.ManagedChannel} exposing a reference count. When instantiated
- * the reference count is 1. {@link ManagedChannel#shutdown()} will be called on the wrapped channel
- * when the reference count reaches 0.
+ * A wrapper around a {@link DynamicConnectionPool} exposing {@link Channel} and a reference count.
+ * When instantiated the reference count is 1. {@link DynamicConnectionPool#close()} will be called
+ * on the wrapped channel when the reference count reaches 0.
  *
  * <p>See {@link ReferenceCounted} for more information about reference counting.
  */
-public class ReferenceCountedChannel extends ManagedChannel implements ReferenceCounted {
-
-  private final ManagedChannel channel;
-  private final AbstractReferenceCounted referenceCounted;
-
-  public ReferenceCountedChannel(ManagedChannel channel) {
-    this(
-        channel,
-        new AbstractReferenceCounted() {
-          @Override
-          protected void deallocate() {
-            channel.shutdown();
+public class ReferenceCountedChannel extends Channel implements ReferenceCounted {
+  private final DynamicConnectionPool dynamicConnectionPool;
+  private final AbstractReferenceCounted referenceCounted =
+      new AbstractReferenceCounted() {
+        @Override
+        protected void deallocate() {
+          try {
+            dynamicConnectionPool.close();
+          } catch (IOException e) {
+            throw new AssertionError(e.getMessage(), e);
           }
+        }
 
-          @Override
-          public ReferenceCounted touch(Object o) {
-            return this;
-          }
-        });
-  }
-
-  protected ReferenceCountedChannel(
-      ManagedChannel channel, AbstractReferenceCounted referenceCounted) {
-    this.channel = channel;
-    this.referenceCounted = referenceCounted;
-  }
+        @Override
+        public ReferenceCounted touch(Object o) {
+          return this;
+        }
+      };
 
-  @Override
-  public ManagedChannel shutdown() {
-    throw new UnsupportedOperationException("Don't call shutdown() directly, but use release() "
-        + "instead.");
+  public ReferenceCountedChannel(ChannelConnectionFactory connectionFactory) {
+    this.dynamicConnectionPool =
+        new DynamicConnectionPool(connectionFactory, connectionFactory.maxConcurrency());
   }
 
-  @Override
   public boolean isShutdown() {
-    return channel.isShutdown();
-  }
-
-  @Override
-  public boolean isTerminated() {
-    return channel.isTerminated();
-  }
-
-  @Override
-  public ManagedChannel shutdownNow() {
-    throw new UnsupportedOperationException("Don't call shutdownNow() directly, but use release() "
-        + "instead.");
-  }
-
-  @Override
-  public boolean awaitTermination(long timeout, TimeUnit timeUnit) throws InterruptedException {
-    return channel.awaitTermination(timeout, timeUnit);
+    return dynamicConnectionPool.isClosed();
+  }
+
+  /** A {@link ClientCall} which call {@link SharedConnection#close()} after the RPC is closed. */
+  static class ConnectionCleanupCall<ReqT, RespT>
+      extends ForwardingClientCall.SimpleForwardingClientCall<ReqT, RespT> {
+    private final SharedConnection connection;
+
+    protected ConnectionCleanupCall(ClientCall<ReqT, RespT> delegate, SharedConnection connection) {
+      super(delegate);
+      this.connection = connection;
+    }
+
+    @Override
+    public void start(Listener<RespT> responseListener, Metadata headers) {
+      super.start(
+          new ForwardingClientCallListener.SimpleForwardingClientCallListener<RespT>(
+              responseListener) {
+            @Override
+            public void onClose(Status status, Metadata trailers) {
+              super.onClose(status, trailers);
+
+              try {
+                connection.close();
+              } catch (IOException e) {
+                throw new AssertionError(e.getMessage(), e);
+              }
+            }
+          },
+          headers);
+    }
   }
 
   @Override
   public <RequestT, ResponseT> ClientCall<RequestT, ResponseT> newCall(
       MethodDescriptor<RequestT, ResponseT> methodDescriptor, CallOptions callOptions) {
-    return channel.<RequestT, ResponseT>newCall(methodDescriptor, callOptions);
+    SharedConnection sharedConnection = dynamicConnectionPool.create().blockingGet();
+    ChannelConnection connection = (ChannelConnection) sharedConnection.getUnderlyingConnection();
+    return new ConnectionCleanupCall<>(
+        connection.getChannel().newCall(methodDescriptor, callOptions), sharedConnection);
   }
 
   @Override
   public String authority() {
-    return channel.authority();
+    SharedConnection sharedConnection = dynamicConnectionPool.create().blockingGet();
+    ChannelConnection connection = (ChannelConnection) sharedConnection.getUnderlyingConnection();
+    return connection.getChannel().authority();
   }
 
   @Override
@@ -131,4 +148,4 @@ public boolean release() {
   public boolean release(int decrement) {
     return referenceCounted.release(decrement);
   }
-}
+}