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,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.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 = CookiesHolder.fromCallOptions(callOptions);
+        if (cookie != null) {
+          cookie.addRoutingCookieToHeaders(headers);
+          responseListener = new UpdateCookieListener<>(responseListener, callOptions);
+        }
+        super.start(responseListener, 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 = CookiesHolder.fromCallOptions(callOptions);
+      if (cookiesHolder != null) {
+        cookiesHolder.setRoutingCookieFromTrailers(trailers);
+      }
+      super.onClose(status, trailers);
+    }
+  }
+}

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

@@ -0,0 +1,63 @@
+/*
+ * 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.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+/** 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 ConcurrentHashMap<>();
+
+  /** Returns CookiesHolder if presents in CallOptions. Otherwise returns null. */
+  static CookiesHolder fromCallOptions(CallOptions options) {
+    return options.getOption(COOKIES_HOLDER_KEY);
+  }
+
+  /** Add all the routing cookies to headers if any. */
+  Metadata addRoutingCookieToHeaders(Metadata headers) {
+    if (headers != null && !cookies.isEmpty()) {
+      for (Metadata.Key<String> key : cookies.keySet()) headers.put(key, cookies.get(key));
+    }
+    return headers;
+  }
+
+  /**
+   * Iterate through all the keys in trailing metadata, and add all the keys that match
+   * COOKIE_KEY_PREFIX to cookies.
+   */
+  void setRoutingCookieFromTrailers(Metadata trailers) {
+    if (trailers != null) {
+      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);
+        }
+      }
+    }
+  }
+}

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());
   }
 
   /**
@@ -642,7 +648,9 @@ private UnaryCallable<BulkMutation, Void> createBulkMutateRowsCallable() {
         new TracedUnaryCallable<>(
             tracedBatcherUnaryCallable, clientContext.getTracerFactory(), spanName);
 
-    return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
+    UnaryCallable<BulkMutation, Void> withCookie = new CookiesUnaryCallable<>(traced);
+
+    return withCookie.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
 
   /**
@@ -1017,7 +1025,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")