From ec035598068ad739b6f88d4bda7fe93b16ece2d7 Mon Sep 17 00:00:00 2001
From: Igor Bernstein <igorbernstein@google.com>
Date: Sat, 2 Nov 2024 14:04:58 -0400
Subject: [PATCH 01/16] chore: remodel unary callables as server streaming
 callables with an adapter at the end

Change-Id: I8708dff0e192d7647ef2cb361fc0992e1ddd2b24
---
 .../stub/BigtableUnaryOperationCallable.java  | 183 ++++++++++++++++++
 .../data/v2/stub/EnhancedBigtableStub.java    |  87 ++++++++-
 .../v2/stub/EnhancedBigtableStubSettings.java |   6 +
 .../TransformingServerStreamingCallable.java  |  72 +++++++
 4 files changed, 338 insertions(+), 10 deletions(-)
 create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java
 create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/TransformingServerStreamingCallable.java

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java
new file mode 100644
index 0000000000..202498987d
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java
@@ -0,0 +1,183 @@
+/*
+ * Copyright 2024 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 com.google.api.core.AbstractApiFuture;
+import com.google.api.core.ApiFuture;
+import com.google.api.gax.grpc.GrpcStatusCode;
+import com.google.api.gax.rpc.ApiCallContext;
+import com.google.api.gax.rpc.InternalException;
+import com.google.api.gax.rpc.ResponseObserver;
+import com.google.api.gax.rpc.ServerStreamingCallable;
+import com.google.api.gax.rpc.StreamController;
+import com.google.api.gax.rpc.UnaryCallable;
+import com.google.api.gax.tracing.ApiTracer;
+import com.google.api.gax.tracing.ApiTracerFactory;
+import com.google.api.gax.tracing.SpanName;
+import com.google.common.base.Preconditions;
+import io.grpc.Status;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import javax.annotation.Nullable;
+
+/**
+ * Helper to convert a fake {@link ServerStreamingCallable} (ie only up to 1 response) into a {@link
+ * UnaryCallable}. It is intended to be the outermost callable of a chain.
+ *
+ * <p>Responsibilities:
+ *
+ * <ul>
+ *   <li>Operation level metrics
+ *   <li>Configuring the default call context
+ *   <li>Converting the result to a future
+ */
+class BigtableUnaryOperationCallable<ReqT, RespT> extends UnaryCallable<ReqT, RespT> {
+  private static final Logger LOGGER =
+      Logger.getLogger(BigtableUnaryOperationCallable.class.getName());
+
+  private final ServerStreamingCallable<ReqT, RespT> inner;
+  private final ApiCallContext defaultCallContext;
+  private final ApiTracerFactory tracerFactory;
+  private final SpanName spanName;
+  private final boolean allowNoResponse;
+
+  public BigtableUnaryOperationCallable(
+      ServerStreamingCallable<ReqT, RespT> inner,
+      ApiCallContext defaultCallContext,
+      ApiTracerFactory tracerFactory,
+      SpanName spanName,
+      boolean allowNoResponse) {
+    this.inner = inner;
+    this.defaultCallContext = defaultCallContext;
+    this.tracerFactory = tracerFactory;
+    this.spanName = spanName;
+    this.allowNoResponse = allowNoResponse;
+  }
+
+  @Override
+  public ApiFuture<RespT> futureCall(ReqT req, ApiCallContext apiCallContext) {
+    apiCallContext = defaultCallContext.merge(apiCallContext);
+
+    ApiTracer apiTracer =
+        tracerFactory.newTracer(
+            apiCallContext.getTracer(), spanName, ApiTracerFactory.OperationType.Unary);
+
+    apiCallContext = apiCallContext.withTracer(apiTracer);
+
+    UnaryFuture f = new UnaryFuture(apiTracer, allowNoResponse);
+    inner.call(req, f, apiCallContext);
+    return f;
+  }
+
+  class UnaryFuture extends AbstractApiFuture<RespT> implements ResponseObserver<RespT> {
+    private final ApiTracer tracer;
+    private final boolean allowNoResponse;
+
+    private StreamController controller;
+    private boolean responseReceived;
+    private @Nullable RespT response;
+
+    private UnaryFuture(ApiTracer tracer, boolean allowNoResponse) {
+      this.tracer = Preconditions.checkNotNull(tracer, "tracer can't be null");
+      this.allowNoResponse = allowNoResponse;
+      this.responseReceived = false;
+    }
+
+    @Override
+    public void onStart(StreamController controller) {
+      this.controller = controller;
+      controller.disableAutoInboundFlowControl();
+      // Request 2 to detect protocol bugs
+      controller.request(2);
+    }
+
+    /**
+     * Immediately cancel the future state and try to cancel the underlying operation. Will return
+     * false if the future is already resolved.
+     */
+    @Override
+    public boolean cancel(boolean mayInterruptIfRunning) {
+      if (super.cancel(mayInterruptIfRunning)) {
+        controller.cancel();
+        return true;
+      }
+      return false;
+    }
+
+    @Override
+    public void onResponse(RespT resp) {
+      tracer.responseReceived();
+
+      if (responseReceived) {
+        String msg =
+            String.format(
+                "Received multiple responses for a %s unary operation. Previous: %s, New: %s",
+                spanName, resp, response);
+        LOGGER.log(Level.WARNING, msg);
+
+        InternalException error =
+            new InternalException(msg, null, GrpcStatusCode.of(Status.Code.INTERNAL), false);
+        if (setException(error)) {
+          tracer.operationFailed(error);
+        }
+
+        controller.cancel();
+        return;
+      }
+
+      responseReceived = true;
+      this.response = resp;
+      this.tracer.responseReceived();
+    }
+
+    @Override
+    public void onError(Throwable throwable) {
+      if (this.setException(throwable)) {
+        tracer.operationFailed(throwable);
+      } else if (isCancelled()) {
+        tracer.operationCancelled();
+      }
+      // The future might've been resolved due to double response
+    }
+
+    @Override
+    public void onComplete() {
+      // Sanity check server response counts
+      if (!allowNoResponse && !responseReceived) {
+        String msg = spanName + " unary operation completed without a response message";
+        InternalException e =
+            new InternalException(msg, null, GrpcStatusCode.of(Status.Code.INTERNAL), false);
+
+        if (setException(e)) {
+          tracer.operationFailed(e);
+        }
+        return;
+      }
+
+      // check cancellation race
+      if (isCancelled()) {
+        tracer.operationCancelled();
+        return;
+      }
+
+      // try to resolve the future
+      if (this.set(response)) {
+        tracer.operationSucceeded();
+        return;
+      }
+    }
+  }
+}
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
index 2b2ca38540..6e5b1234c4 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
@@ -41,10 +41,12 @@
 import com.google.api.gax.retrying.RetryAlgorithm;
 import com.google.api.gax.retrying.RetryingExecutorWithContext;
 import com.google.api.gax.retrying.ScheduledRetryingExecutor;
+import com.google.api.gax.retrying.SimpleStreamResumptionStrategy;
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.Callables;
 import com.google.api.gax.rpc.ClientContext;
 import com.google.api.gax.rpc.RequestParamsExtractor;
+import com.google.api.gax.rpc.ResponseObserver;
 import com.google.api.gax.rpc.ServerStreamingCallSettings;
 import com.google.api.gax.rpc.ServerStreamingCallable;
 import com.google.api.gax.rpc.StatusCode.Code;
@@ -124,7 +126,6 @@
 import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsRetryingCallable;
 import com.google.cloud.bigtable.data.v2.stub.readrows.FilterMarkerRowsCallable;
 import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsBatchingDescriptor;
-import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsFirstCallable;
 import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsResumptionStrategy;
 import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsRetryCompletedCallable;
 import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsUserCallable;
@@ -155,6 +156,7 @@
 import java.io.IOException;
 import java.net.URI;
 import java.net.URISyntaxException;
