Skip to content

Commit 834bf5b

Browse files
feat: add basic proxy
stack-info: PR: #9689, branch: igorbernstein2/stack/2
1 parent 67b594d commit 834bf5b

File tree

12 files changed

+1100
-1
lines changed

12 files changed

+1100
-1
lines changed

bigtable/bigtable-proxy/pom.xml

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,47 @@
4444
</dependencyManagement>
4545

4646
<dependencies>
47+
<!-- gRPC -->
48+
<dependency>
49+
<groupId>io.grpc</groupId>
50+
<artifactId>grpc-api</artifactId>
51+
</dependency>
52+
<dependency>
53+
<groupId>io.grpc</groupId>
54+
<artifactId>grpc-core</artifactId>
55+
</dependency>
56+
<dependency>
57+
<groupId>io.grpc</groupId>
58+
<artifactId>grpc-netty-shaded</artifactId>
59+
</dependency>
60+
61+
62+
<!-- service defs -->
63+
<dependency>
64+
<groupId>com.google.api.grpc</groupId>
65+
<artifactId>grpc-google-cloud-bigtable-v2</artifactId>
66+
</dependency>
67+
<dependency>
68+
<groupId>com.google.api.grpc</groupId>
69+
<artifactId>proto-google-cloud-bigtable-v2</artifactId>
70+
</dependency>
71+
<dependency>
72+
<groupId>com.google.api.grpc</groupId>
73+
<artifactId>grpc-google-cloud-bigtable-admin-v2</artifactId>
74+
</dependency>
75+
<dependency>
76+
<groupId>com.google.api.grpc</groupId>
77+
<artifactId>proto-google-cloud-bigtable-admin-v2</artifactId>
78+
</dependency>
79+
<dependency>
80+
<groupId>com.google.api.grpc</groupId>
81+
<artifactId>grpc-google-common-protos</artifactId>
82+
</dependency>
83+
<dependency>
84+
<groupId>com.google.api.grpc</groupId>
85+
<artifactId>proto-google-common-protos</artifactId>
86+
</dependency>
87+
4788
<!-- Logging -->
4889
<dependency>
4990
<groupId>org.slf4j</groupId>
@@ -80,6 +121,11 @@
80121
</dependency>
81122

82123
<!-- Test -->
124+
<dependency>
125+
<groupId>io.grpc</groupId>
126+
<artifactId>grpc-testing</artifactId>
127+
<scope>test</scope>
128+
</dependency>
83129
<dependency>
84130
<groupId>junit</groupId>
85131
<artifactId>junit</artifactId>

bigtable/bigtable-proxy/src/main/java/com/google/cloud/bigtable/examples/proxy/Main.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
package com.google.cloud.bigtable.examples.proxy;
1818

