Skip to content

Commit da1b9e5

Browse files
chore: introduce PrimingKey to encapsulate PingAndWarm request & metadata
1 parent 3389325 commit da1b9e5

File tree

3 files changed

+95
-44
lines changed

3 files changed

+95
-44
lines changed

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

Lines changed: 10 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import com.google.bigtable.v2.PingAndWarmRequest;
2121
import com.google.bigtable.v2.PingAndWarmResponse;
2222
import com.google.cloud.bigtable.examples.proxy.core.CallLabels;
23+
import com.google.cloud.bigtable.examples.proxy.core.CallLabels.PrimingKey;
2324
import com.google.cloud.bigtable.examples.proxy.metrics.Metrics;
2425
import com.google.cloud.bigtable.examples.proxy.metrics.Tracer;
2526
import com.google.common.util.concurrent.ListenableFuture;
@@ -37,8 +38,6 @@
3738
import io.grpc.MethodDescriptor;
3839
import io.grpc.Status;
3940
import io.grpc.StatusRuntimeException;
40-
import java.net.URLEncoder;
41-
import java.nio.charset.StandardCharsets;
4241
import java.util.List;
4342
import java.util.Optional;
4443
import java.util.concurrent.ExecutionException;
@@ -107,18 +106,18 @@ private void warmQuietly() {
107106
}
108107

