Skip to content

Commit

Permalink
chore: refactor settings (#2040)
Browse files Browse the repository at this point in the history 
* chore: refactor settings

* fix format

* create a shared context

* remove patchSettings

* reformat

* update
  • Loading branch information
@mutianf
mutianf authored Jan 9, 2024
1 parent dad7517 commit 6b48606
Show file tree
Hide file tree
Showing 9 changed files with 203 additions and 129 deletions.
6 changes: 6 additions & 0 deletions google-cloud-bigtable/clirr-ignored-differences.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,12 @@
<className>com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub</className>
<method>*</method>
</difference>
<difference>
<!-- method name change is ok because EnhancedBigtableStub is InternalApi -->
<differenceType>7002</differenceType>
<className>com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub</className>
<method>*</method>
</difference>
<!-- InternalApi that was renamed -->
<difference>
<differenceType>8001</differenceType>
Expand Down
13 changes: 13 additions & 0 deletions ...le-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
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;
Expand Down Expand Up @@ -166,6 +167,18 @@ public static BigtableDataClient create(BigtableDataSettings settings) throws IO
return new BigtableDataClient(stub);
}

/**
* Constructs an instance of BigtableDataClient with the provided client context. This is used by
* {@link BigtableDataClientFactory} and the client context will not be closed unless {@link
* BigtableDataClientFactory#close()} is called.
*/
static BigtableDataClient createWithClientContext(
BigtableDataSettings settings, ClientContext context) throws IOException {
EnhancedBigtableStub stub =
EnhancedBigtableStub.createWithClientContext(settings.getStubSettings(), context);
return new BigtableDataClient(stub);
}

