Visitar URL original
feat: client sends routing cookie back to server (#1888) · googleapis/java-bigtable@4c73abd · GitHub
Skip to content
This repository was archived by the owner on Aug 30, 2026. It is now read-only.

Commit 4c73abd

Browse files
mutianfgcf-owl-bot[bot]
mutianf
and
gcf-owl-bot[bot]
authored
feat: client sends routing cookie back to server (#1888)
* feat: client sends retry cookie back to server * udpate to use trailer instead of error info * updating the header name * address some comments * udpate * update tests and handling of retry cookie * address comments * address comments * add cookie to readChangeStream * also check headers and add a test * simplify code * clean up test * clean up test * update dependency * test * move MetadataSubject to a separate file * add the file * add license * address comments * close client * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
1 parent 0827252 commit 4c73abd

8 files changed

Lines changed: 924 additions & 21 deletions

File tree

‎README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ If you are using Maven without the BOM, add this to your dependencies:
5050
If you are using Gradle 5.x or later, add this to your dependencies:
5151

5252
```Groovy
53-
implementation platform('com.google.cloud:libraries-bom:26.26.0')
53+
implementation platform('com.google.cloud:libraries-bom:26.27.0')
5454
5555
implementation 'com.google.cloud:google-cloud-bigtable'
5656
```

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolder.java‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
/*
2+
* Copyright 2023 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.data.v2.stub;
17+
18+
import io.grpc.CallOptions;
19+
import io.grpc.Metadata;
20+
import java.util.HashMap;
21+
import java.util.Map;
22+
import javax.annotation.Nullable;
23+
24+
/** A cookie that holds information for retry or routing */
25+
class CookiesHolder {
26+
27+
static final CallOptions.Key<CookiesHolder> COOKIES_HOLDER_KEY =
28+
CallOptions.Key.create("bigtable-cookies");
29+
30+
/** Routing cookie key prefix. */
31+
static final String COOKIE_KEY_PREFIX = "x-goog-cbt-cookie";
32+
33+
/** A map that stores all the routing cookies. */
34+
private final Map<Metadata.Key<String>, String> cookies = new HashMap<>();
35+
36+
/** Returns CookiesHolder if presents in CallOptions, otherwise returns null. */
37+
@Nullable
38+
static CookiesHolder fromCallOptions(CallOptions options) {
39+
// CookiesHolder should be added by CookiesServerStreamingCallable and
40+
// CookiesUnaryCallable for most methods. However, methods like PingAndWarm
41+
// doesn't support routing cookie, in which case this will return null.
42+
return options.getOption(COOKIES_HOLDER_KEY);
43+
}
44+
45+
/** Add all the routing cookies to headers if any. */
46+
Metadata injectCookiesInRequestHeaders(Metadata headers) {
47+
for (Metadata.Key<String> key : cookies.keySet()) {
48+
headers.put(key, cookies.get(key));
49+
}
50+
return headers;
51+
}
52+
53+
/**
54+
* Iterate through all the keys in initial or trailing metadata, and add all the keys that match
55+
* COOKIE_KEY_PREFIX to cookies. Values in trailers will override the value set in initial
56+
* metadata for the same keys.
57+
*/
58+
void extractCookiesFromMetadata(@Nullable Metadata trailers) {
59+
if (trailers == null) {
60+
return;
61+
}
62+
for (String key : trailers.keys()) {
63+
if (key.startsWith(COOKIE_KEY_PREFIX)) {
64+
Metadata.Key<String> metadataKey = Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER);
65+
String value = trailers.get(metadataKey);
66+
cookies.put(metadataKey, value);
67+
}
68+
}
69+
}
70+
}

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookiesInterceptor.java‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
/*
2+
* Copyright 2023 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.data.v2.stub;
17+
18+
import io.grpc.CallOptions;
19+
import io.grpc.Channel;
20+
import io.grpc.ClientCall;
21+
import io.grpc.ClientInterceptor;
22+
import io.grpc.ForwardingClientCall;
23+
import io.grpc.ForwardingClientCallListener;
24+
import io.grpc.Metadata;
25+
import io.grpc.MethodDescriptor;
26+
import io.grpc.Status;
27+
import java.util.logging.Level;
28+
import java.util.logging.Logger;
29+
30+
/**
31+
* A cookie interceptor that checks the cookie value from returned trailer, updates the cookie
32+
* holder, and inject it in the header of the next request.
33+
*/
34+
class CookiesInterceptor implements ClientInterceptor {
35+
36+
private static final Logger LOG = Logger.getLogger(CookiesInterceptor.class.getName());
37+
38+
@Override
39+
public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
40+
MethodDescriptor<ReqT, RespT> methodDescriptor, CallOptions callOptions, Channel channel) {
41+
return new ForwardingClientCall.SimpleForwardingClientCall<ReqT, RespT>(
42+
channel.newCall(methodDescriptor, callOptions)) {
43+
@Override
44+
public void start(Listener<RespT> responseListener, Metadata headers) {
45+
// Gets the CookiesHolder added from CookiesServerStreamingCallable and
46+
// CookiesUnaryCallable.
47+
// Add CookiesHolder content to request headers if there's any.
48+
try {
49+
CookiesHolder cookie = CookiesHolder.fromCallOptions(callOptions);
50+
if (cookie != null) {
51+
headers = cookie.injectCookiesInRequestHeaders(headers);
52+
responseListener = new UpdateCookieListener<>(responseListener, cookie);
53+
}
54+
} catch (Throwable e) {
55+
LOG.warning("Failed to inject cookie to request headers: " + e);
56+
} finally {
57+
super.start(responseListener, headers);
58+
}
59+
}
60+
};
61+
}
62+
63+
/** Add headers and trailers to CookiesHolder if there's any. * */
64+
static class UpdateCookieListener<RespT>
65+
extends ForwardingClientCallListener.SimpleForwardingClientCallListener<RespT> {
66+
67+
private final CookiesHolder cookie;
68+
69+
UpdateCookieListener(ClientCall.Listener<RespT> delegate, CookiesHolder cookiesHolder) {
70+
super(delegate);
71+
this.cookie = cookiesHolder;
72+
}
73+
74+
@Override
75+
public void onHeaders(Metadata headers) {
76+
try {
77+
cookie.extractCookiesFromMetadata(headers);
78+
} catch (Throwable e) {
79+
LOG.log(Level.WARNING, "Failed to extract cookie from response headers.", e);
80+
} finally {
81+
super.onHeaders(headers);
82+
}
83+
}
84+
85+
@Override
86+
public void onClose(Status status, Metadata trailers) {
87+
try {
88+
cookie.extractCookiesFromMetadata(trailers);
89+
} catch (Throwable e) {
90+
LOG.log(Level.WARNING, "Failed to extract cookie from response trailers.", e);
91+
} finally {
92+
super.onClose(status, trailers);
93+
}
94+
}
95+
}
96+
}

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookiesServerStreamingCallable.java‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
/*
2+
* Copyright 2023 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.data.v2.stub;
17+
18+
import static com.google.cloud.bigtable.data.v2.stub.CookiesHolder.COOKIES_HOLDER_KEY;
19+
20+
import com.google.api.gax.grpc.GrpcCallContext;
21+
import com.google.api.gax.rpc.ApiCallContext;
22+
import com.google.api.gax.rpc.ResponseObserver;
23+
import com.google.api.gax.rpc.ServerStreamingCallable;
24+
25+
/**
26+
* The cookie holder will act as operation scoped storage for all retry attempts. Each attempt's
27+
* cookies will be merged into the value holder and will be sent out with the next retry attempt.
28+
*/
29+
class CookiesServerStreamingCallable<RequestT, ResponseT>
30+
extends ServerStreamingCallable<RequestT, ResponseT> {
31+
32+
private final ServerStreamingCallable<RequestT, ResponseT> callable;
33+
34+
CookiesServerStreamingCallable(ServerStreamingCallable<RequestT, ResponseT> innerCallable) {
35+
this.callable = innerCallable;
36+
}
37+
38+
@Override
39+
public void call(
40+
RequestT request, ResponseObserver<ResponseT> responseObserver, ApiCallContext context) {
41+
GrpcCallContext grpcCallContext = (GrpcCallContext) context;
42+
callable.call(
43+
request,
44+
responseObserver,
45+
grpcCallContext.withCallOptions(
46+
grpcCallContext.getCallOptions().withOption(COOKIES_HOLDER_KEY, new CookiesHolder())));
47+
}
48+
}

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookiesUnaryCallable.java‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
/*
2+
* Copyright 2023 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.data.v2.stub;
17+
18+
import static com.google.cloud.bigtable.data.v2.stub.CookiesHolder.COOKIES_HOLDER_KEY;
19+
20+
import com.google.api.core.ApiFuture;
21+
import com.google.api.gax.grpc.GrpcCallContext;
22+
import com.google.api.gax.rpc.ApiCallContext;
23+
import com.google.api.gax.rpc.UnaryCallable;
24+
25+
/**
26+
* The cookie holder will act as operation scoped storage for all retry attempts. Each attempt's
27+
* cookies will be merged into the value holder and will be sent out with the next retry attempt.
28+
*/
29+
class CookiesUnaryCallable<RequestT, ResponseT> extends UnaryCallable<RequestT, ResponseT> {
30+
private final UnaryCallable<RequestT, ResponseT> innerCallable;
31+
32+
CookiesUnaryCallable(UnaryCallable<RequestT, ResponseT> callable) {
33+
this.innerCallable = callable;
34+
}
35+
36+
@Override
37+
public ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext context) {
38+
GrpcCallContext grpcCallContext = (GrpcCallContext) context;
39+
return innerCallable.futureCall(
40+
request,
41+
grpcCallContext.withCallOptions(
42+
grpcCallContext.getCallOptions().withOption(COOKIES_HOLDER_KEY, new CookiesHolder())));
43+
}
44+
}

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java‎