109108
private void warm() {
110-
List<PingAndWarmRequest> requests = resourceCollector.getRequests();
111-
if (requests.isEmpty()) {
109+
List<PrimingKey> primingKeys = resourceCollector.getPrimingKeys();
110+
if (primingKeys.isEmpty()) {
112111
return;
113112
}
114113

115114
List<ListenableFuture<PingAndWarmResponse>> futures =
116-
requests.stream().map(this::sendPingAndWarm).collect(Collectors.toList());
115+
primingKeys.stream().map(this::sendPingAndWarm).collect(Collectors.toList());
117116

118117
int successCount = 0;
119118
int failures = 0;
120119
for (ListenableFuture<PingAndWarmResponse> future : futures) {
121-
PingAndWarmRequest request = requests.get(successCount + failures);
120+
PrimingKey request = primingKeys.get(successCount + failures);
122121
try {
123122
future.get();
124123
successCount++;
@@ -151,13 +150,11 @@ private void warm() {
151150
}
152151
}
153152

154-
private ListenableFuture<PingAndWarmResponse> sendPingAndWarm(PingAndWarmRequest request) {
155-
CallLabels callLabels =
156-
CallLabels.create(
157-
BigtableGrpc.getPingAndWarmMethod(),
158-
Optional.of("bigtableproxy"),
159-
Optional.of(request.getName()),
160-
Optional.of(request.getAppProfileId()));
153+
private ListenableFuture<PingAndWarmResponse> sendPingAndWarm(PrimingKey primingKey) {
154+
Metadata metadata = primingKey.composeMetadata();
155+
PingAndWarmRequest request = primingKey.composeProto();
156+
157+
CallLabels callLabels = CallLabels.create(BigtableGrpc.getPingAndWarmMethod(), metadata);
161158
Tracer tracer = new Tracer(metrics, callLabels);
162159

163160
CallOptions callOptions =
@@ -169,14 +166,6 @@ private ListenableFuture<PingAndWarmResponse> sendPingAndWarm(PingAndWarmRequest
169166
ClientCall<PingAndWarmRequest, PingAndWarmResponse> call =
170167
inner.newCall(BigtableGrpc.getPingAndWarmMethod(), callOptions);
171168

172-
Metadata metadata = new Metadata();
173-
metadata.put(
174-
CallLabels.REQUEST_PARAMS,
175-
String.format(
176-
"name=%s&app_profile_id=%s",
177-
URLEncoder.encode(request.getName(), StandardCharsets.UTF_8),
178-
URLEncoder.encode(request.getAppProfileId(), StandardCharsets.UTF_8)));
179-
180169
SettableFuture<PingAndWarmResponse> f = SettableFuture.create();
181170
call.start(
182171
new Listener<>() {

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

Lines changed: 15 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -16,44 +16,36 @@
1616

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

19-
import com.google.bigtable.v2.PingAndWarmRequest;
2019
import com.google.cloud.bigtable.examples.proxy.core.CallLabels;
20+
import com.google.cloud.bigtable.examples.proxy.core.CallLabels.ParsingException;
21+
import com.google.cloud.bigtable.examples.proxy.core.CallLabels.PrimingKey;
2122
import com.google.common.cache.Cache;
2223
import com.google.common.cache.CacheBuilder;
2324
import com.google.common.collect.ImmutableList;
2425
import java.time.Duration;
2526
import java.util.List;
27+
import org.slf4j.Logger;
28+
import org.slf4j.LoggerFactory;
2629

2730
public class ResourceCollector {
28-
private final Cache<PingAndWarmRequest, Boolean> warmingRequests =
31+
private static final Logger LOG = LoggerFactory.getLogger(ResourceCollector.class);
32+
33+
private final Cache<PrimingKey, Boolean> primingKeys =
2934
CacheBuilder.newBuilder().expireAfterWrite(Duration.ofHours(1)).maximumSize(100).build();
3035

3136
public void collect(CallLabels labels) {
32-
String[] splits = labels.getResourceName().orElse("").split("/", 5);
33-
if (splits.length <= 4) {
34-
return;
35-
}
36-
if (!"projects".equals(splits[0])) {
37-
return;
37+
try {
38+
PrimingKey.from(labels).ifPresent(k -> primingKeys.put(k, true));
39+
} catch (ParsingException e) {
40+
LOG.atWarn().log("Failed to collect priming request for {}", labels, e);
3841
}
39-
if (!"instances".equals(splits[2])) {
40-
return;
41-
}
42-
String appProfile = labels.getAppProfileId().orElse("");
43-
44-
PingAndWarmRequest req =
45-
PingAndWarmRequest.newBuilder()
46-
.setName("projects/" + splits[1] + "/instances/" + splits[3])
47-
.setAppProfileId(appProfile)
48-
.build();
49-
warmingRequests.put(req, true);
5042
}
5143

52-
public List<PingAndWarmRequest> getRequests() {
53-
return ImmutableList.copyOf(warmingRequests.asMap().keySet());
44+
public List<PrimingKey> getPrimingKeys() {
45+
return ImmutableList.copyOf(primingKeys.asMap().keySet());
5446
}
5547

56-
public void evict(PingAndWarmRequest request) {
57-
warmingRequests.invalidate(request);
48+
public void evict(PrimingKey request) {
49+
primingKeys.invalidate(request);
5850
}
5951
}

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

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,17 @@
1717
package com.google.cloud.bigtable.examples.proxy.core;
1818

1919
import com.google.auto.value.AutoValue;
20+
import com.google.bigtable.v2.PingAndWarmRequest;
21+
import com.google.bigtable.v2.PingAndWarmRequest.Builder;
22+
import com.google.common.collect.ImmutableMap;
2023
import io.grpc.Metadata;
2124
import io.grpc.Metadata.Key;
2225
import io.grpc.MethodDescriptor;
2326
import java.net.URLDecoder;
27+
import java.net.URLEncoder;
2428
import java.nio.charset.StandardCharsets;
29+
import java.util.Map;
30+
import java.util.Map.Entry;
2531
import java.util.Optional;
2632

2733
/**
@@ -155,4 +161,68 @@ private static Optional<String> extractAppProfileId(String[] encodedKvPairs) {
155161
private static String percentDecode(String s) {
156162
return URLDecoder.decode(s, StandardCharsets.UTF_8);
157163
}
164+
165+
@AutoValue
166+
public abstract static class PrimingKey {
167+
abstract Map<String, String> getMetadata();
168+
169+
abstract String getName();
170+
171+
abstract Optional<String> getAppProfileId();
172+
173+
public static Optional<PrimingKey> from(CallLabels labels) throws ParsingException {
174+
Optional<String> resourceName = labels.getResourceName();
175+
if (resourceName.isEmpty()) {
176+
return Optional.empty();
177+
}
178+
String[] resourceNameParts = resourceName.get().split("/", 5);
179+
if (resourceNameParts.length < 4
180+
|| !resourceNameParts[0].equals("projects")
181+
|| !resourceNameParts[2].equals("instances")) {
182+
return Optional.empty();
183+
}
184+
String instanceName =
185+
"projects/" + resourceNameParts[1] + "/instances/" + resourceNameParts[3];
186+
StringBuilder reqParams =
187+
new StringBuilder()
188+
.append("name=")
189+
.append(URLEncoder.encode(instanceName, StandardCharsets.UTF_8));
190+
191+
Optional<String> appProfileId = labels.getAppProfileId();
192+
appProfileId.ifPresent(val -> reqParams.append("&app_profile_id=").append(val));
193+
194+
ImmutableMap.Builder<String, String> md = ImmutableMap.builder();
195+
md.put(REQUEST_PARAMS.name(), reqParams.toString());
196+
197+
labels.getApiClient().ifPresent(c -> md.put(API_CLIENT.name(), c));
198+
199+
return Optional.of(
200+
new AutoValue_CallLabels_PrimingKey(md.build(), instanceName, appProfileId));
201+
}
202+
203+
public Metadata composeMetadata() {
204+
Metadata md = new Metadata();
205+
for (Entry<String, String> e : getMetadata().entrySet()) {
206+
md.put(Key.of(e.getKey(), Metadata.ASCII_STRING_MARSHALLER), e.getValue());
207+
}
208+
return md;
209+
}
210+
211+
public PingAndWarmRequest composeProto() {
212+
Builder builder = PingAndWarmRequest.newBuilder().setName(getName());
213+
getAppProfileId().ifPresent(builder::setAppProfileId);
214+
return builder.build();
215+
}
216+
}
217+
218+
public static class ParsingException extends Exception {
219+
220+
public ParsingException(String message) {
221+
super(message);
222+
}
223+
224+
public ParsingException(String message, Throwable cause) {
225+
super(message, cause);
226+
}
227+
}
158228
}

0 commit comments

Comments
 (0)