+import java.time.Duration;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -568,16 +570,22 @@ public <RowT> UnaryCallable<Query, RowT> createReadRowCallable(RowAdapter<RowT>
                 .build(),
             rowAdapter);
 
-    ReadRowsUserCallable<RowT> readRowCallable =
-        new ReadRowsUserCallable<>(readRowsCallable, requestContext);
-
-    ReadRowsFirstCallable<RowT> firstRow = new ReadRowsFirstCallable<>(readRowCallable);
-
-    UnaryCallable<Query, RowT> traced =
-        new TracedUnaryCallable<>(
-            firstRow, clientContext.getTracerFactory(), getSpanName("ReadRow"));
+    ServerStreamingCallable<Query, RowT> readRowCallable =
+        new ServerStreamingCallable<Query, RowT>() {
+          @Override
+          public void call(
+              Query query, ResponseObserver<RowT> responseObserver, ApiCallContext apiCallContext) {
+            readRowsCallable.call(
+                query.limit(1).toProto(requestContext), responseObserver, apiCallContext);
+          }
+        };
 
-    return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
+    return new BigtableUnaryOperationCallable<>(
+        readRowCallable,
+        clientContext.getDefaultCallContext(),
+        clientContext.getTracerFactory(),
+        getSpanName("ReadRow"),
+        /*allowNoResponses=*/ true);
   }
 
   /**
@@ -1325,6 +1333,21 @@ private <BaseReqT, BaseRespT, ReqT, RespT> UnaryCallable<ReqT, RespT> createUnar
       UnaryCallSettings<ReqT, RespT> callSettings,
       Function<ReqT, BaseReqT> requestTransformer,
       Function<BaseRespT, RespT> responseTranformer) {
+    if (EnhancedBigtableStubSettings.SKIP_TRAILERS) {
+      return createUnaryCallableNew(
+          methodDescriptor, headerParamsFn, callSettings, requestTransformer, responseTranformer);
+    } else {
+      return createUnaryCallableNew(
+          methodDescriptor, headerParamsFn, callSettings, requestTransformer, responseTranformer);
+    }
+  }
+
+  private <BaseReqT, BaseRespT, ReqT, RespT> UnaryCallable<ReqT, RespT> createUnaryCallableOld(
+      MethodDescriptor<BaseReqT, BaseRespT> methodDescriptor,
+      RequestParamsExtractor<BaseReqT> headerParamsFn,
+      UnaryCallSettings<ReqT, RespT> callSettings,
+      Function<ReqT, BaseReqT> requestTransformer,
+      Function<BaseRespT, RespT> responseTranformer) {
 
     UnaryCallable<BaseReqT, BaseRespT> base =
         GrpcRawCallableFactory.createUnaryCallable(
@@ -1361,6 +1384,50 @@ public ApiFuture<RespT> futureCall(ReqT reqT, ApiCallContext apiCallContext) {
     return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
 
+  private <BaseReqT, BaseRespT, ReqT, RespT> UnaryCallable<ReqT, RespT> createUnaryCallableNew(
+      MethodDescriptor<BaseReqT, BaseRespT> methodDescriptor,
+      RequestParamsExtractor<BaseReqT> headerParamsFn,
+      UnaryCallSettings<ReqT, RespT> callSettings,
+      Function<ReqT, BaseReqT> requestTransformer,
+      Function<BaseRespT, RespT> responseTranformer) {
+
+    ServerStreamingCallable<BaseReqT, BaseRespT> base =
+        GrpcRawCallableFactory.createServerStreamingCallable(
+            GrpcCallSettings.<BaseReqT, BaseRespT>newBuilder()
+                .setMethodDescriptor(methodDescriptor)
+                .setParamsExtractor(headerParamsFn)
+                .build(),
+            callSettings.getRetryableCodes());
+
+    base = new StatsHeadersServerStreamingCallable<>(base);
+
+    base = new BigtableTracerStreamingCallable<>(base);
+
+    base = withRetries(base, convertUnaryToServerStreamingSettings(settings.mutateRowSettings()));
+
+    ServerStreamingCallable<ReqT, RespT> transformed =
+        new TransformingServerStreamingCallable<>(base, requestTransformer, responseTranformer);
+
+    return new BigtableUnaryOperationCallable<>(
+        transformed,
+        clientContext.getDefaultCallContext(),
+        clientContext.getTracerFactory(),
+        getSpanName(methodDescriptor.getBareMethodName()),
+        /* allowNoResponse= */ false);
+  }
+
+  private static <ReqT, RespT>
+      ServerStreamingCallSettings<ReqT, RespT> convertUnaryToServerStreamingSettings(
+          UnaryCallSettings<?, ?> unarySettings) {
+    return ServerStreamingCallSettings.<ReqT, RespT>newBuilder()
+        .setResumptionStrategy(new SimpleStreamResumptionStrategy<>())
+        .setRetryableCodes(unarySettings.getRetryableCodes())
+        .setRetrySettings(unarySettings.getRetrySettings())
+        .setIdleTimeoutDuration(Duration.ZERO)
+        .setWaitTimeoutDuration(Duration.ZERO)
+        .build();
+  }
+
   private UnaryCallable<PingAndWarmRequest, PingAndWarmResponse> createPingAndWarmCallable() {
     UnaryCallable<PingAndWarmRequest, PingAndWarmResponse> pingAndWarm =
         GrpcRawCallableFactory.createUnaryCallable(
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java
index 5e5dc64fd9..37c18952ee 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java
@@ -62,6 +62,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.logging.Logger;
 import javax.annotation.Nonnull;
@@ -108,6 +109,11 @@ public class EnhancedBigtableStubSettings extends StubSettings<EnhancedBigtableS
   private static final boolean DIRECT_PATH_ENABLED =
       Boolean.parseBoolean(System.getenv("CBT_ENABLE_DIRECTPATH"));
 
+  static final boolean SKIP_TRAILERS =
+      Optional.of(System.getenv("CBT_SKIP_HEADERS"))
+          .map(Boolean::parseBoolean)
+          .orElse(DIRECT_PATH_ENABLED);
+
   private static final Set<Code> IDEMPOTENT_RETRY_CODES =
       ImmutableSet.of(Code.DEADLINE_EXCEEDED, Code.UNAVAILABLE);
 
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/TransformingServerStreamingCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/TransformingServerStreamingCallable.java
new file mode 100644
index 0000000000..576d8257d7
--- /dev/null
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/TransformingServerStreamingCallable.java
@@ -0,0 +1,72 @@
+/*
+ * Copyright 2024 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 com.google.api.gax.rpc.ApiCallContext;
+import com.google.api.gax.rpc.ResponseObserver;
+import com.google.api.gax.rpc.ServerStreamingCallable;
+import com.google.api.gax.rpc.StreamController;
+import java.util.function.Function;
+
+/** Callable to help crossing api boundary lines between models and protos */
+class TransformingServerStreamingCallable<OuterReqT, OuterRespT, InnerReqT, InnerRespT>
+    extends ServerStreamingCallable<OuterReqT, OuterRespT> {
+  private final ServerStreamingCallable<InnerReqT, InnerRespT> inner;
+  private final Function<OuterReqT, InnerReqT> requestTransformer;
+  private final Function<InnerRespT, OuterRespT> responseTransformer;
+
+  public TransformingServerStreamingCallable(
+      ServerStreamingCallable<InnerReqT, InnerRespT> inner,
+      Function<OuterReqT, InnerReqT> requestTransformer,
+      Function<InnerRespT, OuterRespT> responseTransformer) {
+    this.inner = inner;
+    this.requestTransformer = requestTransformer;
+    this.responseTransformer = responseTransformer;
+  }
+
+  @Override
+  public void call(
+      OuterReqT outerReqT,
+      ResponseObserver<OuterRespT> outerObserver,
+      ApiCallContext apiCallContext) {
+    InnerReqT innerReq = requestTransformer.apply(outerReqT);
+
+    inner.call(
+        innerReq,
+        new ResponseObserver<InnerRespT>() {
+          @Override
+          public void onStart(StreamController streamController) {
+            outerObserver.onStart(streamController);
+          }
+
+          @Override
+          public void onResponse(InnerRespT innerResp) {
+            outerObserver.onResponse(responseTransformer.apply(innerResp));
+          }
+
+          @Override
+          public void onError(Throwable throwable) {
+            outerObserver.onError(throwable);
+          }
+
+          @Override
+          public void onComplete() {
+            outerObserver.onComplete();
+          }
+        },
+        apiCallContext);
+  }
+}

From 9b47fdf3a4bf8f2b8b03eef94a086190f80d7754 Mon Sep 17 00:00:00 2001
From: Igor Bernstein <igorbernstein@google.com>
Date: Mon, 4 Nov 2024 17:27:27 -0500
Subject: [PATCH 02/16] test + fixes

Change-Id: Id4c56656a829f5f4c7ab1170f5f980cf3cc3760c
---
 .../stub/BigtableUnaryOperationCallable.java  |   6 +-
 .../data/v2/stub/EnhancedBigtableStub.java    |   2 +
 .../BigtableUnaryOperationCallableTest.java   | 149 ++++++++++++++++++
 3 files changed, 154 insertions(+), 3 deletions(-)
 create mode 100644 google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallableTest.java

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java
index 202498987d..4d9181c41e 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java
@@ -47,6 +47,7 @@
 class BigtableUnaryOperationCallable<ReqT, RespT> extends UnaryCallable<ReqT, RespT> {
   private static final Logger LOGGER =
       Logger.getLogger(BigtableUnaryOperationCallable.class.getName());
+  Logger logger = LOGGER;
 
   private final ServerStreamingCallable<ReqT, RespT> inner;
   private final ApiCallContext defaultCallContext;
@@ -125,8 +126,8 @@ public void onResponse(RespT resp) {
         String msg =
             String.format(
                 "Received multiple responses for a %s unary operation. Previous: %s, New: %s",
-                spanName, resp, response);
-        LOGGER.log(Level.WARNING, msg);
+                spanName, response, resp);
+        logger.log(Level.WARNING, msg);
 
         InternalException error =
             new InternalException(msg, null, GrpcStatusCode.of(Status.Code.INTERNAL), false);
@@ -140,7 +141,6 @@ public void onResponse(RespT resp) {
 
       responseReceived = true;
       this.response = resp;
-      this.tracer.responseReceived();
     }
 
     @Override
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
index 6e5b1234c4..402f63b145 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
@@ -1408,6 +1408,8 @@ private <BaseReqT, BaseRespT, ReqT, RespT> UnaryCallable<ReqT, RespT> createUnar
     ServerStreamingCallable<ReqT, RespT> transformed =
         new TransformingServerStreamingCallable<>(base, requestTransformer, responseTranformer);
 
+    transformed = new BigtableTracerStreamingCallable<>(transformed);
+
     return new BigtableUnaryOperationCallable<>(
         transformed,
         clientContext.getDefaultCallContext(),
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallableTest.java
new file mode 100644
index 0000000000..b4def0085a
--- /dev/null
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallableTest.java
@@ -0,0 +1,149 @@
+package com.google.cloud.bigtable.data.v2.stub;
+
+import com.google.api.core.ApiFuture;
+import com.google.api.gax.grpc.GrpcCallContext;
+import com.google.api.gax.rpc.InternalException;
+import com.google.api.gax.tracing.ApiTracerFactory;
+import com.google.api.gax.tracing.SpanName;
+import com.google.cloud.bigtable.data.v2.stub.metrics.BigtableTracer;
+import com.google.cloud.bigtable.gaxx.testing.FakeStreamingApi;
+import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCall;
+import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCallable;
+import com.google.common.collect.ImmutableList;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+
+import java.util.concurrent.ExecutionException;
+import java.util.logging.Logger;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+
+@RunWith(JUnit4.class)
+public class BigtableUnaryOperationCallableTest {
+    @Rule
+    public final MockitoRule mockitoRule = MockitoJUnit.rule();
+
+    @Mock
+    private ApiTracerFactory tracerFactory;
+    @Mock
+    private BigtableTracer tracer;
+
+    @Before
+    public void setUp() throws Exception {
+        Mockito.when(tracerFactory.newTracer(Mockito.any(), Mockito.any(), Mockito.any()))
+                .thenReturn(tracer);
+    }
+
+    @Test
+    public void testFutureResolve() throws Exception {
+        BigtableUnaryOperationCallable<String, String> callable = new BigtableUnaryOperationCallable<>(
+                new FakeStreamingApi.ServerStreamingStashCallable<>(ImmutableList.of("value")),
+                GrpcCallContext.createDefault(),
+                tracerFactory,
+                SpanName.of("Fake", "method"),
+                false
+        );
+
+        ApiFuture<String> f = callable.futureCall("fake");
+        assertThat(f.get()).isEqualTo("value");
+    }
+
+    @Test
+    public void testMultipleResponses() throws Exception {
+        MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
+
+        BigtableUnaryOperationCallable<String, String> callable = new BigtableUnaryOperationCallable<>(
+                inner,
+                GrpcCallContext.createDefault(),
+                tracerFactory,
+                SpanName.of("Fake", "method"),
+                false
+        );
+        callable.logger = Mockito.mock(Logger.class);
+
+        ApiFuture<String> f = callable.futureCall("fake");
+        MockServerStreamingCall<String, String> call = inner.popLastCall();
+        call.getController().getObserver().onResponse("first");
+        call.getController().getObserver().onResponse("second");
+
+        Throwable e = Assert.assertThrows(ExecutionException.class, f::get).getCause();
+        assertThat(e).isInstanceOf(InternalException.class);
+        assertThat(e).hasMessageThat().contains("Received multiple responses for a Fake.method unary operation. Previous: first, New: second");
+
+        ArgumentCaptor<String> msgCaptor = ArgumentCaptor.forClass(String.class);
+        verify(callable.logger).log(Mockito.any(), msgCaptor.capture());
+        assertThat(msgCaptor.getValue()).isEqualTo("Received multiple responses for a Fake.method unary operation. Previous: first, New: second");
+
+        assertThat(call.getController().isCancelled()).isTrue();
+
+    }
+
+    @Test
+    public void testCancel() {
+        MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
+        BigtableUnaryOperationCallable<String, String> callable = new BigtableUnaryOperationCallable<>(
+                inner,
+                GrpcCallContext.createDefault(),
+                tracerFactory,
+                SpanName.of("Fake", "method"),
+                false
+        );
+        ApiFuture<String> f = callable.futureCall("req");
+        f.cancel(true);
+
+        MockServerStreamingCall<String, String> call = inner.popLastCall();
+        assertThat(call.getController().isCancelled()).isTrue();
+    }
+
+    @Test
+    public void testMissingResponse() {
+        MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
+        BigtableUnaryOperationCallable<String, String> callable = new BigtableUnaryOperationCallable<>(
+                inner,
+                GrpcCallContext.createDefault(),
+                tracerFactory,
+                SpanName.of("Fake", "method"),
+                false
+        );
+        ApiFuture<String> f = callable.futureCall("req");
+        MockServerStreamingCall<String, String> call = inner.popLastCall();
+        call.getController().getObserver().onComplete();
+
+        Throwable cause = Assert.assertThrows(ExecutionException.class, f::get).getCause();
+        assertThat(cause).hasMessageThat().isEqualTo("Fake.method unary operation completed without a response message");
+    }
+
+    @Test
+    public void testTracing() throws Exception {
+        MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
+        BigtableUnaryOperationCallable<String, String> callable = new BigtableUnaryOperationCallable<>(
+                inner,
+                GrpcCallContext.createDefault(),
+                tracerFactory,
+                SpanName.of("Fake", "method"),
+                false
+        );
+        ApiFuture<String> f = callable.futureCall("req");
+        MockServerStreamingCall<String, String> call = inner.popLastCall();
+        call.getController().getObserver().onResponse("value");
+        call.getController().getObserver().onComplete();
+
+        f.get();
+        verify(tracer).responseReceived();
+        verify(tracer).operationSucceeded();
+
+        // afterResponse is the responsibility of BigtableTracerStreamingCallable
+        verify(tracer, never()).afterResponse(Mockito.anyLong());
+    }
+}
\ No newline at end of file

From 0e8ae3aceabaa403789f5f1d67e4ac1fc826532e Mon Sep 17 00:00:00 2001
From: cloud-java-bot <cloud-java-bot@google.com>
Date: Mon, 4 Nov 2024 22:32:45 +0000
Subject: [PATCH 03/16] chore: generate libraries at Mon Nov  4 22:30:01 UTC
 2024

---
 .../BigtableUnaryOperationCallableTest.java   | 248 +++++++++---------
 1 file changed, 125 insertions(+), 123 deletions(-)

diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallableTest.java
index b4def0085a..d509fc1c5d 100644
--- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallableTest.java
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallableTest.java
@@ -1,5 +1,9 @@
 package com.google.cloud.bigtable.data.v2.stub;
 
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+
 import com.google.api.core.ApiFuture;
 import com.google.api.gax.grpc.GrpcCallContext;
 import com.google.api.gax.rpc.InternalException;
@@ -10,6 +14,8 @@
 import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCall;
 import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCallable;
 import com.google.common.collect.ImmutableList;
+import java.util.concurrent.ExecutionException;
+import java.util.logging.Logger;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Rule;
@@ -22,128 +28,124 @@
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
 
-import java.util.concurrent.ExecutionException;
-import java.util.logging.Logger;
-
-import static com.google.common.truth.Truth.assertThat;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.verify;
-
 @RunWith(JUnit4.class)
 public class BigtableUnaryOperationCallableTest {
-    @Rule
-    public final MockitoRule mockitoRule = MockitoJUnit.rule();
-
-    @Mock
-    private ApiTracerFactory tracerFactory;
-    @Mock
-    private BigtableTracer tracer;
-
-    @Before
-    public void setUp() throws Exception {
-        Mockito.when(tracerFactory.newTracer(Mockito.any(), Mockito.any(), Mockito.any()))
-                .thenReturn(tracer);
-    }
-
-    @Test
-    public void testFutureResolve() throws Exception {
-        BigtableUnaryOperationCallable<String, String> callable = new BigtableUnaryOperationCallable<>(
-                new FakeStreamingApi.ServerStreamingStashCallable<>(ImmutableList.of("value")),
-                GrpcCallContext.createDefault(),
-                tracerFactory,
-                SpanName.of("Fake", "method"),
-                false
-        );
-
-        ApiFuture<String> f = callable.futureCall("fake");
-        assertThat(f.get()).isEqualTo("value");
-    }
-
-    @Test
-    public void testMultipleResponses() throws Exception {
-        MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
-
-        BigtableUnaryOperationCallable<String, String> callable = new BigtableUnaryOperationCallable<>(
-                inner,
-                GrpcCallContext.createDefault(),
-                tracerFactory,
-                SpanName.of("Fake", "method"),
-                false
-        );
-        callable.logger = Mockito.mock(Logger.class);
-
-        ApiFuture<String> f = callable.futureCall("fake");
-        MockServerStreamingCall<String, String> call = inner.popLastCall();
-        call.getController().getObserver().onResponse("first");
-        call.getController().getObserver().onResponse("second");
-
-        Throwable e = Assert.assertThrows(ExecutionException.class, f::get).getCause();
-        assertThat(e).isInstanceOf(InternalException.class);
-        assertThat(e).hasMessageThat().contains("Received multiple responses for a Fake.method unary operation. Previous: first, New: second");
-
-        ArgumentCaptor<String> msgCaptor = ArgumentCaptor.forClass(String.class);
-        verify(callable.logger).log(Mockito.any(), msgCaptor.capture());
-        assertThat(msgCaptor.getValue()).isEqualTo("Received multiple responses for a Fake.method unary operation. Previous: first, New: second");
-
-        assertThat(call.getController().isCancelled()).isTrue();
-
-    }
-
-    @Test
-    public void testCancel() {
-        MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
-        BigtableUnaryOperationCallable<String, String> callable = new BigtableUnaryOperationCallable<>(
-                inner,
-                GrpcCallContext.createDefault(),
-                tracerFactory,
-                SpanName.of("Fake", "method"),
-                false
-        );
-        ApiFuture<String> f = callable.futureCall("req");
-        f.cancel(true);
-
-        MockServerStreamingCall<String, String> call = inner.popLastCall();
-        assertThat(call.getController().isCancelled()).isTrue();
-    }
-
-    @Test
-    public void testMissingResponse() {
-        MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
-        BigtableUnaryOperationCallable<String, String> callable = new BigtableUnaryOperationCallable<>(
-                inner,
-                GrpcCallContext.createDefault(),
-                tracerFactory,
-                SpanName.of("Fake", "method"),
-                false
-        );
-        ApiFuture<String> f = callable.futureCall("req");
-        MockServerStreamingCall<String, String> call = inner.popLastCall();
-        call.getController().getObserver().onComplete();
-
-        Throwable cause = Assert.assertThrows(ExecutionException.class, f::get).getCause();
-        assertThat(cause).hasMessageThat().isEqualTo("Fake.method unary operation completed without a response message");
-    }
-
-    @Test
-    public void testTracing() throws Exception {
-        MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
-        BigtableUnaryOperationCallable<String, String> callable = new BigtableUnaryOperationCallable<>(
-                inner,
-                GrpcCallContext.createDefault(),
-                tracerFactory,
-                SpanName.of("Fake", "method"),
-                false
-        );
-        ApiFuture<String> f = callable.futureCall("req");
-        MockServerStreamingCall<String, String> call = inner.popLastCall();
-        call.getController().getObserver().onResponse("value");
-        call.getController().getObserver().onComplete();
-
-        f.get();
-        verify(tracer).responseReceived();
-        verify(tracer).operationSucceeded();
-
-        // afterResponse is the responsibility of BigtableTracerStreamingCallable
-        verify(tracer, never()).afterResponse(Mockito.anyLong());
-    }
-}
\ No newline at end of file
+  @Rule public final MockitoRule mockitoRule = MockitoJUnit.rule();
+
+  @Mock private ApiTracerFactory tracerFactory;
+  @Mock private BigtableTracer tracer;
+
+  @Before
+  public void setUp() throws Exception {
+    Mockito.when(tracerFactory.newTracer(Mockito.any(), Mockito.any(), Mockito.any()))
+        .thenReturn(tracer);
+  }
+
+  @Test
+  public void testFutureResolve() throws Exception {
+    BigtableUnaryOperationCallable<String, String> callable =
+        new BigtableUnaryOperationCallable<>(
+            new FakeStreamingApi.ServerStreamingStashCallable<>(ImmutableList.of("value")),
+            GrpcCallContext.createDefault(),
+            tracerFactory,
+            SpanName.of("Fake", "method"),
+            false);
+
+    ApiFuture<String> f = callable.futureCall("fake");
+    assertThat(f.get()).isEqualTo("value");
+  }
+
+  @Test
+  public void testMultipleResponses() throws Exception {
+    MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
+
+    BigtableUnaryOperationCallable<String, String> callable =
+        new BigtableUnaryOperationCallable<>(
+            inner,
+            GrpcCallContext.createDefault(),
+            tracerFactory,
+            SpanName.of("Fake", "method"),
+            false);
+    callable.logger = Mockito.mock(Logger.class);
+
+    ApiFuture<String> f = callable.futureCall("fake");
+    MockServerStreamingCall<String, String> call = inner.popLastCall();
+    call.getController().getObserver().onResponse("first");
+    call.getController().getObserver().onResponse("second");
+
+    Throwable e = Assert.assertThrows(ExecutionException.class, f::get).getCause();
+    assertThat(e).isInstanceOf(InternalException.class);
+    assertThat(e)
+        .hasMessageThat()
+        .contains(
+            "Received multiple responses for a Fake.method unary operation. Previous: first, New: second");
+
+    ArgumentCaptor<String> msgCaptor = ArgumentCaptor.forClass(String.class);
+    verify(callable.logger).log(Mockito.any(), msgCaptor.capture());
+    assertThat(msgCaptor.getValue())
+        .isEqualTo(
+            "Received multiple responses for a Fake.method unary operation. Previous: first, New: second");
+
+    assertThat(call.getController().isCancelled()).isTrue();
+  }
+
+  @Test
+  public void testCancel() {
+    MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
+    BigtableUnaryOperationCallable<String, String> callable =
+        new BigtableUnaryOperationCallable<>(
+            inner,
+            GrpcCallContext.createDefault(),
+            tracerFactory,
+            SpanName.of("Fake", "method"),
+            false);
+    ApiFuture<String> f = callable.futureCall("req");
+    f.cancel(true);
+
+    MockServerStreamingCall<String, String> call = inner.popLastCall();
+    assertThat(call.getController().isCancelled()).isTrue();
+  }
+
+  @Test
+  public void testMissingResponse() {
+    MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
+    BigtableUnaryOperationCallable<String, String> callable =
+        new BigtableUnaryOperationCallable<>(
+            inner,
+            GrpcCallContext.createDefault(),
+            tracerFactory,
+            SpanName.of("Fake", "method"),
+            false);
+    ApiFuture<String> f = callable.futureCall("req");
+    MockServerStreamingCall<String, String> call = inner.popLastCall();
+    call.getController().getObserver().onComplete();
+
+    Throwable cause = Assert.assertThrows(ExecutionException.class, f::get).getCause();
+    assertThat(cause)
+        .hasMessageThat()
+        .isEqualTo("Fake.method unary operation completed without a response message");
+  }
+
+  @Test
+  public void testTracing() throws Exception {
+    MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
+    BigtableUnaryOperationCallable<String, String> callable =
+        new BigtableUnaryOperationCallable<>(
+            inner,
+            GrpcCallContext.createDefault(),
+            tracerFactory,
+            SpanName.of("Fake", "method"),
+            false);
+    ApiFuture<String> f = callable.futureCall("req");
+    MockServerStreamingCall<String, String> call = inner.popLastCall();
+    call.getController().getObserver().onResponse("value");
+    call.getController().getObserver().onComplete();
+
+    f.get();
+    verify(tracer).responseReceived();
+    verify(tracer).operationSucceeded();
+
+    // afterResponse is the responsibility of BigtableTracerStreamingCallable
+    verify(tracer, never()).afterResponse(Mockito.anyLong());
+  }
+}

From 5f535976b14b8e773da62713429afa6a14f6bece Mon Sep 17 00:00:00 2001
From: Igor Bernstein <igorbernstein@google.com>
Date: Mon, 4 Nov 2024 17:33:05 -0500
Subject: [PATCH 04/16] oops

Change-Id: I1bd8c318b3272925cd6b81601d7b1d7c772a853f
---
 .../data/v2/stub/EnhancedBigtableStub.java        |  2 --
 .../stub/BigtableUnaryOperationCallableTest.java  | 15 +++++++++++++++
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
index 402f63b145..6e5b1234c4 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
@@ -1408,8 +1408,6 @@ private <BaseReqT, BaseRespT, ReqT, RespT> UnaryCallable<ReqT, RespT> createUnar
     ServerStreamingCallable<ReqT, RespT> transformed =
         new TransformingServerStreamingCallable<>(base, requestTransformer, responseTranformer);
 
-    transformed = new BigtableTracerStreamingCallable<>(transformed);
-
     return new BigtableUnaryOperationCallable<>(
         transformed,
         clientContext.getDefaultCallContext(),
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallableTest.java
index b4def0085a..62f03f36f4 100644
--- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallableTest.java
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallableTest.java
@@ -1,3 +1,18 @@
+/*
+ * Copyright 2024 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 com.google.api.core.ApiFuture;

From e2bb1834e3e0bf224d170101568d798788580007 Mon Sep 17 00:00:00 2001
From: Igor Bernstein <igorbernstein@google.com>
Date: Mon, 4 Nov 2024 17:49:28 -0500
Subject: [PATCH 05/16] more tests

Change-Id: I1c45f2058cadc1acb9c6abd87222be9eb233778c
---
 ...ansformingServerStreamingCallableTest.java | 74 +++++++++++++++++++
 1 file changed, 74 insertions(+)
 create mode 100644 google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/TransformingServerStreamingCallableTest.java

diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/TransformingServerStreamingCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/TransformingServerStreamingCallableTest.java
new file mode 100644
index 0000000000..856d732f5c
--- /dev/null
+++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/TransformingServerStreamingCallableTest.java
@@ -0,0 +1,74 @@
+/*
+ * Copyright 2024 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.common.truth.Truth.assertThat;
+
+import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockResponseObserver;
+import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCall;
+import com.google.cloud.bigtable.gaxx.testing.MockStreamingApi.MockServerStreamingCallable;
+import com.google.common.base.Functions;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class TransformingServerStreamingCallableTest {
+  @Test
+  public void testReqTransform() {
+    MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
+    TransformingServerStreamingCallable<Integer, String, String, String> xform =
+        new TransformingServerStreamingCallable<>(inner, Object::toString, Functions.identity());
+
+    MockResponseObserver<String> responseObserver = new MockResponseObserver<>(true);
+    xform.call(37, responseObserver);
+
+    MockServerStreamingCall<String, String> call = inner.popLastCall();
+    assertThat(call.getRequest()).isEqualTo("37");
+  }
+
+  @Test
+  public void testRespTransform() {
+    MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
+    TransformingServerStreamingCallable<String, Integer, String, String> xform =
+        new TransformingServerStreamingCallable<>(inner, Functions.identity(), Integer::parseInt);
+
+    MockResponseObserver<Integer> outerObserver = new MockResponseObserver<>(true);
+    xform.call("req", outerObserver);
+
+    MockServerStreamingCall<String, String> call = inner.popLastCall();
+    call.getController().getObserver().onResponse("37");
+
+    assertThat(outerObserver.popNextResponse()).isEqualTo(37);
+  }
+
+  @Test
+  public void testError() {
+    MockServerStreamingCallable<String, String> inner = new MockServerStreamingCallable<>();
+    TransformingServerStreamingCallable<String, String, String, String> xform =
+        new TransformingServerStreamingCallable<>(
+            inner, Functions.identity(), Functions.identity());
+
+    MockResponseObserver<String> outerObserver = new MockResponseObserver<>(true);
+    xform.call("req", outerObserver);
+
+    MockServerStreamingCall<String, String> call = inner.popLastCall();
+    RuntimeException e = new RuntimeException("fake error");
+    call.getController().getObserver().onError(e);
+
+    assertThat(outerObserver.getFinalError()).isEqualTo(e);
+  }
+}

From 42cf8d7944ada501d87e3c9c0b738d12371907f9 Mon Sep 17 00:00:00 2001
From: Igor Bernstein <igorbernstein@google.com>
Date: Mon, 4 Nov 2024 19:21:32 -0500
Subject: [PATCH 06/16] avoid multiple cancels

Change-Id: I4e05efaac6ae60f5827c6d666c3c6f6cebebaa54
---
 .../v2/stub/BigtableUnaryOperationCallable.java     | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java
index 4d9181c41e..aa46062a9a 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java
@@ -29,6 +29,8 @@
 import com.google.api.gax.tracing.SpanName;
 import com.google.common.base.Preconditions;
 import io.grpc.Status;
+
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
@@ -88,6 +90,7 @@ class UnaryFuture extends AbstractApiFuture<RespT> implements ResponseObserver<R
     private final boolean allowNoResponse;
 
     private StreamController controller;
+    private final AtomicBoolean upstreamCancelled = new AtomicBoolean();
     private boolean responseReceived;
     private @Nullable RespT response;
 
@@ -112,12 +115,18 @@ public void onStart(StreamController controller) {
     @Override
     public boolean cancel(boolean mayInterruptIfRunning) {
       if (super.cancel(mayInterruptIfRunning)) {
-        controller.cancel();
+        cancelUpstream();
         return true;
       }
       return false;
     }
 
+    private void cancelUpstream() {
+      if (upstreamCancelled.compareAndSet(false, true)) {
+        controller.cancel();
+      }
+    }
+
     @Override
     public void onResponse(RespT resp) {
       tracer.responseReceived();
@@ -135,7 +144,7 @@ public void onResponse(RespT resp) {
           tracer.operationFailed(error);
         }
 
-        controller.cancel();
+        cancelUpstream();
         return;
       }
 

From e2c4391354533fe59975b03a38990147d45386e4 Mon Sep 17 00:00:00 2001
From: cloud-java-bot <cloud-java-bot@google.com>
Date: Tue, 5 Nov 2024 00:26:39 +0000
Subject: [PATCH 07/16] chore: generate libraries at Tue Nov  5 00:23:44 UTC
 2024

---
 .../bigtable/data/v2/stub/BigtableUnaryOperationCallable.java    | 1 -
 1 file changed, 1 deletion(-)

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java
index aa46062a9a..59e8451967 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java
@@ -29,7 +29,6 @@
 import com.google.api.gax.tracing.SpanName;
 import com.google.common.base.Preconditions;
 import io.grpc.Status;
-
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.logging.Level;
 import java.util.logging.Logger;

From a84d67ed3fa13bb3a80a451d7a42c23f4352af73 Mon Sep 17 00:00:00 2001
From: Igor Bernstein <igorbernstein@google.com>
Date: Mon, 4 Nov 2024 19:42:01 -0500
Subject: [PATCH 08/16] fix fallback

Change-Id: I654e70f0b34f5d4c3071ba3c2fed64ea183a865e
---
 .../data/v2/stub/EnhancedBigtableStub.java    | 22 ++++++++++++++-----
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
index 6e5b1234c4..fda54079a9 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
@@ -126,6 +126,7 @@
 import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsRetryingCallable;
 import com.google.cloud.bigtable.data.v2.stub.readrows.FilterMarkerRowsCallable;
 import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsBatchingDescriptor;
+import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsFirstCallable;
 import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsResumptionStrategy;
 import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsRetryCompletedCallable;
 import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsUserCallable;
@@ -580,12 +581,21 @@ public void call(
           }
         };
 
-    return new BigtableUnaryOperationCallable<>(
-        readRowCallable,
-        clientContext.getDefaultCallContext(),
-        clientContext.getTracerFactory(),
-        getSpanName("ReadRow"),
-        /*allowNoResponses=*/ true);
+    if (EnhancedBigtableStubSettings.SKIP_TRAILERS) {
+      return new BigtableUnaryOperationCallable<>(
+              readRowCallable,
+              clientContext.getDefaultCallContext(),
+              clientContext.getTracerFactory(),
+              getSpanName("ReadRow"),
+              /*allowNoResponses=*/ true);
+    } else {
+      readRowCallable = new ReadRowsUserCallable<>(readRowsCallable, requestContext);
+      UnaryCallable<Query, RowT> firstRow = new ReadRowsFirstCallable<>(readRowCallable);
+      return new TracedUnaryCallable<>(firstRow, clientContext.getTracerFactory(), getSpanName("ReadRow"))
+        .withDefaultCallContext(clientContext.getDefaultCallContext());
+    }
+
+
   }
 
   /**

From c6ff6727e6b9a6b364b5d47fcf4046082297b7d8 Mon Sep 17 00:00:00 2001
From: cloud-java-bot <cloud-java-bot@google.com>
Date: Tue, 5 Nov 2024 00:45:23 +0000
Subject: [PATCH 09/16] chore: generate libraries at Tue Nov  5 00:42:37 UTC
 2024

---
 .../data/v2/stub/EnhancedBigtableStub.java      | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
index fda54079a9..8e8aa6c2f5 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
@@ -583,19 +583,18 @@ public void call(
 
     if (EnhancedBigtableStubSettings.SKIP_TRAILERS) {
       return new BigtableUnaryOperationCallable<>(
-              readRowCallable,
-              clientContext.getDefaultCallContext(),
-              clientContext.getTracerFactory(),
-              getSpanName("ReadRow"),
-              /*allowNoResponses=*/ true);
+          readRowCallable,
+          clientContext.getDefaultCallContext(),
+          clientContext.getTracerFactory(),
+          getSpanName("ReadRow"),
+          /*allowNoResponses=*/ true);
     } else {
       readRowCallable = new ReadRowsUserCallable<>(readRowsCallable, requestContext);
       UnaryCallable<Query, RowT> firstRow = new ReadRowsFirstCallable<>(readRowCallable);
-      return new TracedUnaryCallable<>(firstRow, clientContext.getTracerFactory(), getSpanName("ReadRow"))
-        .withDefaultCallContext(clientContext.getDefaultCallContext());
+      return new TracedUnaryCallable<>(
+              firstRow, clientContext.getTracerFactory(), getSpanName("ReadRow"))
+          .withDefaultCallContext(clientContext.getDefaultCallContext());
     }
