feat: client sends routing cookie back to server

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

@@ -0,0 +1,75 @@
+/*
+ * 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 static com.google.cloud.bigtable.data.v2.stub.CookiesHolder.ROUTING_COOKIE_METADATA_KEY;
+
+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;
+
+/**
+ * A cookie interceptor that checks the cookie value from returned ErrorInfo, updates the cookie
+ * holder, and inject it in the header of the next request.
+ */
+class CookieInterceptor implements ClientInterceptor {
+
+  @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) {
+        CookiesHolder cookie = callOptions.getOption(COOKIES_HOLDER_KEY);
+        if (cookie != null && cookie.getRoutingCookie() != null) {
+          headers.put(ROUTING_COOKIE_METADATA_KEY, cookie.getRoutingCookie());
+        }
+        super.start(new UpdateCookieListener<>(responseListener, callOptions), headers);
+      }
+    };
+  }
+
+  static class UpdateCookieListener<RespT>
+      extends ForwardingClientCallListener.SimpleForwardingClientCallListener<RespT> {
+
+    private final CallOptions callOptions;
+
+    UpdateCookieListener(ClientCall.Listener<RespT> delegate, CallOptions callOptions) {
+      super(delegate);
+      this.callOptions = callOptions;
+    }
+
+    @Override
+    public void onClose(Status status, Metadata trailers) {
+      CookiesHolder cookiesHolder = callOptions.getOption(COOKIES_HOLDER_KEY);
+      if (cookiesHolder != null
+          && trailers != null
+          && trailers.containsKey(ROUTING_COOKIE_METADATA_KEY)) {
+        String cookie = trailers.get(ROUTING_COOKIE_METADATA_KEY);
+        cookiesHolder.setRoutingCookie(cookie);
+      }
+      super.onClose(status, trailers);
+    }
+  }
+}

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

@@ -0,0 +1,42 @@
+/*
+ * 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 javax.annotation.Nullable;
+
+/** A cookie that holds information for the retry */
+class CookiesHolder {
+
+  static final CallOptions.Key<CookiesHolder> COOKIES_HOLDER_KEY =
+      CallOptions.Key.create("bigtable-cookies");
+
+  static final String ROUTING_COOKIE_KEY = "x-goog-cbt-cookie-routing";
+
+  static final Metadata.Key<String> ROUTING_COOKIE_METADATA_KEY =
+      Metadata.Key.of(ROUTING_COOKIE_KEY, Metadata.ASCII_STRING_MARSHALLER);
+
+  @Nullable private String routingCookie;
+
+  void setRoutingCookie(@Nullable String routingCookie) {
+    this.routingCookie = routingCookie;
+  }
+
+  String getRoutingCookie() {
+    return this.routingCookie;
+  }
+}

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

@@ -0,0 +1,45 @@
+/*
+ * 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;
+
+/** Cookie callable injects a placeholder for bigtable retry cookie. */
+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())));
+  }
+}

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

@@ -0,0 +1,41 @@
+/*
+ * 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;
+
+/** Cookie callable injects a placeholder for bigtable retry cookie. */
+class CookiesUnaryCallable<RequestT, ResponseT> extends UnaryCallable<RequestT, ResponseT> {
+  private 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())));
+  }
+}

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

@@ -365,7 +365,11 @@ public <RowT> ServerStreamingCallable<Query, RowT> createReadRowsCallable(
         new TracedServerStreamingCallable<>(
             readRowsUserCallable, clientContext.getTracerFactory(), span);
 
-    return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
+    // CookieHolder needs to be injected to the CallOptions outside of retries, otherwise retry
+    // attempts won't see a CookieHolder.
+    ServerStreamingCallable<Query, RowT> withCookie = new CookiesServerStreamingCallable<>(traced);
+
+    return withCookie.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
 
   /**
@@ -401,7 +405,9 @@ public <RowT> UnaryCallable<Query, RowT> createReadRowCallable(RowAdapter<RowT>
         new TracedUnaryCallable<>(
             firstRow, clientContext.getTracerFactory(), getSpanName("ReadRow"));
 
-    return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
+    UnaryCallable<Query, RowT> withCookie = new CookiesUnaryCallable<>(traced);
+
+    return withCookie.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
 
   /**
@@ -1017,7 +1023,11 @@ private <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> createUserFacin
     UnaryCallable<RequestT, ResponseT> traced =
         new TracedUnaryCallable<>(inner, clientContext.getTracerFactory(), getSpanName(methodName));
 
-    return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
+    // CookieHolder needs to be injected to the CallOptions outside of retries, otherwise retry
+    // attempts won't see a CookieHolder.
+    UnaryCallable<RequestT, ResponseT> withCookie = new CookiesUnaryCallable<>(traced);
+
+    return withCookie.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
 
   private UnaryCallable<PingAndWarmRequest, PingAndWarmResponse> createPingAndWarmCallable() {

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

@@ -329,7 +329,8 @@ public static InstantiatingGrpcChannelProvider.Builder defaultGrpcTransportProvi
             Duration.ofSeconds(10)) // wait this long before considering the connection dead
         // Attempts direct access to CBT service over gRPC to improve throughput,
         // whether the attempt is allowed is totally controlled by service owner.
-        .setAttemptDirectPath(true);
+        .setAttemptDirectPath(true)
+        .setInterceptorProvider(() -> ImmutableList.of(new CookieInterceptor()));
   }
 
   @SuppressWarnings("WeakerAccess")