diff --git a/google-cloud-bigtable/clirr-ignored-differences.xml b/google-cloud-bigtable/clirr-ignored-differences.xml index c3a0fa05e1..c9a6f3762c 100644 --- a/google-cloud-bigtable/clirr-ignored-differences.xml +++ b/google-cloud-bigtable/clirr-ignored-differences.xml @@ -34,6 +34,13 @@ com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub * + + + 7005 + com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub + * + * + 7002 diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java index 889d36e383..cef5e58f3a 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java @@ -25,7 +25,6 @@ import com.google.api.gax.batching.Batcher; import com.google.api.gax.grpc.GrpcCallContext; import com.google.api.gax.rpc.ApiExceptions; -import com.google.api.gax.rpc.ClientContext; import com.google.api.gax.rpc.ResponseObserver; import com.google.api.gax.rpc.ServerStream; import com.google.api.gax.rpc.ServerStreamingCallable; @@ -55,6 +54,7 @@ import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement; import com.google.cloud.bigtable.data.v2.models.sql.ResultSet; import com.google.cloud.bigtable.data.v2.models.sql.SqlType; +import com.google.cloud.bigtable.data.v2.stub.BigtableClientContext; import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub; import com.google.cloud.bigtable.data.v2.stub.sql.SqlServerStream; import com.google.common.util.concurrent.MoreExecutors; @@ -186,7 +186,7 @@ public static BigtableDataClient create(BigtableDataSettings settings) throws IO * BigtableDataClientFactory#close()} is called. */ static BigtableDataClient createWithClientContext( - BigtableDataSettings settings, ClientContext context) throws IOException { + BigtableDataSettings settings, BigtableClientContext context) throws IOException { EnhancedBigtableStub stub = EnhancedBigtableStub.createWithClientContext(settings.getStubSettings(), context); return new BigtableDataClient(stub); diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java index cddea20c7d..5628d58417 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java @@ -114,7 +114,8 @@ public BigtableDataClient createDefault() { defaultSettings.getStubSettings(), sharedClientContext.getOpenTelemetry())) .build(); - return BigtableDataClient.createWithClientContext(defaultSettings, clientContext); + return BigtableDataClient.createWithClientContext( + defaultSettings, sharedClientContext.setClientContext(clientContext)); } catch (IOException e) { // Should never happen because the connection has been established already throw new RuntimeException( @@ -141,7 +142,8 @@ public BigtableDataClient createForAppProfile(@Nonnull String appProfileId) thro EnhancedBigtableStub.createBigtableTracerFactory( settings.getStubSettings(), sharedClientContext.getOpenTelemetry())) .build(); - return BigtableDataClient.createWithClientContext(settings, clientContext); + return BigtableDataClient.createWithClientContext( + settings, sharedClientContext.setClientContext(clientContext)); } /** @@ -169,7 +171,8 @@ public BigtableDataClient createForInstance(@Nonnull String projectId, @Nonnull settings.getStubSettings(), sharedClientContext.getOpenTelemetry())) .build(); - return BigtableDataClient.createWithClientContext(settings, clientContext); + return BigtableDataClient.createWithClientContext( + settings, sharedClientContext.setClientContext(clientContext)); } /** @@ -196,6 +199,7 @@ public BigtableDataClient createForInstance( EnhancedBigtableStub.createBigtableTracerFactory( settings.getStubSettings(), sharedClientContext.getOpenTelemetry())) .build(); - return BigtableDataClient.createWithClientContext(settings, clientContext); + return BigtableDataClient.createWithClientContext( + settings, sharedClientContext.setClientContext(clientContext)); } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java index bac1ec4a06..f67ca1b8c0 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java @@ -197,6 +197,11 @@ public ClientContext getClientContext() { return this.clientContext; } + public BigtableClientContext setClientContext(ClientContext clientContext) { + return new BigtableClientContext( + clientContext, openTelemetry, internalOpenTelemetry, metricsProvider); + } + public void close() throws Exception { for (BackgroundResource resource : clientContext.getBackgroundResources()) { resource.close(); 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 5f6b69dea8..1231c4e6dc 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 @@ -27,7 +27,6 @@ import com.google.api.gax.batching.Batcher; import com.google.api.gax.batching.BatcherImpl; import com.google.api.gax.batching.FlowController; -import com.google.api.gax.core.BackgroundResource; import com.google.api.gax.grpc.GaxGrpcProperties; import com.google.api.gax.grpc.GrpcCallContext; import com.google.api.gax.grpc.GrpcCallSettings; @@ -170,7 +169,7 @@ public class EnhancedBigtableStub implements AutoCloseable { private static final String CLIENT_NAME = "Bigtable"; private static final long FLOW_CONTROL_ADJUSTING_INTERVAL_MS = TimeUnit.SECONDS.toMillis(20); private final EnhancedBigtableStubSettings settings; - private final ClientContext clientContext; + private final BigtableClientContext bigtableClientContext; private final boolean closeClientContext; private final RequestContext requestContext; @@ -209,11 +208,13 @@ public static EnhancedBigtableStub create(EnhancedBigtableStubSettings settings) bigtableClientContext.getClientContext().toBuilder() .setTracerFactory(createBigtableTracerFactory(settings, openTelemetry)) .build(); - return new EnhancedBigtableStub(settings, contextWithTracer); + bigtableClientContext = bigtableClientContext.setClientContext(contextWithTracer); + return new EnhancedBigtableStub(settings, bigtableClientContext); } public static EnhancedBigtableStub createWithClientContext( - EnhancedBigtableStubSettings settings, ClientContext clientContext) throws IOException { + EnhancedBigtableStubSettings settings, BigtableClientContext clientContext) + throws IOException { return new EnhancedBigtableStub(settings, clientContext, false); } @@ -289,16 +290,17 @@ static Attributes createBuiltinAttributes(EnhancedBigtableStubSettings settings) "bigtable-java/" + Version.VERSION); } - public EnhancedBigtableStub(EnhancedBigtableStubSettings settings, ClientContext clientContext) { + public EnhancedBigtableStub( + EnhancedBigtableStubSettings settings, BigtableClientContext clientContext) { this(settings, clientContext, true); } public EnhancedBigtableStub( EnhancedBigtableStubSettings settings, - ClientContext clientContext, + BigtableClientContext clientContext, boolean closeClientContext) { this.settings = settings; - this.clientContext = clientContext; + this.bigtableClientContext = clientContext; this.closeClientContext = closeClientContext; this.requestContext = RequestContext.create( @@ -347,7 +349,7 @@ public EnhancedBigtableStub( public ServerStreamingCallable createReadRowsRawCallable( RowAdapter rowAdapter) { return createReadRowsBaseCallable(settings.readRowsSettings(), rowAdapter) - .withDefaultCallContext(clientContext.getDefaultCallContext()); + .withDefaultCallContext(bigtableClientContext.getClientContext().getDefaultCallContext()); } /** @@ -375,10 +377,13 @@ public ServerStreamingCallable createReadRowsCallable( SpanName span = getSpanName("ReadRows"); ServerStreamingCallable traced = new TracedServerStreamingCallable<>( - readRowsUserCallable, clientContext.getTracerFactory(), span); + readRowsUserCallable, + bigtableClientContext.getClientContext().getTracerFactory(), + span); return traced.withDefaultCallContext( - clientContext + bigtableClientContext + .getClientContext() .getDefaultCallContext() .withRetrySettings(settings.readRowsSettings().getRetrySettings())); } @@ -398,6 +403,7 @@ public ServerStreamingCallable createReadRowsCallable( * */ public UnaryCallable createReadRowCallable(RowAdapter rowAdapter) { + ClientContext clientContext = bigtableClientContext.getClientContext(); if (!settings.getEnableSkipTrailers()) { ServerStreamingCallable readRowsCallable = createReadRowsBaseCallable( @@ -506,7 +512,7 @@ private ServerStreamingCallable createReadRo .build(); ServerStreamingCallable watched = - Callables.watched(merging, innerSettings, clientContext); + Callables.watched(merging, innerSettings, bigtableClientContext.getClientContext()); ServerStreamingCallable withBigtableTracer = new BigtableTracerStreamingCallable<>(watched); @@ -582,7 +588,7 @@ public ServerStreamingCallable createSkipLargeRowsCall .build(); ServerStreamingCallable watched = - Callables.watched(merging, innerSettings, clientContext); + Callables.watched(merging, innerSettings, bigtableClientContext.getClientContext()); ServerStreamingCallable withBigtableTracer = new BigtableTracerStreamingCallable<>(watched); @@ -604,10 +610,13 @@ public ServerStreamingCallable createSkipLargeRowsCall SpanName span = getSpanName("ReadRows"); ServerStreamingCallable traced = new TracedServerStreamingCallable<>( - readRowsUserCallable, clientContext.getTracerFactory(), span); + readRowsUserCallable, + bigtableClientContext.getClientContext().getTracerFactory(), + span); return traced.withDefaultCallContext( - clientContext + bigtableClientContext + .getClientContext() .getDefaultCallContext() .withRetrySettings(readRowsSettings.getRetrySettings())); } @@ -644,10 +653,12 @@ private UnaryCallable> createBulkReadRowsCallable( new TracedBatcherUnaryCallable<>(readRowsUserCallable.all()); UnaryCallable> traced = - new TracedUnaryCallable<>(tracedBatcher, clientContext.getTracerFactory(), span); + new TracedUnaryCallable<>( + tracedBatcher, bigtableClientContext.getClientContext().getTracerFactory(), span); return traced.withDefaultCallContext( - clientContext + bigtableClientContext + .getClientContext() .getDefaultCallContext() .withRetrySettings(settings.readRowsSettings().getRetrySettings())); } @@ -717,7 +728,8 @@ public ApiFuture> futureCall(String s, ApiCallContext apiCallCon methodName, new SampleRowKeysCallableWithRequest(retryable, requestContext) .withDefaultCallContext( - clientContext + bigtableClientContext + .getClientContext() .getDefaultCallContext() .withRetrySettings(settings.sampleRowKeysSettings().getRetrySettings()))); } @@ -760,6 +772,7 @@ private UnaryCallable createMutateRowCallable() { *