-
-
   }
 
   /**

From e4c36b4ccc0af04abec7caa5c8fdedc7a4d8bad2 Mon Sep 17 00:00:00 2001
From: Igor Bernstein <igorbernstein@google.com>
Date: Mon, 4 Nov 2024 19:51:17 -0500
Subject: [PATCH 10/16] proper fallback

Change-Id: Ic0106f3c6983edbb032aeba6e107e4324952397d
---
 .../data/v2/stub/EnhancedBigtableStub.java    | 48 ++++++++++---------
 1 file changed, 25 insertions(+), 23 deletions(-)

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
index fda54079a9..5ee43c8208 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
@@ -571,31 +571,33 @@ public <RowT> UnaryCallable<Query, RowT> createReadRowCallable(RowAdapter<RowT>
                 .build(),
             rowAdapter);
 
-    ServerStreamingCallable<Query, RowT> readRowCallable =
-        new ServerStreamingCallable<Query, RowT>() {
-          @Override
-          public void call(
-              Query query, ResponseObserver<RowT> responseObserver, ApiCallContext apiCallContext) {
-            readRowsCallable.call(
-                query.limit(1).toProto(requestContext), responseObserver, apiCallContext);
-          }
-        };
-
-    if (EnhancedBigtableStubSettings.SKIP_TRAILERS) {
-      return new BigtableUnaryOperationCallable<>(
-              readRowCallable,
-              clientContext.getDefaultCallContext(),
-              clientContext.getTracerFactory(),
-              getSpanName("ReadRow"),
-              /*allowNoResponses=*/ true);
+    if (!EnhancedBigtableStubSettings.SKIP_TRAILERS) {
+      ReadRowsUserCallable<RowT> readRowCallable =
+          new ReadRowsUserCallable<>(readRowsCallable, requestContext);
+      ReadRowsFirstCallable<RowT> firstRow = new ReadRowsFirstCallable<>(readRowCallable);
+      UnaryCallable<Query, RowT> traced =
+          new TracedUnaryCallable<>(
+              firstRow, clientContext.getTracerFactory(), getSpanName("ReadRow"));
+      return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
     } else {
-      readRowCallable = new ReadRowsUserCallable<>(readRowsCallable, requestContext);
-      UnaryCallable<Query, RowT> firstRow = new ReadRowsFirstCallable<>(readRowCallable);
-      return new TracedUnaryCallable<>(firstRow, clientContext.getTracerFactory(), getSpanName("ReadRow"))
-        .withDefaultCallContext(clientContext.getDefaultCallContext());
+      ServerStreamingCallable<Query, RowT> readRowCallable =
+          new ServerStreamingCallable<Query, RowT>() {
+            @Override
+            public void call(
+                Query query,
+                ResponseObserver<RowT> responseObserver,
+                ApiCallContext apiCallContext) {
+              readRowsCallable.call(
+                  query.limit(1).toProto(requestContext), responseObserver, apiCallContext);
+            }
+          };
+      return new BigtableUnaryOperationCallable<>(
+          readRowCallable,
+          clientContext.getDefaultCallContext(),
+          clientContext.getTracerFactory(),
+          getSpanName("ReadRow"),
+          /*allowNoResponses=*/ true);
     }
