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 79cfee703a98..cf8b65684ebe 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 @@ -21,7 +21,6 @@ import static com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants.INSTANCE_ID_KEY; import com.google.api.core.ApiFuture; -import com.google.api.core.ApiFutures; import com.google.api.core.BetaApi; import com.google.api.core.InternalApi; import com.google.api.gax.batching.Batcher; @@ -112,7 +111,6 @@ import com.google.cloud.bigtable.data.v2.stub.readrows.FilterMarkerRowsCallable; import com.google.cloud.bigtable.data.v2.stub.readrows.LargeReadRowsResumptionStrategy; 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; @@ -131,7 +129,6 @@ import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import com.google.common.util.concurrent.MoreExecutors; import com.google.protobuf.ByteString; import io.grpc.MethodDescriptor; import io.opencensus.stats.Stats; @@ -404,52 +401,31 @@ public ServerStreamingCallable createReadRowsCallable( */ public UnaryCallable createReadRowCallable(RowAdapter rowAdapter) { ClientContext clientContext = bigtableClientContext.getClientContext(); - if (!settings.getEnableSkipTrailers()) { - ServerStreamingCallable readRowsCallable = - createReadRowsBaseCallable( - ServerStreamingCallSettings.newBuilder() - .setRetryableCodes(settings.readRowSettings().getRetryableCodes()) - .setRetrySettings(settings.readRowSettings().getRetrySettings()) - .setIdleTimeout(settings.readRowSettings().getRetrySettings().getTotalTimeout()) - .build(), - rowAdapter); - - ReadRowsUserCallable readRowCallable = - new ReadRowsUserCallable<>(readRowsCallable, requestContext); - ReadRowsFirstCallable firstRow = new ReadRowsFirstCallable<>(readRowCallable); - UnaryCallable traced = - new TracedUnaryCallable<>( - firstRow, clientContext.getTracerFactory(), getSpanName("ReadRow")); - return traced.withDefaultCallContext( - clientContext - .getDefaultCallContext() - .withRetrySettings(settings.readRowSettings().getRetrySettings())); - } else { - ServerStreamingCallable readRowsCallable = - createReadRowsBaseCallable( - ServerStreamingCallSettings.newBuilder() - .setRetryableCodes(settings.readRowSettings().getRetryableCodes()) - .setRetrySettings(settings.readRowSettings().getRetrySettings()) - .setIdleTimeoutDuration(Duration.ZERO) - .setWaitTimeoutDuration(Duration.ZERO) - .build(), - rowAdapter, - new SimpleStreamResumptionStrategy<>()); - ServerStreamingCallable readRowCallable = - new TransformingServerStreamingCallable<>( - readRowsCallable, - (query) -> query.limit(1).toProto(requestContext), - Functions.identity()); - - return new BigtableUnaryOperationCallable<>( - readRowCallable, - clientContext - .getDefaultCallContext() - .withRetrySettings(settings.readRowSettings().getRetrySettings()), - clientContext.getTracerFactory(), - getSpanName("ReadRow"), - /* allowNoResponses= */ true); - } + + ServerStreamingCallable readRowsCallable = + createReadRowsBaseCallable( + ServerStreamingCallSettings.newBuilder() + .setRetryableCodes(settings.readRowSettings().getRetryableCodes()) + .setRetrySettings(settings.readRowSettings().getRetrySettings()) + .setIdleTimeoutDuration(Duration.ZERO) + .setWaitTimeoutDuration(Duration.ZERO) + .build(), + rowAdapter, + new SimpleStreamResumptionStrategy<>()); + ServerStreamingCallable readRowCallable = + new TransformingServerStreamingCallable<>( + readRowsCallable, + (query) -> query.limit(1).toProto(requestContext), + Functions.identity()); + + return new BigtableUnaryOperationCallable<>( + readRowCallable, + clientContext + .getDefaultCallContext() + .withRetrySettings(settings.readRowSettings().getRetrySettings()), + clientContext.getTracerFactory(), + getSpanName("ReadRow"), + /* allowNoResponses= */ true); } private ServerStreamingCallable createReadRowsBaseCallable( @@ -1292,67 +1268,6 @@ private UnaryCallable createUnar UnaryCallSettings callSettings, Function requestTransformer, Function responseTranformer) { - if (settings.getEnableSkipTrailers()) { - return createUnaryCallableNew( - methodDescriptor, headerParamsFn, callSettings, requestTransformer, responseTranformer); - } else { - return createUnaryCallableOld( - methodDescriptor, headerParamsFn, callSettings, requestTransformer, responseTranformer); - } - } - - private UnaryCallable createUnaryCallableOld( - MethodDescriptor methodDescriptor, - RequestParamsExtractor headerParamsFn, - UnaryCallSettings callSettings, - Function requestTransformer, - Function responseTranformer) { - - UnaryCallable base = - GrpcRawCallableFactory.createUnaryCallable( - GrpcCallSettings.newBuilder() - .setMethodDescriptor(methodDescriptor) - .setParamsExtractor(headerParamsFn) - .build(), - callSettings.getRetryableCodes()); - - UnaryCallable withStatsHeaders = new StatsHeadersUnaryCallable<>(base); - - UnaryCallable withBigtableTracer = - new BigtableTracerUnaryCallable<>(withStatsHeaders); - - UnaryCallable retrying = withRetries(withBigtableTracer, callSettings); - - UnaryCallable transformed = - new UnaryCallable() { - @Override - public ApiFuture futureCall(ReqT reqT, ApiCallContext apiCallContext) { - ApiFuture f = - retrying.futureCall(requestTransformer.apply(reqT), apiCallContext); - return ApiFutures.transform( - f, responseTranformer::apply, MoreExecutors.directExecutor()); - } - }; - - UnaryCallable traced = - new TracedUnaryCallable<>( - transformed, - bigtableClientContext.getClientContext().getTracerFactory(), - getSpanName(methodDescriptor.getBareMethodName())); - - return traced.withDefaultCallContext( - bigtableClientContext - .getClientContext() - .getDefaultCallContext() - .withRetrySettings(callSettings.getRetrySettings())); - } - - private UnaryCallable createUnaryCallableNew( - MethodDescriptor methodDescriptor, - RequestParamsExtractor headerParamsFn, - UnaryCallSettings callSettings, - Function requestTransformer, - Function responseTranformer) { ServerStreamingCallable base = GrpcRawCallableFactory.createServerStreamingCallable( 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 acd3323957cd..d1fe259ea1e3 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 @@ -67,7 +67,6 @@ import java.util.Collections; import java.util.List; import java.util.Map; -import java.util.Optional; import java.util.Set; import java.util.concurrent.ScheduledExecutorService; import java.util.logging.Logger; @@ -119,11 +118,6 @@ public class EnhancedBigtableStubSettings extends StubSettings IDEMPOTENT_RETRY_CODES = ImmutableSet.of(Code.DEADLINE_EXCEEDED, Code.UNAVAILABLE); @@ -268,7 +262,6 @@ public class EnhancedBigtableStubSettings extends StubSettings primedTableIds; private final boolean enableRoutingCookie; private final boolean enableRetryInfo; - private final boolean enableSkipTrailers; private final ServerStreamingCallSettings readRowsSettings; private final UnaryCallSettings readRowSettings; @@ -318,7 +311,6 @@ private EnhancedBigtableStubSettings(Builder builder) { primedTableIds = builder.primedTableIds; enableRoutingCookie = builder.enableRoutingCookie; enableRetryInfo = builder.enableRetryInfo; - enableSkipTrailers = builder.enableSkipTrailers; metricsProvider = builder.metricsProvider; metricsEndpoint = builder.metricsEndpoint; internalMetricsProvider = builder.internalMetricsProvider; @@ -413,10 +405,6 @@ public boolean getEnableRetryInfo() { return enableRetryInfo; } - boolean getEnableSkipTrailers() { - return enableSkipTrailers; - } - /** * Gets the Google Cloud Monitoring endpoint for publishing client side metrics. If it's null, * client will publish metrics to the default monitoring endpoint. @@ -769,7 +757,6 @@ public static class Builder extends StubSettings.Builder readRowsSettings; private final UnaryCallSettings.Builder readRowSettings; @@ -810,7 +797,6 @@ private Builder() { setCredentialsProvider(defaultCredentialsProviderBuilder().build()); this.enableRoutingCookie = true; this.enableRetryInfo = true; - this.enableSkipTrailers = SKIP_TRAILERS; metricsProvider = DefaultMetricsProvider.INSTANCE; this.internalMetricsProvider = DEFAULT_INTERNAL_OTEL_PROVIDER; this.jwtAudience = DEFAULT_DATA_JWT_AUDIENCE; @@ -1223,11 +1209,6 @@ public boolean getEnableRetryInfo() { return enableRetryInfo; } - Builder setEnableSkipTrailers(boolean enabled) { - this.enableSkipTrailers = enabled; - return this; - } - /** Returns the builder for the settings used for calls to readRows. */ public ServerStreamingCallSettings.Builder readRowsSettings() { return readRowsSettings; @@ -1360,7 +1341,6 @@ public String toString() { .add("primedTableIds", primedTableIds) .add("enableRoutingCookie", enableRoutingCookie) .add("enableRetryInfo", enableRetryInfo) - .add("enableSkipTrailers", enableSkipTrailers) .add("readRowsSettings", readRowsSettings) .add("readRowSettings", readRowSettings) .add("sampleRowKeysSettings", sampleRowKeysSettings) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallable.java deleted file mode 100644 index 2ef26605b478..000000000000 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallable.java +++ /dev/null @@ -1,79 +0,0 @@ -/* - * Copyright 2018 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.readrows; - -import com.google.api.core.ApiFuture; -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.ServerStreamingCallable; -import com.google.api.gax.rpc.StateCheckingResponseObserver; -import com.google.api.gax.rpc.StreamController; -import com.google.api.gax.rpc.UnaryCallable; -import com.google.cloud.bigtable.data.v2.models.Query; - -/** - * Enhancement for `readRowsCallable().first()` to gracefully limit the row count instead of - * cancelling the RPC - */ -@InternalApi -public class ReadRowsFirstCallable extends UnaryCallable { - - private final ServerStreamingCallable inner; - - public ReadRowsFirstCallable(ServerStreamingCallable inner) { - this.inner = inner; - } - - @Override - public ApiFuture futureCall(Query query, ApiCallContext context) { - ReadRowsFirstResponseObserver observer = new ReadRowsFirstResponseObserver<>(); - this.inner.call(query.limit(1), observer, context); - return observer.getFuture(); - } - - private class ReadRowsFirstResponseObserver extends StateCheckingResponseObserver { - private StreamController innerController; - private RowT firstRow; - private SettableApiFuture settableFuture = SettableApiFuture.create(); - - @Override - protected void onStartImpl(StreamController streamController) { - this.innerController = streamController; - } - - @Override - protected void onResponseImpl(RowT response) { - if (firstRow == null) { - this.firstRow = response; - } - } - - @Override - protected void onErrorImpl(Throwable throwable) { - settableFuture.setException(throwable); - } - - @Override - protected void onCompleteImpl() { - settableFuture.set(firstRow); - } - - protected ApiFuture getFuture() { - return settableFuture; - } - } -} diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java index 888dbc0f340b..aecad0cc123a 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettingsTest.java @@ -995,7 +995,6 @@ public void enableRetryInfoFalseValueTest() throws IOException { "primedTableIds", "enableRoutingCookie", "enableRetryInfo", - "enableSkipTrailers", "readRowsSettings", "readRowSettings", "sampleRowKeysSettings", diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/SkipTrailersTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/SkipTrailersTest.java index b05abca34ef4..9759f798c4bd 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/SkipTrailersTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/SkipTrailersTest.java @@ -102,7 +102,7 @@ public void setUp() throws Exception { .setInstanceId(INSTANCE_ID) .setMetricsProvider(NoopMetricsProvider.INSTANCE) .setCredentialsProvider(NoCredentialsProvider.create()); - clientBuilder.stubSettings().setEnableSkipTrailers(true).setTracerFactory(tracerFactory); + clientBuilder.stubSettings().setTracerFactory(tracerFactory); client = BigtableDataClient.create(clientBuilder.build()); } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerTest.java index 864a801e6ee1..df63ff8019d6 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerTest.java @@ -79,9 +79,16 @@ import com.google.protobuf.ByteString; import com.google.protobuf.BytesValue; import com.google.protobuf.StringValue; +import io.grpc.CallOptions; +import io.grpc.Channel; +import io.grpc.ClientCall; +import io.grpc.ClientInterceptor; +import io.grpc.ForwardingClientCall; +import io.grpc.ForwardingClientCallListener; import io.grpc.ForwardingServerCall; import io.grpc.ManagedChannelBuilder; import io.grpc.Metadata; +import io.grpc.MethodDescriptor; import io.grpc.ProxiedSocketAddress; import io.grpc.ProxyDetector; import io.grpc.Server; @@ -160,6 +167,8 @@ public class BuiltinMetricsTracerTest { private DelayProxyDetector delayProxyDetector; + private final OutstandingRpcCounter outstandingRpcCounter = new OutstandingRpcCounter(); + @Before public void setUp() throws Exception { metricReader = InMemoryMetricReader.create(); @@ -272,7 +281,7 @@ public void sendHeaders(Metadata headers) { if (oldConfigurator != null) { builder = oldConfigurator.apply(builder); } - return builder.proxyDetector(delayProxyDetector); + return builder.proxyDetector(delayProxyDetector).intercept(outstandingRpcCounter); }); stubSettingsBuilder.setTransportChannelProvider(channelProvider.build()); EnhancedBigtableStubSettings stubSettings = stubSettingsBuilder.build(); @@ -541,10 +550,11 @@ public void testRetryCount() throws InterruptedException { } @Test - public void testMutateRowAttemptsTagValues() { + public void testMutateRowAttemptsTagValues() throws InterruptedException { stub.mutateRowCallable() .call(RowMutation.create(TABLE, "random-row").setCell("cf", "q", "value")); + outstandingRpcCounter.waitUntilRpcsDone(); MetricData metricData = getMetricData(metricReader, ATTEMPT_LATENCIES_NAME); Attributes expected1 = @@ -724,6 +734,7 @@ public void testQueuedOnChannelUnaryLatencies() throws Exception { Duration proxyDelayPriorTest = delayProxyDetector.getCurrentDelayUsed(); f.get(); + outstandingRpcCounter.waitUntilRpcsDone(); MetricData clientLatency = getMetricData(metricReader, CLIENT_BLOCKING_LATENCIES_NAME); Attributes attributes = @@ -1123,6 +1134,46 @@ public AtomicInteger getResponseCounter() { } } + static class OutstandingRpcCounter implements ClientInterceptor { + private int numOutstandingRpcs = 0; + private final Object lock = new Object(); + + @Override + public ClientCall interceptCall( + MethodDescriptor methodDescriptor, CallOptions callOptions, Channel channel) { + synchronized (lock) { + numOutstandingRpcs++; + } + return new ForwardingClientCall.SimpleForwardingClientCall( + channel.newCall(methodDescriptor, callOptions)) { + @Override + public void start(Listener responseListener, Metadata headers) { + super.start( + new ForwardingClientCallListener.SimpleForwardingClientCallListener( + responseListener) { + @Override + public void onClose(Status status, Metadata trailers) { + super.onClose(status, trailers); + synchronized (lock) { + numOutstandingRpcs--; + lock.notify(); + } + } + }, + headers); + } + }; + } + + void waitUntilRpcsDone() throws InterruptedException { + synchronized (lock) { + while (numOutstandingRpcs > 0) { + lock.wait(); + } + } + } + } + class DelayProxyDetector implements ProxyDetector { private volatile Instant lastProxyDelay = null; diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallableTest.java deleted file mode 100644 index 07cf3478c1f8..000000000000 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/readrows/ReadRowsFirstCallableTest.java +++ /dev/null @@ -1,76 +0,0 @@ -/* - * Copyright 2018 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.readrows; - -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.verify; - -import com.google.api.gax.grpc.GrpcCallContext; -import com.google.api.gax.rpc.ApiCallContext; -import com.google.api.gax.rpc.ResponseObserver; -import com.google.api.gax.rpc.ServerStreamingCallable; -import com.google.cloud.bigtable.data.v2.internal.RequestContext; -import com.google.cloud.bigtable.data.v2.models.Query; -import com.google.cloud.bigtable.data.v2.models.Row; -import com.google.common.truth.Truth; -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.Mockito; -import org.mockito.junit.MockitoJUnit; -import org.mockito.junit.MockitoRule; -import org.mockito.quality.Strictness; - -@RunWith(JUnit4.class) -public class ReadRowsFirstCallableTest { - @Rule public MockitoRule mockitoRule = MockitoJUnit.rule().strictness(Strictness.WARN); - - private static final RequestContext REQUEST_CONTEXT = - RequestContext.create("fake-project", "fake-instance", "fake-profile"); - - private ServerStreamingCallable innerCallable; - private ArgumentCaptor innerQuery; - - @SuppressWarnings("unchecked") - @Before - public void setUp() { - innerCallable = Mockito.mock(ServerStreamingCallable.class); - innerQuery = ArgumentCaptor.forClass(Query.class); - } - - @Test - public void testLimitAdded() { - ReadRowsFirstCallable callable = new ReadRowsFirstCallable<>(innerCallable); - callable.futureCall(Query.create("fake-table"), GrpcCallContext.createDefault()); - verify(innerCallable) - .call(innerQuery.capture(), any(ResponseObserver.class), any(ApiCallContext.class)); - Truth.assertThat(innerQuery.getValue().toProto(REQUEST_CONTEXT)) - .isEqualTo(Query.create("fake-table").limit(1).toProto(REQUEST_CONTEXT)); - } - - @Test - public void testLimitChanged() { - ReadRowsFirstCallable callable = new ReadRowsFirstCallable<>(innerCallable); - callable.futureCall(Query.create("fake-table").limit(10), GrpcCallContext.createDefault()); - verify(innerCallable) - .call(innerQuery.capture(), any(ResponseObserver.class), any(ApiCallContext.class)); - Truth.assertThat(innerQuery.getValue().toProto(REQUEST_CONTEXT)) - .isEqualTo(Query.create("fake-table").limit(1).toProto(REQUEST_CONTEXT)); - } -}