From 6e7d29b8a7cad5e9bf76e802a273c33bf4460124 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Thu, 12 Nov 2026 14:49:50 -0500 Subject: [PATCH 1/5] fix(bigtable): enforce StreamController cancellation and flow control in MaybePointReadCallable --- .../stub/readrows/MaybePointReadCallable.java | 54 +++++-- .../readrows/MaybePointReadCallableTest.java | 141 +++++++++++++++++- 2 files changed, 178 insertions(+), 17 deletions(-) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java index 5530cd346612..4ffb12038d01 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java @@ -19,13 +19,16 @@ import com.google.api.core.ApiFutureCallback; import com.google.api.core.ApiFutures; import com.google.api.core.InternalApi; +import com.google.api.core.SettableApiFuture; 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 com.google.api.gax.rpc.UnaryCallable; import com.google.cloud.bigtable.data.v2.models.Query; +import com.google.common.base.Preconditions; import com.google.common.util.concurrent.MoreExecutors; +import java.util.concurrent.CancellationException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; @@ -52,7 +55,12 @@ public void call(Query request, ResponseObserver responseObserver, ApiCall return; } + // TODO: move StreamController and callback state to an inner class so hasStarted and + // autoFlowControl can be plain booleans instead of AtomicBooleans. + AtomicBoolean hasStarted = new AtomicBoolean(); + AtomicBoolean autoFlowControl = new AtomicBoolean(true); AtomicBoolean cancelled = new AtomicBoolean(); + SettableApiFuture demand = SettableApiFuture.create(); AtomicReference> futureRef = new AtomicReference<>(); responseObserver.onStart( @@ -60,6 +68,7 @@ public void call(Query request, ResponseObserver responseObserver, ApiCall @Override public void cancel() { cancelled.set(true); + demand.cancel(false); ApiFuture f = futureRef.get(); if (f != null) { f.cancel(false); @@ -67,19 +76,42 @@ public void cancel() { } @Override - public void disableAutoInboundFlowControl() {} + public void disableAutoInboundFlowControl() { + Preconditions.checkState( + !hasStarted.get(), + "Can't disable automatic flow control after the stream has started."); + autoFlowControl.set(false); + } @Override - public void request(int count) {} + public void request(int count) { + Preconditions.checkState(!autoFlowControl.get(), "Automatic flow control is enabled."); + Preconditions.checkArgument(count > 0, "Count must be > 0"); + demand.set(null); + } }); + hasStarted.set(true); + if (autoFlowControl.get()) { + demand.set(null); + } + + // Start the underlying RPC eagerly on the caller thread (rather than waiting for demand) to: + // 1. Preserve the caller thread's ambient io.grpc.Context (deadlines, trace spans). + // 2. Match GAX/gRPC inbound flow-control semantics, where the request is sent immediately, + // RPC errors surface even when demand is 0, and request(n) only gates inbound row delivery. + // Chaining via transformAsync also lets future.cancel(false) propagate cancellation to both the + // in-flight RPC future and the demand future. ApiFuture future; try { - future = pointReader.futureCall(request, context); + future = + ApiFutures.transformAsync( + pointReader.futureCall(request, context), + row -> ApiFutures.transform(demand, unused -> row, MoreExecutors.directExecutor()), + MoreExecutors.directExecutor()); } catch (Throwable t) { - if (!cancelled.get()) { - responseObserver.onError(t); - } + responseObserver.onError( + cancelled.get() ? new CancellationException("User cancelled stream") : t); return; } futureRef.set(future); @@ -92,9 +124,6 @@ public void request(int count) {} new ApiFutureCallback() { @Override public void onSuccess(RowT row) { - if (cancelled.get()) { - return; - } if (row != null) { try { responseObserver.onResponse(row); @@ -102,15 +131,16 @@ public void onSuccess(RowT row) { responseObserver.onError(t); return; } + if (cancelled.get()) { + responseObserver.onError(new CancellationException("User cancelled stream")); + return; + } } responseObserver.onComplete(); } @Override public void onFailure(Throwable t) { - if (cancelled.get()) { - return; - } responseObserver.onError(t); } }, diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallableTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallableTest.java index 42a6c28a9d88..cfc34ac99210 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallableTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallableTest.java @@ -16,11 +16,13 @@ package com.google.cloud.bigtable.data.v2.stub.readrows; import static com.google.common.truth.Truth.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; import com.google.api.core.ApiFuture; import com.google.api.core.SettableApiFuture; import com.google.api.gax.rpc.ApiCallContext; import com.google.api.gax.rpc.ResponseObserver; +import com.google.api.gax.rpc.ServerStream; import com.google.api.gax.rpc.StreamController; import com.google.api.gax.rpc.UnaryCallable; import com.google.cloud.bigtable.data.v2.models.Query; @@ -33,6 +35,13 @@ import com.google.protobuf.ByteString; import java.util.ArrayList; import java.util.List; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -132,14 +141,14 @@ public void pointReaderFails_propagatesErrorToObserver() { } @Test - public void observerCancel_cancelsFutureAndSuppressesError() { + public void observerCancel_cancelsFutureAndPropagatesCancellation() { Query query = Query.create(TABLE_ID).rowKey("a"); callable.call(query, observer, null); observer.controller.cancel(); assertThat(pointReader.response.isCancelled()).isTrue(); - assertThat(observer.error).isNull(); + assertThat(observer.error).isInstanceOf(CancellationException.class); assertThat(observer.completed).isFalse(); } @@ -151,7 +160,8 @@ public void cancelBeforeFutureReturns_cancelsAfterFutureAttaches() { callable.call(query, observer, null); assertThat(pointReader.response.isCancelled()).isTrue(); - assertThat(observer.error).isNull(); + assertThat(observer.error).isInstanceOf(CancellationException.class); + assertThat(observer.completed).isFalse(); } @Test @@ -168,14 +178,135 @@ public void futureCallThrows_routesThroughOnError() { } @Test - public void futureCallThrowsAfterCancel_suppressesError() { + public void futureCallThrowsAfterCancel_propagatesCancellation() { Query query = Query.create(TABLE_ID).rowKey("a"); pointReader.syncFailure = new RuntimeException("sync boom"); pointReader.onCall = () -> observer.controller.cancel(); callable.call(query, observer, null); - assertThat(observer.error).isNull(); + assertThat(observer.error).isInstanceOf(CancellationException.class); + assertThat(observer.completed).isFalse(); + } + + @Test + public void cancelInsideOnResponse_deliversCancellationAndNotOnComplete() { + Query query = Query.create(TABLE_ID).rowKey("a"); + RecordingObserver cancellingObserver = + new RecordingObserver() { + @Override + public void onResponse(Row response) { + super.onResponse(response); + controller.cancel(); + } + }; + + callable.call(query, cancellingObserver, null); + pointReader.response.set(ROW_A); + + assertThat(cancellingObserver.responses).containsExactly(ROW_A); + assertThat(cancellingObserver.completed).isFalse(); + assertThat(cancellingObserver.error).isInstanceOf(CancellationException.class); + } + + @Test + public void serverStreamCancel_unblocksBlockedHasNext() throws Exception { + Query query = Query.create(TABLE_ID).rowKey("a"); + ServerStream stream = callable.call(query); + + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + CountDownLatch entered = new CountDownLatch(1); + Future hasNextFuture = + executor.submit( + () -> { + entered.countDown(); + return stream.iterator().hasNext(); + }); + assertThat(entered.await(5, TimeUnit.SECONDS)).isTrue(); + Thread.sleep(20); + + stream.cancel(); + + ExecutionException thrown = + assertThrows(ExecutionException.class, () -> hasNextFuture.get(5, TimeUnit.SECONDS)); + assertThat(thrown.getCause()).isInstanceOf(CancellationException.class); + } finally { + executor.shutdownNow(); + } + } + + @Test + public void manualFlowControl_gatesDeliveryUntilRequested() { + Query query = Query.create(TABLE_ID).rowKey("a"); + RecordingObserver manualObserver = + new RecordingObserver() { + @Override + public void onStart(StreamController controller) { + super.onStart(controller); + controller.disableAutoInboundFlowControl(); + } + }; + + callable.call(query, manualObserver, null); + pointReader.response.set(ROW_A); + + // Demand is 0, so row must not be delivered yet + assertThat(manualObserver.responses).isEmpty(); + assertThat(manualObserver.completed).isFalse(); + assertThat(manualObserver.error).isNull(); + + manualObserver.controller.request(1); + + assertThat(manualObserver.responses).containsExactly(ROW_A); + assertThat(manualObserver.completed).isTrue(); + assertThat(manualObserver.error).isNull(); + } + + @Test + public void manualFlowControl_cancelBeforeRequestDeliversCancellationWithoutRow() { + Query query = Query.create(TABLE_ID).rowKey("a"); + RecordingObserver manualObserver = + new RecordingObserver() { + @Override + public void onStart(StreamController controller) { + super.onStart(controller); + controller.disableAutoInboundFlowControl(); + } + }; + + callable.call(query, manualObserver, null); + pointReader.response.set(ROW_A); + + manualObserver.controller.cancel(); + + assertThat(manualObserver.responses).isEmpty(); + assertThat(manualObserver.completed).isFalse(); + assertThat(manualObserver.error).isInstanceOf(CancellationException.class); + } + + @Test + public void flowControlPreconditions_enforced() { + Query query = Query.create(TABLE_ID).rowKey("a"); + callable.call(query, observer, null); + + // Cannot disable auto flow control after onStart has finished + assertThrows( + IllegalStateException.class, () -> observer.controller.disableAutoInboundFlowControl()); + // Cannot call request(int) while auto flow control is enabled + assertThrows(IllegalStateException.class, () -> observer.controller.request(1)); + + RecordingObserver manualObserver = + new RecordingObserver() { + @Override + public void onStart(StreamController controller) { + super.onStart(controller); + controller.disableAutoInboundFlowControl(); + } + }; + callable.call(query, manualObserver, null); + assertThrows(IllegalArgumentException.class, () -> manualObserver.controller.request(0)); + assertThrows(IllegalArgumentException.class, () -> manualObserver.controller.request(-1)); } private static class FakePointReader extends UnaryCallable { From a9a592d6b0fc8f238f27fa908cc8509b0a4ad01a Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Thu, 12 Nov 2026 15:29:53 -0500 Subject: [PATCH 2/5] fix(bigtable): capture CancellationException in cancel() and check at onSuccess entry --- .../stub/readrows/MaybePointReadCallable.java | 33 ++++++++++++------- .../readrows/MaybePointReadCallableTest.java | 13 ++++++++ 2 files changed, 34 insertions(+), 12 deletions(-) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java index 4ffb12038d01..1b4a26f39167 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java @@ -59,7 +59,7 @@ public void call(Query request, ResponseObserver responseObserver, ApiCall // autoFlowControl can be plain booleans instead of AtomicBooleans. AtomicBoolean hasStarted = new AtomicBoolean(); AtomicBoolean autoFlowControl = new AtomicBoolean(true); - AtomicBoolean cancelled = new AtomicBoolean(); + AtomicReference cancellation = new AtomicReference<>(); SettableApiFuture demand = SettableApiFuture.create(); AtomicReference> futureRef = new AtomicReference<>(); @@ -67,11 +67,13 @@ public void call(Query request, ResponseObserver responseObserver, ApiCall new StreamController() { @Override public void cancel() { - cancelled.set(true); - demand.cancel(false); - ApiFuture f = futureRef.get(); - if (f != null) { - f.cancel(false); + if (cancellation.compareAndSet( + null, new CancellationException("User cancelled stream"))) { + demand.cancel(false); + ApiFuture f = futureRef.get(); + if (f != null) { + f.cancel(false); + } } } @@ -110,12 +112,12 @@ public void request(int count) { row -> ApiFutures.transform(demand, unused -> row, MoreExecutors.directExecutor()), MoreExecutors.directExecutor()); } catch (Throwable t) { - responseObserver.onError( - cancelled.get() ? new CancellationException("User cancelled stream") : t); + CancellationException e = cancellation.get(); + responseObserver.onError(e != null ? e : t); return; } futureRef.set(future); - if (cancelled.get()) { + if (cancellation.get() != null) { future.cancel(false); } @@ -124,6 +126,11 @@ public void request(int count) { new ApiFutureCallback() { @Override public void onSuccess(RowT row) { + CancellationException e = cancellation.get(); + if (e != null) { + responseObserver.onError(e); + return; + } if (row != null) { try { responseObserver.onResponse(row); @@ -131,8 +138,9 @@ public void onSuccess(RowT row) { responseObserver.onError(t); return; } - if (cancelled.get()) { - responseObserver.onError(new CancellationException("User cancelled stream")); + e = cancellation.get(); + if (e != null) { + responseObserver.onError(e); return; } } @@ -141,7 +149,8 @@ public void onSuccess(RowT row) { @Override public void onFailure(Throwable t) { - responseObserver.onError(t); + CancellationException e = cancellation.get(); + responseObserver.onError(e != null ? e : t); } }, MoreExecutors.directExecutor()); diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallableTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallableTest.java index cfc34ac99210..19e2f4176b7e 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallableTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallableTest.java @@ -164,6 +164,19 @@ public void cancelBeforeFutureReturns_cancelsAfterFutureAttaches() { assertThat(observer.completed).isFalse(); } + @Test + public void cancelWhenFutureAlreadyCompleted_propagatesCancellationWithoutResponse() { + Query query = Query.create(TABLE_ID).rowKey("a"); + pointReader.response.set(ROW_A); + pointReader.onCall = () -> observer.controller.cancel(); + + callable.call(query, observer, null); + + assertThat(observer.responses).isEmpty(); + assertThat(observer.error).isInstanceOf(CancellationException.class); + assertThat(observer.completed).isFalse(); + } + @Test public void futureCallThrows_routesThroughOnError() { Query query = Query.create(TABLE_ID).rowKey("a"); From 386cd381aa32934531db9d28c329cf643e851def Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Thu, 12 Nov 2026 15:56:08 -0500 Subject: [PATCH 3/5] chore(bigtable): use MoreObjects.firstNonNull for cancellation fallback --- .../data/v2/stub/readrows/MaybePointReadCallable.java | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java index 1b4a26f39167..402d5aecb646 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java @@ -26,6 +26,7 @@ import com.google.api.gax.rpc.StreamController; import com.google.api.gax.rpc.UnaryCallable; import com.google.cloud.bigtable.data.v2.models.Query; +import com.google.common.base.MoreObjects; import com.google.common.base.Preconditions; import com.google.common.util.concurrent.MoreExecutors; import java.util.concurrent.CancellationException; @@ -112,8 +113,7 @@ public void request(int count) { row -> ApiFutures.transform(demand, unused -> row, MoreExecutors.directExecutor()), MoreExecutors.directExecutor()); } catch (Throwable t) { - CancellationException e = cancellation.get(); - responseObserver.onError(e != null ? e : t); + responseObserver.onError(MoreObjects.firstNonNull(cancellation.get(), t)); return; } futureRef.set(future); @@ -149,8 +149,7 @@ public void onSuccess(RowT row) { @Override public void onFailure(Throwable t) { - CancellationException e = cancellation.get(); - responseObserver.onError(e != null ? e : t); + responseObserver.onError(MoreObjects.firstNonNull(cancellation.get(), t)); } }, MoreExecutors.directExecutor()); From decef54c4074bb9983f1a6f39f562274f5bdea2d Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Thu, 12 Nov 2026 22:21:47 +0000 Subject: [PATCH 4/5] fix(bigtable): complete immediately without waiting for demand when row is null --- .../stub/readrows/MaybePointReadCallable.java | 5 ++++- .../readrows/MaybePointReadCallableTest.java | 20 +++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java index 402d5aecb646..8ba338d2f30b 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java @@ -110,7 +110,10 @@ public void request(int count) { future = ApiFutures.transformAsync( pointReader.futureCall(request, context), - row -> ApiFutures.transform(demand, unused -> row, MoreExecutors.directExecutor()), + row -> + row == null + ? ApiFutures.immediateFuture(null) + : ApiFutures.transform(demand, unused -> row, MoreExecutors.directExecutor()), MoreExecutors.directExecutor()); } catch (Throwable t) { responseObserver.onError(MoreObjects.firstNonNull(cancellation.get(), t)); diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallableTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallableTest.java index 19e2f4176b7e..a7105f4d987e 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallableTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallableTest.java @@ -276,6 +276,26 @@ public void onStart(StreamController controller) { assertThat(manualObserver.error).isNull(); } + @Test + public void manualFlowControl_nullRowCompletesWithoutWaitingForDemand() { + Query query = Query.create(TABLE_ID).rowKey("missing"); + RecordingObserver manualObserver = + new RecordingObserver() { + @Override + public void onStart(StreamController controller) { + super.onStart(controller); + controller.disableAutoInboundFlowControl(); + } + }; + + callable.call(query, manualObserver, null); + pointReader.response.set(null); + + assertThat(manualObserver.responses).isEmpty(); + assertThat(manualObserver.completed).isTrue(); + assertThat(manualObserver.error).isNull(); + } + @Test public void manualFlowControl_cancelBeforeRequestDeliversCancellationWithoutRow() { Query query = Query.create(TABLE_ID).rowKey("a"); From 8ce8487ecc37af9042dd2716e1d6967548d43e2a Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Fri, 13 Nov 2026 15:49:26 +0000 Subject: [PATCH 5/5] chore(bigtable): extract waitForDemand helper in MaybePointReadCallable --- .../v2/stub/readrows/MaybePointReadCallable.java | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java index 8ba338d2f30b..21cf7d75a910 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/MaybePointReadCallable.java @@ -110,10 +110,7 @@ public void request(int count) { future = ApiFutures.transformAsync( pointReader.futureCall(request, context), - row -> - row == null - ? ApiFutures.immediateFuture(null) - : ApiFutures.transform(demand, unused -> row, MoreExecutors.directExecutor()), + row -> waitForDemand(row, demand), MoreExecutors.directExecutor()); } catch (Throwable t) { responseObserver.onError(MoreObjects.firstNonNull(cancellation.get(), t)); @@ -157,4 +154,12 @@ public void onFailure(Throwable t) { }, MoreExecutors.directExecutor()); } + + private static ApiFuture waitForDemand(RowT row, ApiFuture demand) { + if (row == null) { + return ApiFutures.immediateFuture(null); + } + // Wait for demand to complete, then forward the row. + return ApiFutures.transform(demand, unused -> row, MoreExecutors.directExecutor()); + } }