This function should not be exposed to external users, as it could cause a data loss. */ private UnaryCallable createMutateRowsBaseCallable() { + ClientContext clientContext = bigtableClientContext.getClientContext(); ServerStreamingCallable base = GrpcRawCallableFactory.createServerStreamingCallable( GrpcCallSettings.newBuilder() @@ -874,9 +887,10 @@ public Batcher newMutateRowsBatcher( bulkMutateRowsCallable, BulkMutation.create(tableId), settings.bulkMutateRowsSettings().getBatchingSettings(), - clientContext.getExecutor(), + bigtableClientContext.getClientContext().getExecutor(), bulkMutationFlowController, - MoreObjects.firstNonNull(ctx, clientContext.getDefaultCallContext())); + MoreObjects.firstNonNull( + ctx, bigtableClientContext.getClientContext().getDefaultCallContext())); } /** @@ -905,9 +919,10 @@ public Batcher newMutateRowsBatcher( bulkMutateRowsCallable, BulkMutation.create(targetId), settings.bulkMutateRowsSettings().getBatchingSettings(), - clientContext.getExecutor(), + bigtableClientContext.getClientContext().getExecutor(), bulkMutationFlowController, - MoreObjects.firstNonNull(ctx, clientContext.getDefaultCallContext())); + MoreObjects.firstNonNull( + ctx, bigtableClientContext.getClientContext().getDefaultCallContext())); } /** @@ -933,9 +948,10 @@ public Batcher newBulkReadRowsBatcher( bulkReadRowsCallable, query, settings.bulkReadRowsSettings().getBatchingSettings(), - clientContext.getExecutor(), + bigtableClientContext.getClientContext().getExecutor(), null, - MoreObjects.firstNonNull(ctx, clientContext.getDefaultCallContext())); + MoreObjects.firstNonNull( + ctx, bigtableClientContext.getClientContext().getDefaultCallContext())); } /** @@ -995,6 +1011,7 @@ private UnaryCallable createReadModifyWriteRowCallable( */ private ServerStreamingCallable createGenerateInitialChangeStreamPartitionsCallable() { + ClientContext clientContext = bigtableClientContext.getClientContext(); ServerStreamingCallable< GenerateInitialChangeStreamPartitionsRequest, GenerateInitialChangeStreamPartitionsResponse> @@ -1075,6 +1092,7 @@ private UnaryCallable createReadModifyWriteRowCallable( ServerStreamingCallable createReadChangeStreamCallable( ChangeStreamRecordAdapter changeStreamRecordAdapter) { + ClientContext clientContext = bigtableClientContext.getClientContext(); ServerStreamingCallable base = GrpcRawCallableFactory.createServerStreamingCallable( GrpcCallSettings.newBuilder() @@ -1151,6 +1169,7 @@ private UnaryCallable createReadModifyWriteRowCallable( */ @InternalApi("For internal use only") public ExecuteQueryCallable createExecuteQueryCallable() { + ClientContext clientContext = bigtableClientContext.getClientContext(); ServerStreamingCallable base = GrpcRawCallableFactory.createServerStreamingCallable( GrpcCallSettings.newBuilder() @@ -1245,9 +1264,13 @@ private UnaryCallable createUserFacin String methodName, UnaryCallable inner) { UnaryCallable traced = - new TracedUnaryCallable<>(inner, clientContext.getTracerFactory(), getSpanName(methodName)); + new TracedUnaryCallable<>( + inner, + bigtableClientContext.getClientContext().getTracerFactory(), + getSpanName(methodName)); - return traced.withDefaultCallContext(clientContext.getDefaultCallContext()); + return traced.withDefaultCallContext( + bigtableClientContext.getClientContext().getDefaultCallContext()); } private Map composeRequestParams( @@ -1314,11 +1337,14 @@ public ApiFuture futureCall(ReqT reqT, ApiCallContext apiCallContext) { UnaryCallable traced = new TracedUnaryCallable<>( transformed, - clientContext.getTracerFactory(), + bigtableClientContext.getClientContext().getTracerFactory(), getSpanName(methodDescriptor.getBareMethodName())); return traced.withDefaultCallContext( - clientContext.getDefaultCallContext().withRetrySettings(callSettings.getRetrySettings())); + bigtableClientContext + .getClientContext() + .getDefaultCallContext() + .withRetrySettings(callSettings.getRetrySettings())); } private UnaryCallable createUnaryCallableNew( @@ -1347,8 +1373,11 @@ private UnaryCallable createUnar return new BigtableUnaryOperationCallable<>( transformed, - clientContext.getDefaultCallContext().withRetrySettings(callSettings.getRetrySettings()), - clientContext.getTracerFactory(), + bigtableClientContext + .getClientContext() + .getDefaultCallContext() + .withRetrySettings(callSettings.getRetrySettings()), + bigtableClientContext.getClientContext().getTracerFactory(), getSpanName(methodDescriptor.getBareMethodName()), /* allowNoResponse= */ false); } @@ -1371,9 +1400,11 @@ private UnaryCallable withRetries( if (settings.getEnableRetryInfo()) { retrying = com.google.cloud.bigtable.gaxx.retrying.Callables.retrying( - innerCallable, unaryCallSettings, clientContext); + innerCallable, unaryCallSettings, bigtableClientContext.getClientContext()); } else { - retrying = Callables.retrying(innerCallable, unaryCallSettings, clientContext); + retrying = + Callables.retrying( + innerCallable, unaryCallSettings, bigtableClientContext.getClientContext()); } if (settings.getEnableRoutingCookie()) { return new CookiesUnaryCallable<>(retrying); @@ -1389,9 +1420,11 @@ private ServerStreamingCallable withR if (settings.getEnableRetryInfo()) { retrying = com.google.cloud.bigtable.gaxx.retrying.Callables.retrying( - innerCallable, serverStreamingCallSettings, clientContext); + innerCallable, serverStreamingCallSettings, bigtableClientContext.getClientContext()); } else { - retrying = Callables.retrying(innerCallable, serverStreamingCallSettings, clientContext); + retrying = + Callables.retrying( + innerCallable, serverStreamingCallSettings, bigtableClientContext.getClientContext()); } if (settings.getEnableRoutingCookie()) { return new CookiesServerStreamingCallable<>(retrying); @@ -1408,7 +1441,7 @@ private ServerStreamingCallable large ServerStreamingCallable retrying; retrying = com.google.cloud.bigtable.gaxx.retrying.Callables.retryingForLargeRows( - innerCallable, serverStreamingCallSettings, clientContext); + innerCallable, serverStreamingCallSettings, bigtableClientContext.getClientContext()); if (settings.getEnableRoutingCookie()) { return new CookiesServerStreamingCallable<>(retrying); } @@ -1507,12 +1540,10 @@ private SpanName getSpanName(String methodName) { @Override public void close() { if (closeClientContext) { - for (BackgroundResource backgroundResource : clientContext.getBackgroundResources()) { - try { - backgroundResource.close(); - } catch (Exception e) { - throw new IllegalStateException("Failed to close resource", e); - } + try { + bigtableClientContext.close(); + } catch (Exception e) { + throw new IllegalStateException("failed to close client context", e); } } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerCallableTest.java index b0966a2166..1a09951b77 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerCallableTest.java @@ -139,7 +139,9 @@ public void sendHeaders(Metadata headers) { null)) .build(); attempts = settings.getStubSettings().readRowsSettings().getRetrySettings().getMaxAttempts(); - stub = new EnhancedBigtableStub(settings.getStubSettings(), clientContext); + stub = + new EnhancedBigtableStub( + settings.getStubSettings(), bigtableClientContext.setClientContext(clientContext)); // Create another server without injecting the server-timing header and another stub that // connects to it. @@ -164,7 +166,9 @@ public void sendHeaders(Metadata headers) { null)) .build(); noHeaderStub = - new EnhancedBigtableStub(noHeaderSettings.getStubSettings(), noHeaderClientContext); + new EnhancedBigtableStub( + noHeaderSettings.getStubSettings(), + noHeaderBigtableClientContext.setClientContext(noHeaderClientContext)); } @After 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 a59cf84751..864a801e6e 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 @@ -50,7 +50,6 @@ import com.google.api.gax.batching.BatchingSettings; import com.google.api.gax.batching.FlowControlSettings; import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider; -import com.google.api.gax.rpc.ClientContext; import com.google.api.gax.rpc.NotFoundException; import com.google.api.gax.rpc.ResponseObserver; import com.google.api.gax.rpc.StreamController; @@ -71,6 +70,7 @@ import com.google.cloud.bigtable.data.v2.models.RowMutation; import com.google.cloud.bigtable.data.v2.models.RowMutationEntry; import com.google.cloud.bigtable.data.v2.models.TableId; +import com.google.cloud.bigtable.data.v2.stub.BigtableClientContext; import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub; import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings; import com.google.common.base.Stopwatch; @@ -276,7 +276,7 @@ public void sendHeaders(Metadata headers) { }); stubSettingsBuilder.setTransportChannelProvider(channelProvider.build()); EnhancedBigtableStubSettings stubSettings = stubSettingsBuilder.build(); - stub = new EnhancedBigtableStub(stubSettings, ClientContext.create(stubSettings)); + stub = new EnhancedBigtableStub(stubSettings, BigtableClientContext.create(stubSettings)); } @After diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java index a9f3aa038b..8446f80f87 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java @@ -133,7 +133,9 @@ public void setUp() throws Exception { localStats.getStatsRecorder(), null)) .build(); - stub = new EnhancedBigtableStub(settings.getStubSettings(), clientContext); + stub = + new EnhancedBigtableStub( + settings.getStubSettings(), bigtableClientContext.setClientContext(clientContext)); } @After