Lines changed: 44 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,14 @@ public static EnhancedBigtableStubSettings finalizeSettings(
185185
// workaround JWT audience issues
186186
patchCredentials(builder);
187187

188+
// patch cookies interceptor
189+
InstantiatingGrpcChannelProvider.Builder transportProvider = null;
190+
if (builder.getTransportChannelProvider() instanceof InstantiatingGrpcChannelProvider) {
191+
transportProvider =
192+
((InstantiatingGrpcChannelProvider) builder.getTransportChannelProvider()).toBuilder();
193+
transportProvider.setInterceptorProvider(() -> ImmutableList.of(new CookiesInterceptor()));
194+
}
195+
188196
// Inject channel priming
189197
if (settings.isRefreshingChannel()) {
190198
// Fix the credentials so that they can be shared
@@ -194,20 +202,18 @@ public static EnhancedBigtableStubSettings finalizeSettings(
194202
}
195203
builder.setCredentialsProvider(FixedCredentialsProvider.create(credentials));
196204

197-
// Inject the primer
198-
InstantiatingGrpcChannelProvider transportProvider =
199-
(InstantiatingGrpcChannelProvider) settings.getTransportChannelProvider();
200-
201-
builder.setTransportChannelProvider(
202-
transportProvider
203-
.toBuilder()
204-
.setChannelPrimer(
205-
BigtableChannelPrimer.create(
206-
credentials,
207-
settings.getProjectId(),
208-
settings.getInstanceId(),
209-
settings.getAppProfileId()))
210-
.build());
205+
if (transportProvider != null) {
206+
transportProvider.setChannelPrimer(
207+
BigtableChannelPrimer.create(
208+
credentials,
209+
settings.getProjectId(),
210+
settings.getInstanceId(),
211+
settings.getAppProfileId()));
212+
}
213+
}
214+
215+
if (transportProvider != null) {
216+
builder.setTransportChannelProvider(transportProvider.build());
211217
}
212218

213219
ImmutableMap<TagKey, TagValue> attributes =
@@ -365,7 +371,11 @@ public <RowT> ServerStreamingCallable<Query, RowT> createReadRowsCallable(
365371
new TracedServerStreamingCallable<>(
366372
readRowsUserCallable, clientContext.getTracerFactory(), span);
367373

368-
return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
374+
// CookieHolder needs to be injected to the CallOptions outside of retries, otherwise retry
375+
// attempts won't see a CookieHolder.
376+
ServerStreamingCallable<Query, RowT> withCookie = new CookiesServerStreamingCallable<>(traced);
377+
378+
return withCookie.withDefaultCallContext(clientContext.getDefaultCallContext());
369379
}
370380

371381
/**
@@ -401,7 +411,9 @@ public <RowT> UnaryCallable<Query, RowT> createReadRowCallable(RowAdapter<RowT>
401411
new TracedUnaryCallable<>(
402412
firstRow, clientContext.getTracerFactory(), getSpanName("ReadRow"));
403413

404-
return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
414+
UnaryCallable<Query, RowT> withCookie = new CookiesUnaryCallable<>(traced);
415+
416+
return withCookie.withDefaultCallContext(clientContext.getDefaultCallContext());
405417
}
406418

407419
/**
@@ -642,7 +654,9 @@ private UnaryCallable<BulkMutation, Void> createBulkMutateRowsCallable() {
642654
new TracedUnaryCallable<>(
643655
tracedBatcherUnaryCallable, clientContext.getTracerFactory(), spanName);
644656

645-
return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
657+
UnaryCallable<BulkMutation, Void> withCookie = new CookiesUnaryCallable<>(traced);
658+
659+
return withCookie.withDefaultCallContext(clientContext.getDefaultCallContext());
646660
}
647661

648662
/**
@@ -924,7 +938,10 @@ public Map<String, String> extract(
924938
ServerStreamingCallable<String, ByteStringRange> traced =
925939
new TracedServerStreamingCallable<>(retrying, clientContext.getTracerFactory(), span);
926940

927-
return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
941+
ServerStreamingCallable<String, ByteStringRange> withCookie =
942+
new CookiesServerStreamingCallable<>(traced);
943+
944+
return withCookie.withDefaultCallContext(clientContext.getDefaultCallContext());
928945
}
929946

930947
/**
@@ -1004,7 +1021,10 @@ public Map<String, String> extract(
10041021
new TracedServerStreamingCallable<>(
10051022
readChangeStreamUserCallable, clientContext.getTracerFactory(), span);
10061023

1007-
return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
1024+
ServerStreamingCallable<ReadChangeStreamQuery, ChangeStreamRecordT> withCookie =
1025+
new CookiesServerStreamingCallable<>(traced);
1026+
1027+
return withCookie.withDefaultCallContext(clientContext.getDefaultCallContext());
10081028
}
10091029

10101030
/**
@@ -1017,7 +1037,11 @@ private <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> createUserFacin
10171037
UnaryCallable<RequestT, ResponseT> traced =
10181038
new TracedUnaryCallable<>(inner, clientContext.getTracerFactory(), getSpanName(methodName));
10191039

1020-
return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
1040+
// CookieHolder needs to be injected to the CallOptions outside of retries, otherwise retry
1041+
// attempts won't see a CookieHolder.
1042+
UnaryCallable<RequestT, ResponseT> withCookie = new CookiesUnaryCallable<>(traced);
1043+
1044+
return withCookie.withDefaultCallContext(clientContext.getDefaultCallContext());
10211045
}
10221046

10231047
private UnaryCallable<PingAndWarmRequest, PingAndWarmResponse> createPingAndWarmCallable() {

0 commit comments

Comments
 (0)