Skip to content

Commit

Permalink
Remote: gRPC load balancing. (Part 5)
Browse files Browse the repository at this point in the history
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
  • Loading branch information
Googler authored and philwo committed Mar 15, 2021
1 parent b529f8d commit a6293b3
Show file tree
Hide file tree
Showing 19 changed files with 446 additions and 364 deletions.
6 changes: 3 additions & 3 deletions src/main/java/com/google/devtools/build/lib/remote/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand All @@ -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",
Expand Down Expand Up @@ -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",
],
Expand Down
Original file line number Diff line number Diff line change
@@ -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;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -131,4 +148,4 @@ public boolean release() {
public boolean release(int decrement) {
return referenceCounted.release(decrement);
}
}
}

This file was deleted.

Loading

0 comments on commit a6293b3

Please sign in to comment.