-
-
   }
 
   /**

From 99b6044d6c9570aa17d759d4eb33fad24daccd05 Mon Sep 17 00:00:00 2001
From: Igor Bernstein <igorbernstein@google.com>
Date: Mon, 4 Nov 2024 21:27:54 -0500
Subject: [PATCH 11/16] Use transforming callable

Change-Id: I8d8474050e40cd819d3be2a5b251448f6eb8c94f
---
 .../data/v2/stub/EnhancedBigtableStub.java      | 17 ++++++-----------
 1 file changed, 6 insertions(+), 11 deletions(-)

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
index 5ee43c8208..3809fc704f 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
@@ -46,7 +46,6 @@
 import com.google.api.gax.rpc.Callables;
 import com.google.api.gax.rpc.ClientContext;
 import com.google.api.gax.rpc.RequestParamsExtractor;
-import com.google.api.gax.rpc.ResponseObserver;
 import com.google.api.gax.rpc.ServerStreamingCallSettings;
 import com.google.api.gax.rpc.ServerStreamingCallable;
 import com.google.api.gax.rpc.StatusCode.Code;
@@ -138,6 +137,7 @@
 import com.google.cloud.bigtable.gaxx.retrying.ApiResultRetryAlgorithm;
 import com.google.cloud.bigtable.gaxx.retrying.RetryInfoRetryAlgorithm;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Functions;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
@@ -581,16 +581,11 @@ public <RowT> UnaryCallable<Query, RowT> createReadRowCallable(RowAdapter<RowT>
       return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
     } else {
       ServerStreamingCallable<Query, RowT> readRowCallable =
-          new ServerStreamingCallable<Query, RowT>() {
-            @Override
-            public void call(
-                Query query,
-                ResponseObserver<RowT> responseObserver,
-                ApiCallContext apiCallContext) {
-              readRowsCallable.call(
-                  query.limit(1).toProto(requestContext), responseObserver, apiCallContext);
-            }
-          };
+          new TransformingServerStreamingCallable<>(
+              readRowsCallable,
+              (query) -> query.limit(1).toProto(requestContext),
+              Functions.identity());
+
       return new BigtableUnaryOperationCallable<>(
           readRowCallable,
           clientContext.getDefaultCallContext(),

From b3ddb672ce10613761be29c5c1dd8fdb80d394fb Mon Sep 17 00:00:00 2001
From: Igor Bernstein <igorbernstein@google.com>
Date: Mon, 4 Nov 2024 21:28:02 -0500
Subject: [PATCH 12/16] fix npe

Change-Id: Ib589ca063369e26ef214eb89099e459981dafe83
---
 .../bigtable/data/v2/stub/EnhancedBigtableStubSettings.java     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java
index 37c18952ee..863389166f 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java
@@ -110,7 +110,7 @@ public class EnhancedBigtableStubSettings extends StubSettings<EnhancedBigtableS
       Boolean.parseBoolean(System.getenv("CBT_ENABLE_DIRECTPATH"));
 
   static final boolean SKIP_TRAILERS =
-      Optional.of(System.getenv("CBT_SKIP_HEADERS"))
+      Optional.ofNullable(System.getenv("CBT_SKIP_HEADERS"))
           .map(Boolean::parseBoolean)
           .orElse(DIRECT_PATH_ENABLED);
 

From 14b97b4d450d0f501ac9cada10da71688df85860 Mon Sep 17 00:00:00 2001
From: Igor Bernstein <igorbernstein@google.com>
Date: Mon, 4 Nov 2024 21:28:49 -0500
Subject: [PATCH 13/16] clean up logic

Change-Id: I4504c47143000d97554a96469d5f3fd368d08ef1
---
 .../stub/BigtableUnaryOperationCallable.java  | 50 +++++++++----------
 1 file changed, 24 insertions(+), 26 deletions(-)

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java
index 59e8451967..19f7a5224c 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableUnaryOperationCallable.java
@@ -130,25 +130,26 @@ private void cancelUpstream() {
     public void onResponse(RespT resp) {
       tracer.responseReceived();
 
-      if (responseReceived) {
-        String msg =
-            String.format(
-                "Received multiple responses for a %s unary operation. Previous: %s, New: %s",
-                spanName, response, resp);
-        logger.log(Level.WARNING, msg);
-
-        InternalException error =
-            new InternalException(msg, null, GrpcStatusCode.of(Status.Code.INTERNAL), false);
-        if (setException(error)) {
-          tracer.operationFailed(error);
-        }
-
-        cancelUpstream();
+      // happy path - buffer the only responsse
+      if (!responseReceived) {
+        responseReceived = true;
+        this.response = resp;
         return;
       }
 
-      responseReceived = true;
-      this.response = resp;
+      String msg =
+          String.format(
+              "Received multiple responses for a %s unary operation. Previous: %s, New: %s",
+              spanName, response, resp);
+      logger.log(Level.WARNING, msg);
+
+      InternalException error =
+          new InternalException(msg, null, GrpcStatusCode.of(Status.Code.INTERNAL), false);
+      if (setException(error)) {
+        tracer.operationFailed(error);
+      }
+
+      cancelUpstream();
     }
 
     @Override
@@ -163,28 +164,25 @@ public void onError(Throwable throwable) {
 
     @Override
     public void onComplete() {
-      // Sanity check server response counts
-      if (!allowNoResponse && !responseReceived) {
+      if (allowNoResponse || responseReceived) {
+        if (set(response)) {
+          tracer.operationSucceeded();
+          return;
+        }
+      } else {
         String msg = spanName + " unary operation completed without a response message";
         InternalException e =
             new InternalException(msg, null, GrpcStatusCode.of(Status.Code.INTERNAL), false);
 
         if (setException(e)) {
           tracer.operationFailed(e);
+          return;
         }
-        return;
       }
 
       // check cancellation race
       if (isCancelled()) {
         tracer.operationCancelled();
-        return;
-      }
-
-      // try to resolve the future
-      if (this.set(response)) {
-        tracer.operationSucceeded();
-        return;
       }
     }
   }

From 35272b6b3ea16bab99ae0e16f8eb4a8d3069f9f7 Mon Sep 17 00:00:00 2001
From: Igor Bernstein <igorbernstein@google.com>
Date: Tue, 5 Nov 2024 10:56:47 -0500
Subject: [PATCH 14/16] oops, messed up splitting commits, this should've been
 part of this pr not the next

Change-Id: I16a35e19c50b7b7b855f4299cf41f0607b3e90bd
---
 .../data/v2/stub/EnhancedBigtableStub.java    | 42 +++++++++++++------
 1 file changed, 29 insertions(+), 13 deletions(-)

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
index 3809fc704f..4d26c3329d 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
@@ -42,6 +42,7 @@
 import com.google.api.gax.retrying.RetryingExecutorWithContext;
 import com.google.api.gax.retrying.ScheduledRetryingExecutor;
 import com.google.api.gax.retrying.SimpleStreamResumptionStrategy;
+import com.google.api.gax.retrying.StreamResumptionStrategy;
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.Callables;
 import com.google.api.gax.rpc.ClientContext;
@@ -562,16 +563,16 @@ public <RowT> ServerStreamingCallable<Query, RowT> createReadRowsCallable(
    * </ul>
    */
   public <RowT> UnaryCallable<Query, RowT> createReadRowCallable(RowAdapter<RowT> rowAdapter) {
-    ServerStreamingCallable<ReadRowsRequest, RowT> readRowsCallable =
-        createReadRowsBaseCallable(
-            ServerStreamingCallSettings.<ReadRowsRequest, Row>newBuilder()
-                .setRetryableCodes(settings.readRowSettings().getRetryableCodes())
-                .setRetrySettings(settings.readRowSettings().getRetrySettings())
-                .setIdleTimeout(settings.readRowSettings().getRetrySettings().getTotalTimeout())
-                .build(),
-            rowAdapter);
-
     if (!EnhancedBigtableStubSettings.SKIP_TRAILERS) {
+      ServerStreamingCallable<ReadRowsRequest, RowT> readRowsCallable =
+          createReadRowsBaseCallable(
+              ServerStreamingCallSettings.<ReadRowsRequest, Row>newBuilder()
+                  .setRetryableCodes(settings.readRowSettings().getRetryableCodes())
+                  .setRetrySettings(settings.readRowSettings().getRetrySettings())
+                  .setIdleTimeout(settings.readRowSettings().getRetrySettings().getTotalTimeout())
+                  .build(),
+              rowAdapter);
+
       ReadRowsUserCallable<RowT> readRowCallable =
           new ReadRowsUserCallable<>(readRowsCallable, requestContext);
       ReadRowsFirstCallable<RowT> firstRow = new ReadRowsFirstCallable<>(readRowCallable);
@@ -580,6 +581,15 @@ public <RowT> UnaryCallable<Query, RowT> createReadRowCallable(RowAdapter<RowT>
               firstRow, clientContext.getTracerFactory(), getSpanName("ReadRow"));
       return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
     } else {
+      ServerStreamingCallable<ReadRowsRequest, RowT> readRowsCallable =
+          createReadRowsBaseCallable(
+              ServerStreamingCallSettings.<ReadRowsRequest, Row>newBuilder()
+                  .setRetryableCodes(settings.readRowSettings().getRetryableCodes())
+                  .setRetrySettings(settings.readRowSettings().getRetrySettings())
+                  .setIdleTimeout(settings.readRowSettings().getRetrySettings().getTotalTimeout())
+                  .build(),
+              rowAdapter,
+              new SimpleStreamResumptionStrategy<>());
       ServerStreamingCallable<Query, RowT> readRowCallable =
           new TransformingServerStreamingCallable<>(
               readRowsCallable,
@@ -595,6 +605,11 @@ public <RowT> UnaryCallable<Query, RowT> createReadRowCallable(RowAdapter<RowT>
     }
   }
 
+  private <ReqT, RowT> ServerStreamingCallable<ReadRowsRequest, RowT> createReadRowsBaseCallable(
+      ServerStreamingCallSettings<ReqT, Row> readRowsSettings, RowAdapter<RowT> rowAdapter) {
+    return createReadRowsBaseCallable(
+        readRowsSettings, rowAdapter, new ReadRowsResumptionStrategy<RowT>(rowAdapter));
+  }
   /**
    * Creates a callable chain to handle ReadRows RPCs. The chain will:
    *
@@ -611,8 +626,9 @@ public <RowT> UnaryCallable<Query, RowT> createReadRowCallable(RowAdapter<RowT>
    * <p>NOTE: the caller is responsible for adding tracing & metrics.
    */
   private <ReqT, RowT> ServerStreamingCallable<ReadRowsRequest, RowT> createReadRowsBaseCallable(
-      ServerStreamingCallSettings<ReqT, Row> readRowsSettings, RowAdapter<RowT> rowAdapter) {
-
+      ServerStreamingCallSettings<ReqT, Row> readRowsSettings,
+      RowAdapter<RowT> rowAdapter,
+      StreamResumptionStrategy<ReadRowsRequest, RowT> resumptionStrategy) {
     ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> base =
         GrpcRawCallableFactory.createServerStreamingCallable(
             GrpcCallSettings.<ReadRowsRequest, ReadRowsResponse>newBuilder()
@@ -653,7 +669,7 @@ public Map<String, String> extract(ReadRowsRequest readRowsRequest) {
     // ReadRowsRequest -> ReadRowsResponse callable).
     ServerStreamingCallSettings<ReadRowsRequest, RowT> innerSettings =
         ServerStreamingCallSettings.<ReadRowsRequest, RowT>newBuilder()
-            .setResumptionStrategy(new ReadRowsResumptionStrategy<>(rowAdapter))
+            .setResumptionStrategy(resumptionStrategy)
             .setRetryableCodes(readRowsSettings.getRetryableCodes())
             .setRetrySettings(readRowsSettings.getRetrySettings())
             .setIdleTimeout(readRowsSettings.getIdleTimeout())
@@ -1344,7 +1360,7 @@ private <BaseReqT, BaseRespT, ReqT, RespT> UnaryCallable<ReqT, RespT> createUnar
       return createUnaryCallableNew(
           methodDescriptor, headerParamsFn, callSettings, requestTransformer, responseTranformer);
     } else {
-      return createUnaryCallableNew(
+      return createUnaryCallableOld(
           methodDescriptor, headerParamsFn, callSettings, requestTransformer, responseTranformer);
     }
   }

From 0c2bc8414b478af4eef6650548d13cc13d5c2245 Mon Sep 17 00:00:00 2001
From: Igor Bernstein <igorbernstein@google.com>
Date: Tue, 5 Nov 2024 11:55:36 -0500
Subject: [PATCH 15/16] typo

Change-Id: I8202e935975e1a55606265c502fe7573b8a4acb0
---
 .../cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java       | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
index 4d26c3329d..580f0a1ae9 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
@@ -1426,7 +1426,7 @@ private <BaseReqT, BaseRespT, ReqT, RespT> UnaryCallable<ReqT, RespT> createUnar
 
     base = new BigtableTracerStreamingCallable<>(base);
 
-    base = withRetries(base, convertUnaryToServerStreamingSettings(settings.mutateRowSettings()));
+    base = withRetries(base, convertUnaryToServerStreamingSettings(callSettings));
 
     ServerStreamingCallable<ReqT, RespT> transformed =
         new TransformingServerStreamingCallable<>(base, requestTransformer, responseTranformer);

From 066ae08e82282af308ba2ebaebd50f01fb3844eb Mon Sep 17 00:00:00 2001
From: Igor Bernstein <igorbernstein@google.com>
Date: Tue, 5 Nov 2024 12:02:51 -0500
Subject: [PATCH 16/16] disable watchdog for the new ReadRow callable chain

Change-Id: I4522719a65f24d27fb9dccde031c3b1cc04042c2
---
 .../cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java      | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
index 580f0a1ae9..a1100231fa 100644
--- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
+++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
@@ -586,7 +586,8 @@ public <RowT> UnaryCallable<Query, RowT> createReadRowCallable(RowAdapter<RowT>
               ServerStreamingCallSettings.<ReadRowsRequest, Row>newBuilder()
                   .setRetryableCodes(settings.readRowSettings().getRetryableCodes())
                   .setRetrySettings(settings.readRowSettings().getRetrySettings())
-                  .setIdleTimeout(settings.readRowSettings().getRetrySettings().getTotalTimeout())
+                  .setIdleTimeoutDuration(Duration.ZERO)
+                  .setWaitTimeoutDuration(Duration.ZERO)
                   .build(),
               rowAdapter,
               new SimpleStreamResumptionStrategy<>());