@InternalApi("Visible for testing")
BigtableDataClient(EnhancedBigtableStub stub) {
this.stub = stub;
Expand Down
85 changes: 41 additions & 44 deletions ...d-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,8 @@

import com.google.api.core.BetaApi;
import com.google.api.gax.core.BackgroundResource;
import com.google.api.gax.core.FixedCredentialsProvider;
import com.google.api.gax.core.FixedExecutorProvider;
import com.google.api.gax.rpc.ClientContext;
import com.google.api.gax.rpc.FixedHeaderProvider;
import com.google.api.gax.rpc.FixedTransportChannelProvider;
import com.google.api.gax.rpc.FixedWatchdogProvider;
import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings;
import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub;
import java.io.IOException;
import javax.annotation.Nonnull;

Expand Down Expand Up @@ -78,7 +73,8 @@ public final class BigtableDataClientFactory implements AutoCloseable {
*/
public static BigtableDataClientFactory create(BigtableDataSettings defaultSettings)
throws IOException {
ClientContext sharedClientContext = ClientContext.create(defaultSettings.getStubSettings());
ClientContext sharedClientContext =
EnhancedBigtableStub.createClientContext(defaultSettings.getStubSettings());
return new BigtableDataClientFactory(sharedClientContext, defaultSettings);
}

Expand Down Expand Up @@ -110,12 +106,16 @@ public void close() throws Exception {
* release all resources, first close all of the created clients and then this factory instance.
*/
public BigtableDataClient createDefault() {
BigtableDataSettings.Builder settingsBuilder = defaultSettings.toBuilder();
patchStubSettings(settingsBuilder.stubSettings());
BigtableDataSettings settings = settingsBuilder.build();

try {
return BigtableDataClient.create(settings);
ClientContext clientContext =
sharedClientContext
.toBuilder()
.setTracerFactory(
EnhancedBigtableStub.createBigtableTracerFactory(
defaultSettings.getStubSettings()))
.build();

return BigtableDataClient.createWithClientContext(defaultSettings, clientContext);
} catch (IOException e) {
// Should never happen because the connection has been established already
throw new RuntimeException(
Expand All @@ -133,12 +133,16 @@ public BigtableDataClient createDefault() {
* release all resources, first close all of the created clients and then this factory instance.
*/
public BigtableDataClient createForAppProfile(@Nonnull String appProfileId) throws IOException {
BigtableDataSettings.Builder settingsBuilder =
defaultSettings.toBuilder().setAppProfileId(appProfileId);

patchStubSettings(settingsBuilder.stubSettings());
BigtableDataSettings settings =
defaultSettings.toBuilder().setAppProfileId(appProfileId).build();

return BigtableDataClient.create(settingsBuilder.build());
ClientContext clientContext =
sharedClientContext
.toBuilder()
.setTracerFactory(
EnhancedBigtableStub.createBigtableTracerFactory(settings.getStubSettings()))
.build();
return BigtableDataClient.createWithClientContext(settings, clientContext);
}

/**
Expand All @@ -152,16 +156,22 @@ public BigtableDataClient createForAppProfile(@Nonnull String appProfileId) thro
*/
public BigtableDataClient createForInstance(@Nonnull String projectId, @Nonnull String instanceId)
throws IOException {
BigtableDataSettings.Builder settingsBuilder =
BigtableDataSettings settings =
defaultSettings
.toBuilder()
.setProjectId(projectId)
.setInstanceId(instanceId)
.setDefaultAppProfileId();
.setDefaultAppProfileId()
.build();

patchStubSettings(settingsBuilder.stubSettings());
ClientContext clientContext =
sharedClientContext
.toBuilder()
.setTracerFactory(
EnhancedBigtableStub.createBigtableTracerFactory(settings.getStubSettings()))
.build();

return BigtableDataClient.create(settingsBuilder.build());
return BigtableDataClient.createWithClientContext(settings, clientContext);
}

/**
Expand All @@ -176,32 +186,19 @@ public BigtableDataClient createForInstance(@Nonnull String projectId, @Nonnull
public BigtableDataClient createForInstance(
@Nonnull String projectId, @Nonnull String instanceId, @Nonnull String appProfileId)
throws IOException {
BigtableDataSettings.Builder settingsBuilder =
BigtableDataSettings settings =
defaultSettings
.toBuilder()
.setProjectId(projectId)
.setInstanceId(instanceId)
.setAppProfileId(appProfileId);

patchStubSettings(settingsBuilder.stubSettings());

return BigtableDataClient.create(settingsBuilder.build());
}

// Update stub settings to use shared resources in this factory
private void patchStubSettings(EnhancedBigtableStubSettings.Builder stubSettings) {
stubSettings
// Channel refreshing will be configured in the shared ClientContext. Derivative clients
// won't be able to reconfigure the refreshing logic
.setRefreshingChannel(false)
.setTransportChannelProvider(
FixedTransportChannelProvider.create(sharedClientContext.getTransportChannel()))
.setCredentialsProvider(
FixedCredentialsProvider.create(sharedClientContext.getCredentials()))
.setExecutorProvider(FixedExecutorProvider.create(sharedClientContext.getExecutor()))
.setStreamWatchdogProvider(
FixedWatchdogProvider.create(sharedClientContext.getStreamWatchdog()))
.setHeaderProvider(FixedHeaderProvider.create(sharedClientContext.getHeaders()))
.setClock(sharedClientContext.getClock());
.setAppProfileId(appProfileId)
.build();
ClientContext clientContext =
sharedClientContext
.toBuilder()
.setTracerFactory(
EnhancedBigtableStub.createBigtableTracerFactory(settings.getStubSettings()))
.build();
return BigtableDataClient.createWithClientContext(settings, clientContext);
}
}
117 changes: 71 additions & 46 deletions ...d-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@
import com.google.api.gax.rpc.ServerStreamingCallable;
import com.google.api.gax.rpc.UnaryCallSettings;
import com.google.api.gax.rpc.UnaryCallable;
import com.google.api.gax.tracing.ApiTracerFactory;
import com.google.api.gax.tracing.OpencensusTracerFactory;
import com.google.api.gax.tracing.SpanName;
import com.google.api.gax.tracing.TracedServerStreamingCallable;
Expand Down Expand Up @@ -110,6 +111,7 @@
import com.google.cloud.bigtable.data.v2.stub.readrows.RowMergingCallable;
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.MoreObjects;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
Expand Down Expand Up @@ -150,6 +152,8 @@ public class EnhancedBigtableStub implements AutoCloseable {

private final EnhancedBigtableStubSettings settings;
private final ClientContext clientContext;

private final boolean closeClientContext;
private final RequestContext requestContext;
private final FlowController bulkMutationFlowController;
private final DynamicFlowControlStats bulkMutationDynamicFlowControlStats;
Expand All @@ -172,13 +176,20 @@ public class EnhancedBigtableStub implements AutoCloseable {

public static EnhancedBigtableStub create(EnhancedBigtableStubSettings settings)
throws IOException {
settings = finalizeSettings(settings, Tags.getTagger(), Stats.getStatsRecorder());

return new EnhancedBigtableStub(settings, ClientContext.create(settings));
settings = settings.toBuilder().setTracerFactory(createBigtableTracerFactory(settings)).build();
ClientContext clientContext = createClientContext(settings);

return new EnhancedBigtableStub(settings, clientContext);
}

public static EnhancedBigtableStubSettings finalizeSettings(
EnhancedBigtableStubSettings settings, Tagger tagger, StatsRecorder stats)
public static EnhancedBigtableStub createWithClientContext(
EnhancedBigtableStubSettings settings, ClientContext clientContext) throws IOException {

return new EnhancedBigtableStub(settings, clientContext, false);
}

public static ClientContext createClientContext(EnhancedBigtableStubSettings settings)
throws IOException {
EnhancedBigtableStubSettings.Builder builder = settings.toBuilder();

Expand Down Expand Up @@ -222,49 +233,53 @@ public static EnhancedBigtableStubSettings finalizeSettings(
builder.setTransportChannelProvider(transportProvider.build());
}

return ClientContext.create(builder.build());
}

public static ApiTracerFactory createBigtableTracerFactory(
EnhancedBigtableStubSettings settings) {
return createBigtableTracerFactory(settings, Tags.getTagger(), Stats.getStatsRecorder());
}

@VisibleForTesting
public static ApiTracerFactory createBigtableTracerFactory(
EnhancedBigtableStubSettings settings, Tagger tagger, StatsRecorder stats) {
String projectId = settings.getProjectId();
String instanceId = settings.getInstanceId();
String appProfileId = settings.getAppProfileId();

ImmutableMap<TagKey, TagValue> attributes =
ImmutableMap.<TagKey, TagValue>builder()
.put(RpcMeasureConstants.BIGTABLE_PROJECT_ID, TagValue.create(settings.getProjectId()))
.put(
RpcMeasureConstants.BIGTABLE_INSTANCE_ID, TagValue.create(settings.getInstanceId()))
.put(
RpcMeasureConstants.BIGTABLE_APP_PROFILE_ID,
TagValue.create(settings.getAppProfileId()))
.put(RpcMeasureConstants.BIGTABLE_PROJECT_ID, TagValue.create(projectId))
.put(RpcMeasureConstants.BIGTABLE_INSTANCE_ID, TagValue.create(instanceId))
.put(RpcMeasureConstants.BIGTABLE_APP_PROFILE_ID, TagValue.create(appProfileId))
.build();
ImmutableMap<String, String> builtinAttributes =
ImmutableMap.<String, String>builder()
.put("project_id", settings.getProjectId())
.put("instance", settings.getInstanceId())
.put("app_profile", settings.getAppProfileId())
.put("project_id", projectId)
.put("instance", instanceId)
.put("app_profile", appProfileId)
.build();
// Inject Opencensus instrumentation
builder.setTracerFactory(
new CompositeTracerFactory(
ImmutableList.of(
// Add OpenCensus Tracing
new OpencensusTracerFactory(
ImmutableMap.<String, String>builder()
// Annotate traces with the same tags as metrics
.put(
RpcMeasureConstants.BIGTABLE_PROJECT_ID.getName(),
settings.getProjectId())
.put(
RpcMeasureConstants.BIGTABLE_INSTANCE_ID.getName(),
settings.getInstanceId())
.put(
RpcMeasureConstants.BIGTABLE_APP_PROFILE_ID.getName(),
settings.getAppProfileId())
// Also annotate traces with library versions
.put("gax", GaxGrpcProperties.getGaxGrpcVersion())
.put("grpc", GaxGrpcProperties.getGrpcVersion())
.put("gapic", Version.VERSION)
.build()),
// Add OpenCensus Metrics
MetricsTracerFactory.create(tagger, stats, attributes),
BuiltinMetricsTracerFactory.create(builtinAttributes),
// Add user configured tracer
settings.getTracerFactory())));
return builder.build();

return new CompositeTracerFactory(
ImmutableList.of(
// Add OpenCensus Tracing
new OpencensusTracerFactory(
ImmutableMap.<String, String>builder()
// Annotate traces with the same tags as metrics
.put(RpcMeasureConstants.BIGTABLE_PROJECT_ID.getName(), projectId)
.put(RpcMeasureConstants.BIGTABLE_INSTANCE_ID.getName(), instanceId)
.put(RpcMeasureConstants.BIGTABLE_APP_PROFILE_ID.getName(), appProfileId)
// Also annotate traces with library versions
.put("gax", GaxGrpcProperties.getGaxGrpcVersion())
.put("grpc", GaxGrpcProperties.getGrpcVersion())
.put("gapic", Version.VERSION)
.build()),
// Add OpenCensus Metrics
MetricsTracerFactory.create(tagger, stats, attributes),
BuiltinMetricsTracerFactory.create(builtinAttributes),
// Add user configured tracer
settings.getTracerFactory()));
}

private static void patchCredentials(EnhancedBigtableStubSettings.Builder settings)
Expand Down Expand Up @@ -303,8 +318,16 @@ private static void patchCredentials(EnhancedBigtableStubSettings.Builder settin
}

public EnhancedBigtableStub(EnhancedBigtableStubSettings settings, ClientContext clientContext) {
this(settings, clientContext, true);
}

public EnhancedBigtableStub(
EnhancedBigtableStubSettings settings,
ClientContext clientContext,
boolean closeClientContext) {
this.settings = settings;
this.clientContext = clientContext;
this.closeClientContext = closeClientContext;
this.requestContext =
RequestContext.create(
settings.getProjectId(), settings.getInstanceId(), settings.getAppProfileId());
Expand Down Expand Up @@ -1166,11 +1189,13 @@ private SpanName getSpanName(String methodName) {

@Override
public void close() {
for (BackgroundResource backgroundResource : clientContext.getBackgroundResources()) {
try {
backgroundResource.close();
} catch (Exception e) {
throw new IllegalStateException("Failed to close resource", e);
if (closeClientContext) {
for (BackgroundResource backgroundResource : clientContext.getBackgroundResources()) {
try {
backgroundResource.close();
} catch (Exception e) {
throw new IllegalStateException("Failed to close resource", e);
}
}
}
}
Expand Down
22 changes: 0 additions & 22 deletions ...le/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
import com.google.api.gax.batching.FlowControlSettings;
import com.google.api.gax.batching.FlowController;
import com.google.api.gax.batching.FlowController.LimitExceededBehavior;
import com.google.api.gax.core.FixedCredentialsProvider;
import com.google.api.gax.core.GoogleCredentialsProvider;
import com.google.api.gax.grpc.ChannelPoolSettings;
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
Expand All @@ -33,7 +32,6 @@
import com.google.api.gax.rpc.StubSettings;
import com.google.api.gax.rpc.TransportChannelProvider;
import com.google.api.gax.rpc.UnaryCallSettings;
import com.google.auth.Credentials;
import com.google.bigtable.v2.FeatureFlags;
import com.google.bigtable.v2.PingAndWarmRequest;
import com.google.cloud.bigtable.Version;
Expand Down Expand Up @@ -1023,26 +1021,6 @@ public UnaryCallSettings.Builder<PingAndWarmRequest, Void> pingAndWarmSettings()
public EnhancedBigtableStubSettings build() {
Preconditions.checkState(projectId != null, "Project id must be set");
Preconditions.checkState(instanceId != null, "Instance id must be set");
if (isRefreshingChannel) {
Preconditions.checkArgument(
getTransportChannelProvider() instanceof InstantiatingGrpcChannelProvider,
"refreshingChannel only works with InstantiatingGrpcChannelProviders");
InstantiatingGrpcChannelProvider.Builder channelProviderBuilder =
((InstantiatingGrpcChannelProvider) getTransportChannelProvider()).toBuilder();
Credentials credentials = null;
if (getCredentialsProvider() != null) {
try {
credentials = getCredentialsProvider().getCredentials();
} catch (IOException e) {
throw new RuntimeException(e);
}
}
// Use shared credentials
this.setCredentialsProvider(FixedCredentialsProvider.create(credentials));
channelProviderBuilder.setChannelPrimer(
BigtableChannelPrimer.create(credentials, projectId, instanceId, appProfileId));
this.setTransportChannelProvider(channelProviderBuilder.build());
}

if (this.bulkMutateRowsSettings().isServerInitiatedFlowControlEnabled()) {
// only set mutate rows feature flag when this feature is enabled
Expand Down
Loading

0 comments on commit 6b48606

Please sign in to comment.