-
Notifications
You must be signed in to change notification settings - Fork 88
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
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>
- Loading branch information
1 parent
0827252
commit 4c73abd
Showing
8 changed files
with
924 additions
and
21 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
70 changes: 70 additions & 0 deletions
70
...le-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolder.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,70 @@ | ||
| /* | ||
| * Copyright 2023 Google LLC | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * https://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
| package com.google.cloud.bigtable.data.v2.stub; | ||
|
|
||
| import io.grpc.CallOptions; | ||
| import io.grpc.Metadata; | ||
| import java.util.HashMap; | ||
| import java.util.Map; | ||
| import javax.annotation.Nullable; | ||
|
|
||
| /** A cookie that holds information for retry or routing */ | ||
| class CookiesHolder { | ||
|
|
||
| static final CallOptions.Key<CookiesHolder> COOKIES_HOLDER_KEY = | ||
| CallOptions.Key.create("bigtable-cookies"); | ||
|
|
||
| /** Routing cookie key prefix. */ | ||
| static final String COOKIE_KEY_PREFIX = "x-goog-cbt-cookie"; | ||
|
|
||
| /** A map that stores all the routing cookies. */ | ||
| private final Map<Metadata.Key<String>, String> cookies = new HashMap<>(); | ||
|
|
||
| /** Returns CookiesHolder if presents in CallOptions, otherwise returns null. */ | ||
| @Nullable | ||
| static CookiesHolder fromCallOptions(CallOptions options) { | ||
| // CookiesHolder should be added by CookiesServerStreamingCallable and | ||
| // CookiesUnaryCallable for most methods. However, methods like PingAndWarm | ||
| // doesn't support routing cookie, in which case this will return null. | ||
| return options.getOption(COOKIES_HOLDER_KEY); | ||
| } | ||
|
|
||
| /** Add all the routing cookies to headers if any. */ | ||
| Metadata injectCookiesInRequestHeaders(Metadata headers) { | ||
| for (Metadata.Key<String> key : cookies.keySet()) { | ||
| headers.put(key, cookies.get(key)); | ||
| } | ||
| return headers; | ||
| } | ||
|
|
||
| /** | ||
| * Iterate through all the keys in initial or trailing metadata, and add all the keys that match | ||
| * COOKIE_KEY_PREFIX to cookies. Values in trailers will override the value set in initial | ||
| * metadata for the same keys. | ||
| */ | ||
| void extractCookiesFromMetadata(@Nullable Metadata trailers) { | ||
| if (trailers == null) { | ||
| return; | ||
| } | ||
| for (String key : trailers.keys()) { | ||
| if (key.startsWith(COOKIE_KEY_PREFIX)) { | ||
| Metadata.Key<String> metadataKey = Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER); | ||
| String value = trailers.get(metadataKey); | ||
| cookies.put(metadataKey, value); | ||
| } | ||
| } | ||
| } | ||
| } |
96 changes: 96 additions & 0 deletions
96
...oud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookiesInterceptor.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,96 @@ | ||
| /* | ||
| * Copyright 2023 Google LLC | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * https://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
| package com.google.cloud.bigtable.data.v2.stub; | ||
|
|
||
| import io.grpc.CallOptions; | ||
| import io.grpc.Channel; | ||
| import io.grpc.ClientCall; | ||
| import io.grpc.ClientInterceptor; | ||
| import io.grpc.ForwardingClientCall; | ||
| import io.grpc.ForwardingClientCallListener; | ||
| import io.grpc.Metadata; | ||
| import io.grpc.MethodDescriptor; | ||
| import io.grpc.Status; | ||
| import java.util.logging.Level; | ||
| import java.util.logging.Logger; | ||
|
|
||
| /** | ||
| * A cookie interceptor that checks the cookie value from returned trailer, updates the cookie | ||
| * holder, and inject it in the header of the next request. | ||
| */ | ||
| class CookiesInterceptor implements ClientInterceptor { | ||
|
|
||
| private static final Logger LOG = Logger.getLogger(CookiesInterceptor.class.getName()); | ||
|
|
||
| @Override | ||
| public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall( | ||
| MethodDescriptor<ReqT, RespT> methodDescriptor, CallOptions callOptions, Channel channel) { | ||
| return new ForwardingClientCall.SimpleForwardingClientCall<ReqT, RespT>( | ||
| channel.newCall(methodDescriptor, callOptions)) { | ||
| @Override | ||
| public void start(Listener<RespT> responseListener, Metadata headers) { | ||
| // Gets the CookiesHolder added from CookiesServerStreamingCallable and | ||
| // CookiesUnaryCallable. | ||
| // Add CookiesHolder content to request headers if there's any. | ||
| try { | ||
| CookiesHolder cookie = CookiesHolder.fromCallOptions(callOptions); | ||
| if (cookie != null) { | ||
| headers = cookie.injectCookiesInRequestHeaders(headers); | ||
| responseListener = new UpdateCookieListener<>(responseListener, cookie); | ||
| } | ||
| } catch (Throwable e) { | ||
| LOG.warning("Failed to inject cookie to request headers: " + e); | ||
| } finally { | ||
| super.start(responseListener, headers); | ||
| } | ||
| } | ||
| }; | ||
| } | ||
|
|
||
| /** Add headers and trailers to CookiesHolder if there's any. * */ | ||
| static class UpdateCookieListener<RespT> | ||
| extends ForwardingClientCallListener.SimpleForwardingClientCallListener<RespT> { | ||
|
|
||
| private final CookiesHolder cookie; | ||
|
|
||
| UpdateCookieListener(ClientCall.Listener<RespT> delegate, CookiesHolder cookiesHolder) { | ||
| super(delegate); | ||
| this.cookie = cookiesHolder; | ||
| } | ||
|
|
||
| @Override | ||
| public void onHeaders(Metadata headers) { | ||
| try { | ||
| cookie.extractCookiesFromMetadata(headers); | ||
| } catch (Throwable e) { | ||
| LOG.log(Level.WARNING, "Failed to extract cookie from response headers.", e); | ||
| } finally { | ||
| super.onHeaders(headers); | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public void onClose(Status status, Metadata trailers) { | ||
| try { | ||
| cookie.extractCookiesFromMetadata(trailers); | ||
| } catch (Throwable e) { | ||
| LOG.log(Level.WARNING, "Failed to extract cookie from response trailers.", e); | ||
| } finally { | ||
| super.onClose(status, trailers); | ||
| } | ||
| } | ||
| } | ||
| } |
48 changes: 48 additions & 0 deletions
48
.../src/main/java/com/google/cloud/bigtable/data/v2/stub/CookiesServerStreamingCallable.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,48 @@ | ||
| /* | ||
| * Copyright 2023 Google LLC | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * https://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
| package com.google.cloud.bigtable.data.v2.stub; | ||
|
|
||
| import static com.google.cloud.bigtable.data.v2.stub.CookiesHolder.COOKIES_HOLDER_KEY; | ||
|
|
||
| import com.google.api.gax.grpc.GrpcCallContext; | ||
| import com.google.api.gax.rpc.ApiCallContext; | ||
| import com.google.api.gax.rpc.ResponseObserver; | ||
| import com.google.api.gax.rpc.ServerStreamingCallable; | ||
|
|
||
| /** | ||
| * The cookie holder will act as operation scoped storage for all retry attempts. Each attempt's | ||
| * cookies will be merged into the value holder and will be sent out with the next retry attempt. | ||
| */ | ||
| class CookiesServerStreamingCallable<RequestT, ResponseT> | ||
| extends ServerStreamingCallable<RequestT, ResponseT> { | ||
|
|
||
| private final ServerStreamingCallable<RequestT, ResponseT> callable; | ||
|
|
||
| CookiesServerStreamingCallable(ServerStreamingCallable<RequestT, ResponseT> innerCallable) { | ||
| this.callable = innerCallable; | ||
| } | ||
|
|
||
| @Override | ||
| public void call( | ||
| RequestT request, ResponseObserver<ResponseT> responseObserver, ApiCallContext context) { | ||
| GrpcCallContext grpcCallContext = (GrpcCallContext) context; | ||
| callable.call( | ||
| request, | ||
| responseObserver, | ||
| grpcCallContext.withCallOptions( | ||
| grpcCallContext.getCallOptions().withOption(COOKIES_HOLDER_KEY, new CookiesHolder()))); | ||
| } | ||
| } |
44 changes: 44 additions & 0 deletions
44
...d-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookiesUnaryCallable.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,44 @@ | ||
| /* | ||
| * Copyright 2023 Google LLC | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * https://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
| package com.google.cloud.bigtable.data.v2.stub; | ||
|
|
||
| import static com.google.cloud.bigtable.data.v2.stub.CookiesHolder.COOKIES_HOLDER_KEY; | ||
|
|
||
| import com.google.api.core.ApiFuture; | ||
| import com.google.api.gax.grpc.GrpcCallContext; | ||
| import com.google.api.gax.rpc.ApiCallContext; | ||
| import com.google.api.gax.rpc.UnaryCallable; | ||
|
|
||
| /** | ||
| * The cookie holder will act as operation scoped storage for all retry attempts. Each attempt's | ||
| * cookies will be merged into the value holder and will be sent out with the next retry attempt. | ||
| */ | ||
| class CookiesUnaryCallable<RequestT, ResponseT> extends UnaryCallable<RequestT, ResponseT> { | ||
| private final UnaryCallable<RequestT, ResponseT> innerCallable; | ||
|
|
||
| CookiesUnaryCallable(UnaryCallable<RequestT, ResponseT> callable) { | ||
| this.innerCallable = callable; | ||
| } | ||
|
|
||
| @Override | ||
| public ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext context) { | ||
| GrpcCallContext grpcCallContext = (GrpcCallContext) context; | ||
| return innerCallable.futureCall( | ||
| request, | ||
| grpcCallContext.withCallOptions( | ||
| grpcCallContext.getCallOptions().withOption(COOKIES_HOLDER_KEY, new CookiesHolder()))); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.