diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolder.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolder.java index 6727cccfd1..88daad5464 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolder.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolder.java @@ -19,6 +19,7 @@ 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 { @@ -32,8 +33,13 @@ class CookiesHolder { /** A map that stores all the routing cookies. */ private final Map, String> cookies = new HashMap<>(); - /** Returns CookiesHolder if presents in CallOptions. Otherwise returns null. */ + /** 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 + // or ReadChangeStream doesn't support routing cookie, in which case this + // method will return null. return options.getOption(COOKIES_HOLDER_KEY); } @@ -49,7 +55,7 @@ Metadata injectCookiesInRequestHeaders(Metadata headers) { * Iterate through all the keys in trailing metadata, and add all the keys that match * COOKIE_KEY_PREFIX to cookies. */ - void extractCookiesFromResponseTrailers(Metadata trailers) { + void extractCookiesFromResponseTrailers(@Nullable Metadata trailers) { if (trailers == null) { return; } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookiesInterceptor.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookiesInterceptor.java index d1d944420f..91bee7ca4a 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookiesInterceptor.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/CookiesInterceptor.java @@ -24,6 +24,7 @@ import io.grpc.Metadata; import io.grpc.MethodDescriptor; import io.grpc.Status; +import java.util.logging.Logger; /** * A cookie interceptor that checks the cookie value from returned trailer, updates the cookie @@ -31,6 +32,8 @@ */ class CookiesInterceptor implements ClientInterceptor { + private static final Logger LOG = Logger.getLogger(CookiesInterceptor.class.getName()); + @Override public ClientCall interceptCall( MethodDescriptor methodDescriptor, CallOptions callOptions, Channel channel) { @@ -41,12 +44,17 @@ public void start(Listener responseListener, Metadata headers) { // Gets the CookiesHolder added from CookiesServerStreamingCallable and // CookiesUnaryCallable. // Add CookiesHolder content to request headers if there's any. - CookiesHolder cookie = CookiesHolder.fromCallOptions(callOptions); - if (cookie != null) { - cookie.injectCookiesInRequestHeaders(headers); - responseListener = new UpdateCookieListener<>(responseListener, cookie); + try { + CookiesHolder cookie = CookiesHolder.fromCallOptions(callOptions); + if (cookie != null) { + 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); } - super.start(responseListener, headers); } }; } @@ -64,8 +72,13 @@ static class UpdateCookieListener @Override public void onClose(Status status, Metadata trailers) { - cookie.extractCookiesFromResponseTrailers(trailers); - super.onClose(status, trailers); + try { + cookie.extractCookiesFromResponseTrailers(trailers); + } catch (Throwable e) { + LOG.warning("Failed to extract cookie from response trailers: " + e); + } finally { + super.onClose(status, trailers); + } } } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java index 370f0dbd58..41182564ab 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/CookiesHolderTest.java @@ -106,6 +106,10 @@ public ServerCall.Listener interceptCall( .setProjectId("fake-project") .setInstanceId("fake-instance"); + // Override CheckAndMutate and ReadModifyWrite retry settings here. These operations + // are currently not retryable but this could change in the future after we + // have routing cookie sends back meaningful information and changes how retry works. + // Routing cookie still needs to be respected and handled by the callables. settings .stubSettings() .checkAndMutateRowSettings() @@ -306,6 +310,10 @@ public void testNoCookieSucceedSampleRowKeys() { @Test public void testAllMethodsAreCalled() { + // This test ensures that all methods respect the retry cookie except for the ones that are + // explicitly added to the methods list. It requires that any newly method is exercised in this + // test. + // This is enforced by introspecting grpc method descriptors. client.readRows(Query.create("fake-table")).iterator().hasNext(); fakeService.count.set(0); client.mutateRow(RowMutation.create("fake-table", "key").setCell("cf", "q", "v")); @@ -351,11 +359,7 @@ public void readRows( ReadRowsRequest request, StreamObserver responseObserver) { if (count.getAndIncrement() < 1) { Metadata trailers = new Metadata(); - if (returnCookie) { - trailers.put(ROUTING_COOKIE_1, "readRows"); - trailers.put(ROUTING_COOKIE_2, testCookie); - trailers.put(BAD_KEY, "bad-key"); - } + maybePopulateCookie(trailers, "readRows"); StatusRuntimeException exception = new StatusRuntimeException(Status.UNAVAILABLE, trailers); responseObserver.onError(exception); return; @@ -369,11 +373,7 @@ public void mutateRow( MutateRowRequest request, StreamObserver responseObserver) { if (count.getAndIncrement() < 1) { Metadata trailers = new Metadata(); - if (returnCookie) { - trailers.put(ROUTING_COOKIE_1, "mutateRow"); - trailers.put(ROUTING_COOKIE_2, testCookie); - trailers.put(BAD_KEY, "bad-key"); - } + maybePopulateCookie(trailers, "mutateRow"); StatusRuntimeException exception = new StatusRuntimeException(Status.UNAVAILABLE, trailers); responseObserver.onError(exception); return; @@ -387,11 +387,7 @@ public void mutateRows( MutateRowsRequest request, StreamObserver responseObserver) { if (count.getAndIncrement() < 1) { Metadata trailers = new Metadata(); - if (returnCookie) { - trailers.put(ROUTING_COOKIE_1, "mutateRows"); - trailers.put(ROUTING_COOKIE_2, testCookie); - trailers.put(BAD_KEY, "bad-key"); - } + maybePopulateCookie(trailers, "mutateRows"); StatusRuntimeException exception = new StatusRuntimeException(Status.UNAVAILABLE, trailers); responseObserver.onError(exception); return; @@ -408,11 +404,7 @@ public void sampleRowKeys( SampleRowKeysRequest request, StreamObserver responseObserver) { if (count.getAndIncrement() < 1) { Metadata trailers = new Metadata(); - if (returnCookie) { - trailers.put(ROUTING_COOKIE_1, "sampleRowKeys"); - trailers.put(ROUTING_COOKIE_2, testCookie); - trailers.put(BAD_KEY, "bad-key"); - } + maybePopulateCookie(trailers, "sampleRowKeys"); StatusRuntimeException exception = new StatusRuntimeException(Status.UNAVAILABLE, trailers); responseObserver.onError(exception); return; @@ -427,11 +419,7 @@ public void checkAndMutateRow( StreamObserver responseObserver) { if (count.getAndIncrement() < 1) { Metadata trailers = new Metadata(); - if (returnCookie) { - trailers.put(ROUTING_COOKIE_1, "checkAndMutate"); - trailers.put(ROUTING_COOKIE_2, testCookie); - trailers.put(BAD_KEY, "bad-key"); - } + maybePopulateCookie(trailers, "checkAndMutate"); StatusRuntimeException exception = new StatusRuntimeException(Status.UNAVAILABLE, trailers); responseObserver.onError(exception); return; @@ -446,11 +434,7 @@ public void readModifyWriteRow( StreamObserver responseObserver) { if (count.getAndIncrement() < 1) { Metadata trailers = new Metadata(); - if (returnCookie) { - trailers.put(ROUTING_COOKIE_1, "readModifyWrite"); - trailers.put(ROUTING_COOKIE_2, testCookie); - trailers.put(BAD_KEY, "bad-key"); - } + maybePopulateCookie(trailers, "readModifyWrite"); StatusRuntimeException exception = new StatusRuntimeException(Status.UNAVAILABLE, trailers); responseObserver.onError(exception); return; @@ -458,5 +442,13 @@ public void readModifyWriteRow( responseObserver.onNext(ReadModifyWriteRowResponse.getDefaultInstance()); responseObserver.onCompleted(); } + + private void maybePopulateCookie(Metadata trailers, String label) { + if (returnCookie) { + trailers.put(ROUTING_COOKIE_1, label); + trailers.put(ROUTING_COOKIE_2, testCookie); + trailers.put(BAD_KEY, "bad-key"); + } + } } }