19+
import com.google.cloud.bigtable.examples.proxy.commands.Serve;
1920
import org.slf4j.Logger;
2021
import org.slf4j.LoggerFactory;
2122
import org.slf4j.bridge.SLF4JBridgeHandler;
@@ -26,7 +27,7 @@
2627
* Main entry point for proxy commands under {@link
2728
* com.google.cloud.bigtable.examples.proxy.commands}.
2829
*/
29-
@Command(subcommands = {})
30+
@Command(subcommands = {Serve.class})
3031
public final class Main {
3132
private static final Logger LOGGER = LoggerFactory.getLogger(Main.class);
3233

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.examples.proxy.commands;
17+
18+
import com.google.auto.value.AutoValue;
19+
import com.google.common.base.Preconditions;
20+
import picocli.CommandLine.ITypeConverter;
21+
22+
@AutoValue
23+
abstract class Endpoint {
24+
abstract String getName();
25+
26+
abstract int getPort();
27+
28+
@Override
29+
public String toString() {
30+
return String.format("%s:%d", getName(), getPort());
31+
}
32+
33+
static Endpoint create(String name, int port) {
34+
return new AutoValue_Endpoint(name, port);
35+
}
36+
37+
static class ArgConverter implements ITypeConverter<Endpoint> {
38+
@Override
39+
public Endpoint convert(String s) throws Exception {
40+
int i = s.lastIndexOf(":");
41+
Preconditions.checkArgument(i > 0, "endpoint must of the form `name:port`");
42+
43+
String name = s.substring(0, i);
44+
int port = Integer.parseInt(s.substring(i + 1));
45+
return Endpoint.create(name, port);
46+
}
47+
}
48+
}
Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.examples.proxy.commands;
17+
18+
import com.google.bigtable.admin.v2.BigtableInstanceAdminGrpc;
19+
import com.google.bigtable.admin.v2.BigtableTableAdminGrpc;
20+
import com.google.bigtable.v2.BigtableGrpc;
21+
import com.google.cloud.bigtable.examples.proxy.core.ProxyHandler;
22+
import com.google.cloud.bigtable.examples.proxy.core.Registry;
23+
import com.google.common.collect.ImmutableMap;
24+
import com.google.longrunning.OperationsGrpc;
25+
import io.grpc.InsecureServerCredentials;
26+
import io.grpc.ManagedChannel;
27+
import io.grpc.ManagedChannelBuilder;
28+
import io.grpc.Server;
29+
import io.grpc.ServerCallHandler;
30+
import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder;
31+
import java.io.IOException;
32+
import java.net.InetSocketAddress;
33+
import java.util.Map;
34+
import java.util.concurrent.Callable;
35+
import java.util.concurrent.TimeUnit;
36+
import org.slf4j.Logger;
37+
import org.slf4j.LoggerFactory;
38+
import picocli.CommandLine.Command;
39+
import picocli.CommandLine.Help.Visibility;
40+
import picocli.CommandLine.Option;
41+
42+
@Command(name = "serve", mixinStandardHelpOptions = true, description = "Start the proxy server")
43+
public class Serve implements Callable<Void> {
44+
private static final Logger LOGGER = LoggerFactory.getLogger(Serve.class);
45+
46+
@Option(
47+
names = "--listen-port",
48+
required = true,
49+
description = "Local port to accept connections on")
50+
int listenPort;
51+
52+
@Option(names = "--useragent", showDefaultValue = Visibility.ALWAYS)
53+
String userAgent = "bigtable-java-proxy";
54+
55+
@Option(
56+
names = "--bigtable-data-endpoint",
57+
converter = Endpoint.ArgConverter.class,
58+
showDefaultValue = Visibility.ALWAYS)
59+
Endpoint dataEndpoint = Endpoint.create("bigtable.googleapis.com", 443);
60+
61+
@Option(
62+
names = "--bigtable-admin-endpoint",
63+
converter = Endpoint.ArgConverter.class,
64+
showDefaultValue = Visibility.ALWAYS)
65+
Endpoint adminEndpoint = Endpoint.create("bigtableadmin.googleapis.com", 443);
66+
67+
ManagedChannel adminChannel = null;
68+
ManagedChannel dataChannel = null;
69+
Server server;
70+
71+
@Override
72+
public Void call() throws Exception {
73+
start();
74+
server.awaitTermination();
75+
cleanup();
76+
return null;
77+
}
78+
79+
void start() throws IOException {
80+
if (dataChannel == null) {
81+
dataChannel =
82+
ManagedChannelBuilder.forAddress(dataEndpoint.getName(), dataEndpoint.getPort())
83+
.maxInboundMessageSize(256 * 1024 * 1024)
84+
.disableRetry()
85+
.keepAliveTime(30, TimeUnit.SECONDS)
86+
.keepAliveTimeout(10, TimeUnit.SECONDS)
87+
.build();
88+
}
89+
if (adminChannel == null) {
90+
adminChannel =
91+
ManagedChannelBuilder.forAddress(adminEndpoint.getName(), adminEndpoint.getPort())
92+
.userAgent(userAgent)
93+
.disableRetry()
94+
.build();
95+
}
96+
97+
Map<String, ServerCallHandler<byte[], byte[]>> serviceMap =
98+
ImmutableMap.of(
99+
BigtableGrpc.SERVICE_NAME,
100+
new ProxyHandler<>(dataChannel),
101+
BigtableInstanceAdminGrpc.SERVICE_NAME,
102+
new ProxyHandler<>(adminChannel),
103+
BigtableTableAdminGrpc.SERVICE_NAME,
104+
new ProxyHandler<>(adminChannel),
105+
OperationsGrpc.SERVICE_NAME,
106+
new ProxyHandler<>(adminChannel));
107+
108+
server =
109+
NettyServerBuilder.forAddress(
110+
new InetSocketAddress("localhost", listenPort), InsecureServerCredentials.create())
111+
.fallbackHandlerRegistry(new Registry(serviceMap))
112+
.maxInboundMessageSize(256 * 1024 * 1024)
113+
.build();
114+
115+
server.start();
116+
LOGGER.info("Listening on port {}", server.getPort());
117+
}
118+
119+
void cleanup() throws InterruptedException {
120+
dataChannel.shutdown();
121+
adminChannel.shutdown();
122+
}
123+
}
Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
/*
2+
* Copyright 2024 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.examples.proxy.core;
17+
18+
import com.google.common.io.ByteStreams;
19+
import io.grpc.MethodDescriptor;
20+
import java.io.ByteArrayInputStream;
21+
import java.io.IOException;
22+
import java.io.InputStream;
23+
24+
class ByteMarshaller implements MethodDescriptor.Marshaller<byte[]> {
25+
26+
@Override
27+
public byte[] parse(InputStream stream) {
28+
try {
29+
return ByteStreams.toByteArray(stream);
30+
} catch (IOException ex) {
31+
throw new RuntimeException(ex);
32+
}
33+
}
34+
35+
@Override
36+
public InputStream stream(byte[] value) {
37+
return new ByteArrayInputStream(value);
38+
}
39+
}

0 commit comments

Comments
 (0)