Skip to content

Commit 3d0a6d9

Browse files
mutianfmutianf
authored
fix: create stub with BigtableClientContext so otels are closed (#2747)
* fix: create stub with BigtableClientContext so otels are closed * address comment
1 parent b3cadab commit 3d0a6d9

8 files changed

Lines changed: 104 additions & 51 deletions

File tree

‎google-cloud-bigtable/clirr-ignored-differences.xml‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,13 @@
3434
<className>com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub</className>
3535
<method>*</method>
3636
</difference>
37+
<difference>
38+
<!-- change method types is ok because EnhancedBigtableStub is InternalApi -->
39+
<differenceType>7005</differenceType>
40+
<className>com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub</className>
41+
<method>*</method>
42+
<to>*</to>
43+
</difference>
3744
<difference>
3845
<!-- method name change is ok because EnhancedBigtableStub is InternalApi -->
3946
<differenceType>7002</differenceType>

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
import com.google.api.gax.batching.Batcher;
2626
import com.google.api.gax.grpc.GrpcCallContext;
2727
import com.google.api.gax.rpc.ApiExceptions;
28-
import com.google.api.gax.rpc.ClientContext;
2928
import com.google.api.gax.rpc.ResponseObserver;
3029
import com.google.api.gax.rpc.ServerStream;
3130
import com.google.api.gax.rpc.ServerStreamingCallable;
@@ -55,6 +54,7 @@
5554
import com.google.cloud.bigtable.data.v2.models.sql.PreparedStatement;
5655
import com.google.cloud.bigtable.data.v2.models.sql.ResultSet;
5756
import com.google.cloud.bigtable.data.v2.models.sql.SqlType;
57+
import com.google.cloud.bigtable.data.v2.stub.BigtableClientContext;
5858
import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub;
5959
import com.google.cloud.bigtable.data.v2.stub.sql.SqlServerStream;
6060
import com.google.common.util.concurrent.MoreExecutors;
@@ -186,7 +186,7 @@ public static BigtableDataClient create(BigtableDataSettings settings) throws IO
186186
* BigtableDataClientFactory#close()} is called.
187187
*/
188188
static BigtableDataClient createWithClientContext(
189-
BigtableDataSettings settings, ClientContext context) throws IOException {
189+
BigtableDataSettings settings, BigtableClientContext context) throws IOException {
190190
EnhancedBigtableStub stub =
191191
EnhancedBigtableStub.createWithClientContext(settings.getStubSettings(), context);
192192
return new BigtableDataClient(stub);

‎google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,8 @@ public BigtableDataClient createDefault() {
114114
defaultSettings.getStubSettings(), sharedClientContext.getOpenTelemetry()))
115115
.build();
116116

117-
return BigtableDataClient.createWithClientContext(defaultSettings, clientContext);
117+
return BigtableDataClient.createWithClientContext(
118+
defaultSettings, sharedClientContext.withClientContext(clientContext));
118119
} catch (IOException e) {
119120
// Should never happen because the connection has been established already
120121
throw new RuntimeException(
@@ -141,7 +142,8 @@ public BigtableDataClient createForAppProfile(@Nonnull String appProfileId) thro
141142
EnhancedBigtableStub.createBigtableTracerFactory(
142143
settings.getStubSettings(), sharedClientContext.getOpenTelemetry()))
143144
.build();
144-
return BigtableDataClient.createWithClientContext(settings, clientContext);
145+
return BigtableDataClient.createWithClientContext(
146+
settings, sharedClientContext.withClientContext(clientContext));
145147
}
146148

147149
/**
@@ -169,7 +171,8 @@ public BigtableDataClient createForInstance(@Nonnull String projectId, @Nonnull
169171
settings.getStubSettings(), sharedClientContext.getOpenTelemetry()))
170172
.build();
171173

172-
return BigtableDataClient.createWithClientContext(settings, clientContext);
174+
return BigtableDataClient.createWithClientContext(
175+
settings, sharedClientContext.withClientContext(clientContext));
173176
}
174177

175178
/**
@@ -196,6 +199,7 @@ public BigtableDataClient createForInstance(
196199
EnhancedBigtableStub.createBigtableTracerFactory(
197200
settings.getStubSettings(), sharedClientContext.getOpenTelemetry()))
198201
.build();
199-
return BigtableDataClient.createWithClientContext(settings, clientContext);
202+
return BigtableDataClient.createWithClientContext(
203+
settings, sharedClientContext.withClientContext(clientContext));
200204
}
201205
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,11 @@ public ClientContext getClientContext() {
197197
return this.clientContext;
198198
}
199199

200+
public BigtableClientContext withClientContext(ClientContext clientContext) {
201+
return new BigtableClientContext(
202+
clientContext, openTelemetry, internalOpenTelemetry, metricsProvider);
203+
}
204+
200205
public void close() throws Exception {
201206
for (BackgroundResource resource : clientContext.getBackgroundResources()) {
202207
resource.close();

0 commit comments

Comments
 (0)