Skip to content

Commit bd647b0

Browse files
author
Bharath Kumarasubramanian
committed
Client bindings and request flows for controller APIs createStore, getStoresInCluster and QueryJobStatus
1 parent 14d05d4 commit bd647b0

File tree

8 files changed

+657
-0
lines changed

8 files changed

+657
-0
lines changed

Diff for: ‎internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/ControllerClient.java

+42
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,9 @@
7878
import com.linkedin.venice.exceptions.ErrorType;
7979
import com.linkedin.venice.exceptions.VeniceException;
8080
import com.linkedin.venice.exceptions.VeniceHttpException;
81+
import com.linkedin.venice.grpc.ControllerGrpcTransport;
82+
import com.linkedin.venice.grpc.GrpcControllerRoute;
83+
import com.linkedin.venice.grpc.GrpcConverters;
8184
import com.linkedin.venice.helix.VeniceJsonSerializer;
8285
import com.linkedin.venice.meta.VeniceUserStoreType;
8386
import com.linkedin.venice.meta.Version;
@@ -98,6 +101,7 @@
98101
import java.util.List;
99102
import java.util.Optional;
100103
import java.util.Set;
104+
import java.util.concurrent.CompletionStage;
101105
import java.util.concurrent.ExecutionException;
102106
import java.util.concurrent.TimeoutException;
103107
import java.util.function.Function;
@@ -121,6 +125,8 @@ public class ControllerClient implements Closeable {
121125
private String leaderControllerUrl;
122126
private final List<String> controllerDiscoveryUrls;
123127

128+
private final boolean useGrpc;
129+
124130
public ControllerClient(String clusterName, String discoveryUrls) {
125131
this(clusterName, discoveryUrls, Optional.empty());
126132
}
@@ -140,6 +146,7 @@ public ControllerClient(String clusterName, String discoveryUrls, Optional<SSLFa
140146
if (this.controllerDiscoveryUrls.isEmpty()) {
141147
throw new VeniceException("Controller discovery url list is empty");
142148
}
149+
this.useGrpc = false;
143150
}
144151

145152
/**
@@ -1421,6 +1428,41 @@ private <T extends ControllerResponse> T request(
14211428
int timeoutMs,
14221429
int maxAttempts,
14231430
byte[] data) {
1431+
if (useGrpc) {
1432+
return fireRequestUsingGrpc(route, params, responseType, timeoutMs, maxAttempts, data);
1433+
} else {
1434+
return fireRequestUsingHttp(route, params, responseType, timeoutMs, maxAttempts, data);
1435+
}
1436+
}
1437+
1438+
private <T extends ControllerResponse> T fireRequestUsingGrpc(
1439+
ControllerRoute route,
1440+
QueryParams params,
1441+
Class<T> responseType,
1442+
int timeoutMs,
1443+
int maxAttempts,
1444+
byte[] data) {
1445+
GrpcControllerRoute grpcControllerRoute = GrpcConverters.mapControllerRouteToGrpcControllerRoute(route);
1446+
try (ControllerGrpcTransport transport = new ControllerGrpcTransport(sslFactory)) {
1447+
try {
1448+
CompletionStage<T> responseFuture =
1449+
transport.request(getLeaderControllerUrl(), params, responseType, grpcControllerRoute);
1450+
return responseFuture.toCompletableFuture().join();
1451+
} catch (Exception e) {
1452+
throw new VeniceException("Encountered error when getting response from server", e);
1453+
}
1454+
} catch (Exception e) {
1455+
throw new VeniceException("Encountered error when fetching response from grpc server", e);
1456+
}
1457+
}
1458+
1459+
private <T extends ControllerResponse> T fireRequestUsingHttp(
1460+
ControllerRoute route,
1461+
QueryParams params,
1462+
Class<T> responseType,
1463+
int timeoutMs,
1464+
int maxAttempts,
1465+
byte[] data) {
14241466
Exception lastException = null;
14251467
boolean logErrorMessage = true;
14261468
try (ControllerTransport transport = new ControllerTransport(sslFactory)) {

Diff for: ‎internal/venice-common/src/main/java/com/linkedin/venice/grpc/ControllerGrpcTransport.java

+157
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
package com.linkedin.venice.grpc;
2+
3+
import com.google.common.annotations.VisibleForTesting;
4+
import com.google.common.base.Preconditions;
5+
import com.linkedin.venice.client.exceptions.VeniceClientException;
6+
import com.linkedin.venice.controllerapi.QueryParams;
7+
import com.linkedin.venice.exceptions.VeniceException;
8+
import com.linkedin.venice.protocols.CreateStoreGrpcRequest;
9+
import com.linkedin.venice.protocols.GetStoresInClusterGrpcRequest;
10+
import com.linkedin.venice.protocols.QueryJobStatusGrpcRequest;
11+
import com.linkedin.venice.protocols.VeniceControllerGrpcServiceGrpc;
12+
import com.linkedin.venice.security.SSLFactory;
13+
import com.linkedin.venice.utils.concurrent.VeniceConcurrentHashMap;
14+
import io.grpc.ChannelCredentials;
15+
import io.grpc.Grpc;
16+
import io.grpc.InsecureChannelCredentials;
17+
import io.grpc.ManagedChannel;
18+
import io.grpc.TlsChannelCredentials;
19+
import io.grpc.stub.StreamObserver;
20+
import java.io.IOException;
21+
import java.util.Map;
22+
import java.util.Optional;
23+
import java.util.concurrent.CompletableFuture;
24+
import java.util.concurrent.CompletionStage;
25+
26+
27+
public class ControllerGrpcTransport implements AutoCloseable {
28+
private static final int PORT = 1234;
29+
private static final String GRPC_ADDRESS_FORMAT = "%s:%s";
30+
private final VeniceConcurrentHashMap<String, ManagedChannel> serverGrpcChannels;
31+
private final ChannelCredentials channelCredentials;
32+
private final VeniceConcurrentHashMap<ManagedChannel, VeniceControllerGrpcServiceGrpc.VeniceControllerGrpcServiceStub> stubCache;
33+
34+
public ControllerGrpcTransport(Optional<SSLFactory> sslFactory) {
35+
this.stubCache = new VeniceConcurrentHashMap<>();
36+
this.serverGrpcChannels = new VeniceConcurrentHashMap<>();
37+
this.channelCredentials = buildChannelCredentials(sslFactory);
38+
}
39+
40+
public <ResT> CompletionStage<ResT> request(
41+
String serverUrl,
42+
QueryParams params,
43+
Class<ResT> responseType,
44+
GrpcControllerRoute route) {
45+
46+
VeniceControllerGrpcServiceGrpc.VeniceControllerGrpcServiceStub stub = getOrCreateStub(serverUrl);
47+
CompletableFuture<ResT> valueFuture = new CompletableFuture<>();
48+
49+
if (GrpcControllerRoute.CREATE_STORE.equals(route)) {
50+
stub.createStore(
51+
(CreateStoreGrpcRequest) GrpcConverters.getRequestConverter(route.getRequestType()).convert(params),
52+
buildStreamObserver(valueFuture, responseType, route));
53+
} else if (GrpcControllerRoute.GET_STORES_IN_CLUSTER.equals(route)) {
54+
stub.getStoresInCluster(
55+
(GetStoresInClusterGrpcRequest) GrpcConverters.getRequestConverter(route.getRequestType()).convert(params),
56+
buildStreamObserver(valueFuture, responseType, route));
57+
} else if (GrpcControllerRoute.QUERY_JOB_STATUS.equals(route)) {
58+
stub.getJobStatus(
59+
(QueryJobStatusGrpcRequest) GrpcConverters.getRequestConverter(route.getRequestType()).convert(params),
60+
buildStreamObserver(valueFuture, responseType, route));
61+
} else {
62+
throw new VeniceException("Unknown gRPC route; Failing the request");
63+
}
64+
65+
return valueFuture;
66+
}
67+
68+
@VisibleForTesting
69+
<T, ResT> ControllerGrpcObserver<T, ResT> buildStreamObserver(
70+
CompletableFuture<ResT> future,
71+
Class<ResT> httpResponseType,
72+
GrpcControllerRoute route) {
73+
return new ControllerGrpcObserver<>(future, httpResponseType, route);
74+
}
75+
76+
@Override
77+
public void close() throws IOException {
78+
for (Map.Entry<String, ManagedChannel> entry: serverGrpcChannels.entrySet()) {
79+
entry.getValue().shutdown();
80+
}
81+
}
82+
83+
@VisibleForTesting
84+
ChannelCredentials buildChannelCredentials(Optional<SSLFactory> sslFactory) {
85+
SSLFactory factory = sslFactory.orElse(null);
86+
87+
if (factory == null) {
88+
return InsecureChannelCredentials.create();
89+
}
90+
91+
try {
92+
TlsChannelCredentials.Builder tlsBuilder = TlsChannelCredentials.newBuilder()
93+
.keyManager(GrpcUtils.getKeyManagers(factory))
94+
.trustManager(GrpcUtils.getTrustManagers(factory));
95+
return tlsBuilder.build();
96+
} catch (Exception e) {
97+
throw new VeniceClientException(
98+
"Failed to initialize SSL channel credentials for Venice gRPC Transport Client",
99+
e);
100+
}
101+
}
102+
103+
VeniceControllerGrpcServiceGrpc.VeniceControllerGrpcServiceStub getOrCreateStub(String serverAddress) {
104+
String grpcAddress = getGrpcAddressFromServerAddress(serverAddress);
105+
106+
ManagedChannel channel = serverGrpcChannels
107+
.computeIfAbsent(serverAddress, k -> Grpc.newChannelBuilder(grpcAddress, channelCredentials).build());
108+
109+
return stubCache.computeIfAbsent(channel, VeniceControllerGrpcServiceGrpc::newStub);
110+
}
111+
112+
@VisibleForTesting
113+
String getGrpcAddressFromServerAddress(String serverAddress) {
114+
String[] serverAddressParts = serverAddress.split(":");
115+
Preconditions.checkState(serverAddressParts.length == 2, "Invalid server address");
116+
117+
return String.format(GRPC_ADDRESS_FORMAT, serverAddressParts[0], PORT);
118+
}
119+
120+
static class ControllerGrpcObserver<ResT, HttpResT> implements StreamObserver<ResT> {
121+
private final CompletableFuture<HttpResT> responseFuture;
122+
private final Class<HttpResT> httpResponseType;
123+
124+
private final GrpcControllerRoute route;
125+
126+
public ControllerGrpcObserver(
127+
CompletableFuture<HttpResT> future,
128+
Class<HttpResT> httpResponseType,
129+
GrpcControllerRoute route) {
130+
this.httpResponseType = httpResponseType;
131+
this.responseFuture = future;
132+
this.route = route;
133+
}
134+
135+
@Override
136+
public void onNext(ResT value) {
137+
if (!responseFuture.isDone()) {
138+
139+
@SuppressWarnings("Unchecked")
140+
HttpResT result = ((GrpcToHttpResponseConverter<ResT, HttpResT>) GrpcConverters
141+
.getResponseConverter(route.getResponseType(), httpResponseType)).convert(value);
142+
143+
responseFuture.complete(result);
144+
}
145+
}
146+
147+
@Override
148+
public void onError(Throwable t) {
149+
responseFuture.completeExceptionally(t);
150+
}
151+
152+
@Override
153+
public void onCompleted() {
154+
// do nothing
155+
}
156+
}
157+
}

Diff for: ‎internal/venice-common/src/main/java/com/linkedin/venice/grpc/GrpcControllerRoute.java

+32
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
package com.linkedin.venice.grpc;
2+
3+
import com.google.protobuf.GeneratedMessageV3;
4+
import com.linkedin.venice.protocols.CreateStoreGrpcRequest;
5+
import com.linkedin.venice.protocols.CreateStoreGrpcResponse;
6+
import com.linkedin.venice.protocols.GetStoresInClusterGrpcRequest;
7+
import com.linkedin.venice.protocols.GetStoresInClusterGrpcResponse;
8+
import com.linkedin.venice.protocols.QueryJobStatusGrpcRequest;
9+
import com.linkedin.venice.protocols.QueryJobStatusGrpcResponse;
10+
11+
12+
public enum GrpcControllerRoute {
13+
CREATE_STORE(CreateStoreGrpcRequest.class, CreateStoreGrpcResponse.class),
14+
GET_STORES_IN_CLUSTER(GetStoresInClusterGrpcRequest.class, GetStoresInClusterGrpcResponse.class),
15+
QUERY_JOB_STATUS(QueryJobStatusGrpcRequest.class, QueryJobStatusGrpcResponse.class);
16+
17+
private final Class<? extends GeneratedMessageV3> requestType;
18+
private final Class<? extends GeneratedMessageV3> responseType;
19+
20+
GrpcControllerRoute(Class<? extends GeneratedMessageV3> reqT, Class<? extends GeneratedMessageV3> resT) {
21+
this.requestType = reqT;
22+
this.responseType = resT;
23+
}
24+
25+
public Class<? extends GeneratedMessageV3> getRequestType() {
26+
return this.requestType;
27+
}
28+
29+
public Class<? extends GeneratedMessageV3> getResponseType() {
30+
return this.responseType;
31+
}
32+
}

0 commit comments

Comments
 (0)