Skip to content

Commit a6293b3

Browse files
Googlerphilwo
Googler
authored and
philwo
committed
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
1 parent b529f8d commit a6293b3

19 files changed

+446
-364
lines changed

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,6 @@ java_library(
4242
":ExecutionStatusException",
4343
":ReferenceCountedChannel",
4444
":Retrier",
45-
"//src/main/java/com/google/devtools/build/lib:build-request-options",
4645
"//src/main/java/com/google/devtools/build/lib:runtime",
4746
"//src/main/java/com/google/devtools/build/lib/actions",
4847
"//src/main/java/com/google/devtools/build/lib/actions:artifacts",
@@ -75,6 +74,7 @@ java_library(
7574
"//src/main/java/com/google/devtools/build/lib/remote/common",
7675
"//src/main/java/com/google/devtools/build/lib/remote/disk",
7776
"//src/main/java/com/google/devtools/build/lib/remote/downloader",
77+
"//src/main/java/com/google/devtools/build/lib/remote/grpc",
7878
"//src/main/java/com/google/devtools/build/lib/remote/http",
7979
"//src/main/java/com/google/devtools/build/lib/remote/logging",
8080
"//src/main/java/com/google/devtools/build/lib/remote/merkletree",
@@ -98,6 +98,7 @@ java_library(
9898
"//third_party:guava",
9999
"//third_party:jsr305",
100100
"//third_party:netty",
101+
"//third_party:rxjava3",
101102
"//third_party/grpc:grpc-jar",
102103
"//third_party/protobuf:protobuf_java",
103104
"//third_party/protobuf:protobuf_java_util",
@@ -128,10 +129,9 @@ java_library(
128129
name = "ReferenceCountedChannel",
129130
srcs = [
130131
"ReferenceCountedChannel.java",
131-
"ReferenceCountedChannelPool.java",
132132
],
133133
deps = [
134-
"//third_party:guava",
134+
"//src/main/java/com/google/devtools/build/lib/remote/grpc",
135135
"//third_party:netty",
136136
"//third_party/grpc:grpc-jar",
137137
],

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

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
// Copyright 2021 The Bazel Authors. All rights reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
package com.google.devtools.build.lib.remote;
15+
16+
import com.google.devtools.build.lib.authandtls.AuthAndTLSOptions;
17+
import com.google.devtools.build.lib.remote.grpc.ChannelConnectionFactory;
18+
import io.grpc.ClientInterceptor;
19+
import io.reactivex.rxjava3.core.Single;
20+
import java.util.List;
21+
22+
/**
23+
* A {@link ChannelConnectionFactory} which creates {@link ChannelConnection} using {@link
24+
* ChannelFactory}.
25+
*/
26+
public class GoogleChannelConnectionFactory implements ChannelConnectionFactory {
27+
private final ChannelFactory channelFactory;
28+
private final String target;
29+
private final String proxy;
30+
private final AuthAndTLSOptions options;
31+
private final List<ClientInterceptor> interceptors;
32+
private final int maxConcurrency;
33+
34+
public GoogleChannelConnectionFactory(
35+
ChannelFactory channelFactory,
36+
String target,
37+
String proxy,
38+
AuthAndTLSOptions options,
39+
List<ClientInterceptor> interceptors,
40+
int maxConcurrency) {
41+
this.channelFactory = channelFactory;
42+
this.target = target;
43+
this.proxy = proxy;
44+
this.options = options;
45+
this.interceptors = interceptors;
46+
this.maxConcurrency = maxConcurrency;
47+
}
48+
49+
@Override
50+
public Single<ChannelConnection> create() {
51+
return Single.fromCallable(
52+
() ->
53+
new ChannelConnection(channelFactory.newChannel(target, proxy, options, interceptors)));
54+
}
55+
56+
@Override
57+
public int maxConcurrency() {
58+
return maxConcurrency;
59+
}
60+
}

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

Lines changed: 71 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -13,84 +13,101 @@
1313
// limitations under the License.
1414
package com.google.devtools.build.lib.remote;
1515

16+
import com.google.devtools.build.lib.remote.grpc.ChannelConnectionFactory;
17+
import com.google.devtools.build.lib.remote.grpc.ChannelConnectionFactory.ChannelConnection;
18+
import com.google.devtools.build.lib.remote.grpc.DynamicConnectionPool;
19+
import com.google.devtools.build.lib.remote.grpc.SharedConnectionFactory.SharedConnection;
1620
import io.grpc.CallOptions;
21+
import io.grpc.Channel;
1722
import io.grpc.ClientCall;
18-
import io.grpc.ManagedChannel;
23+
import io.grpc.ForwardingClientCall;
24+
import io.grpc.ForwardingClientCallListener;
25+
import io.grpc.Metadata;
1926
import io.grpc.MethodDescriptor;
27+
import io.grpc.Status;
2028
import io.netty.util.AbstractReferenceCounted;
2129
import io.netty.util.ReferenceCounted;
22-
import java.util.concurrent.TimeUnit;
30+
import java.io.IOException;
2331

2432
/**
25-
* A wrapper around a {@link io.grpc.ManagedChannel} exposing a reference count. When instantiated
26-
* the reference count is 1. {@link ManagedChannel#shutdown()} will be called on the wrapped channel
27-
* when the reference count reaches 0.
33+
* A wrapper around a {@link DynamicConnectionPool} exposing {@link Channel} and a reference count.
34+
* When instantiated the reference count is 1. {@link DynamicConnectionPool#close()} will be called
35+
* on the wrapped channel when the reference count reaches 0.
2836
*
2937
* <p>See {@link ReferenceCounted} for more information about reference counting.
3038
*/
31-
public class ReferenceCountedChannel extends ManagedChannel implements ReferenceCounted {
32-
33-
private final ManagedChannel channel;
34-
private final AbstractReferenceCounted referenceCounted;
35-
36-
public ReferenceCountedChannel(ManagedChannel channel) {
37-
this(
38-
channel,
39-
new AbstractReferenceCounted() {
40-
@Override
41-
protected void deallocate() {
42-
channel.shutdown();
39+
public class ReferenceCountedChannel extends Channel implements ReferenceCounted {
40+
private final DynamicConnectionPool dynamicConnectionPool;
41+
private final AbstractReferenceCounted referenceCounted =
42+
new AbstractReferenceCounted() {
43+
@Override
44+
protected void deallocate() {
45+
try {
46+
dynamicConnectionPool.close();
47+
} catch (IOException e) {
48+
throw new AssertionError(e.getMessage(), e);
4349
}
50+
}
4451

45-
@Override
46-
public ReferenceCounted touch(Object o) {
47-
return this;
48-
}
49-
});
50-
}
51-
52-
protected ReferenceCountedChannel(
53-
ManagedChannel channel, AbstractReferenceCounted referenceCounted) {
54-
this.channel = channel;
55-
this.referenceCounted = referenceCounted;
56-
}
52+
@Override
53+
public ReferenceCounted touch(Object o) {
54+
return this;
55+
}
56+
};
5757

58-
@Override
59-
public ManagedChannel shutdown() {
60-
throw new UnsupportedOperationException("Don't call shutdown() directly, but use release() "
61-
+ "instead.");
58+
public ReferenceCountedChannel(ChannelConnectionFactory connectionFactory) {
59+
this.dynamicConnectionPool =
60+
new DynamicConnectionPool(connectionFactory, connectionFactory.maxConcurrency());
6261
}
6362

64-
@Override
6563
public boolean isShutdown() {
66-
return channel.isShutdown();
67-
}
68-
69-
@Override
70-
public boolean isTerminated() {
71-
return channel.isTerminated();
72-
}
73-
74-
@Override
75-
public ManagedChannel shutdownNow() {
76-
throw new UnsupportedOperationException("Don't call shutdownNow() directly, but use release() "
77-
+ "instead.");
78-
}
79-
80-
@Override
81-
public boolean awaitTermination(long timeout, TimeUnit timeUnit) throws InterruptedException {
82-
return channel.awaitTermination(timeout, timeUnit);
64+
return dynamicConnectionPool.isClosed();
65+
}
66+
67+
/** A {@link ClientCall} which call {@link SharedConnection#close()} after the RPC is closed. */
68+
static class ConnectionCleanupCall<ReqT, RespT>
69+
extends ForwardingClientCall.SimpleForwardingClientCall<ReqT, RespT> {
70+
private final SharedConnection connection;
71+
72+
protected ConnectionCleanupCall(ClientCall<ReqT, RespT> delegate, SharedConnection connection) {
73+
super(delegate);
74+
this.connection = connection;
75+
}
76+
77+
@Override
78+
public void start(Listener<RespT> responseListener, Metadata headers) {
79+
super.start(
80+
new ForwardingClientCallListener.SimpleForwardingClientCallListener<RespT>(
81+
responseListener) {
82+
@Override
83+
public void onClose(Status status, Metadata trailers) {
84+
super.onClose(status, trailers);
85+
86+
try {
87+
connection.close();
88+
} catch (IOException e) {
89+
throw new AssertionError(e.getMessage(), e);
90+
}
91+
}
92+
},
93+
headers);
94+
}
8395
}
8496

8597
@Override
8698
public <RequestT, ResponseT> ClientCall<RequestT, ResponseT> newCall(
8799
MethodDescriptor<RequestT, ResponseT> methodDescriptor, CallOptions callOptions) {
88-
return channel.<RequestT, ResponseT>newCall(methodDescriptor, callOptions);
100+
SharedConnection sharedConnection = dynamicConnectionPool.create().blockingGet();
101+
ChannelConnection connection = (ChannelConnection) sharedConnection.getUnderlyingConnection();
102+
return new ConnectionCleanupCall<>(
103+
connection.getChannel().newCall(methodDescriptor, callOptions), sharedConnection);
89104
}
90105

91106
@Override
92107
public String authority() {
93-
return channel.authority();
108+
SharedConnection sharedConnection = dynamicConnectionPool.create().blockingGet();
109+
ChannelConnection connection = (ChannelConnection) sharedConnection.getUnderlyingConnection();
110+
return connection.getChannel().authority();
94111
}
95112

96113
@Override
@@ -131,4 +148,4 @@ public boolean release() {
131148
public boolean release(int decrement) {
132149
return referenceCounted.release(decrement);
133150
}
134-
}
151+
}

‎src/main/java/com/google/devtools/build/lib/remote/ReferenceCountedChannelPool.java‎

Lines changed: 0 additions & 120 deletions
This file was deleted.

0 commit comments

Comments
 (0)