From be71100bc93d5394ae22c1a59f893bf26364da8d Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Fri, 31 Jul 2026 14:21:57 +0000 Subject: [PATCH] feat(bigtable): attach client config UUID header to GetClientConfigurationRequest Generate a UUID when the session client is initialized and attach it via the bigtable-client-config-uuid header on every GetClientConfigurationRequest (the initial fetch and all refresh polls) so the server can correlate requests from the same client instance. --- .../util/ClientConfigurationManager.java | 31 ++++++++++++++ .../util/ClientConfigurationManagerTest.java | 40 ++++++++++++++++++- 2 files changed, 70 insertions(+), 1 deletion(-) diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManager.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManager.java index 03a83f034b14..89f2b689d016 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManager.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManager.java @@ -26,6 +26,7 @@ import com.google.cloud.bigtable.data.v2.internal.api.Util; import com.google.cloud.bigtable.data.v2.internal.csm.attributes.ClientInfo; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DebugTagTracer; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableMap; import com.google.protobuf.TextFormat; @@ -52,6 +53,7 @@ import java.util.Objects; import java.util.Optional; import java.util.Properties; +import java.util.UUID; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; @@ -74,6 +76,17 @@ public class ClientConfigurationManager implements AutoCloseable { public static final String OVERRIDE_SYS_PROP_KEY = "bigtable.internal.client-config-override"; + /** + * Header carrying a UUID that uniquely identifies this session client instance. It is attached to + * every {@link GetClientConfigurationRequest} (both the initial fetch and all subsequent refresh + * polls) so the server can correlate requests originating from the same client. + */ + private static final String CLIENT_UUID_HEADER = "bigtable-client-config-uuid"; + + @VisibleForTesting + static final Metadata.Key CLIENT_UUID_KEY = + Metadata.Key.of(CLIENT_UUID_HEADER, Metadata.ASCII_STRING_MARSHALLER); + public interface ConfigListener { void onChange(T newValue); } @@ -118,6 +131,10 @@ public void close() { private final GetClientConfigurationRequest request; private final ChannelProvider channelProvider; + // A UUID generated once when this session client is initialized. It is attached to the + // GetClientConfigurationRequest header for both the initial and all refresh requests. + private final String clientUuid; + @GuardedBy("this") private ManagedChannel channel; @@ -196,6 +213,13 @@ public ClientConfigurationManager( ImmutableMap.of( "instance_name", clientInfo.getInstanceName().toString(), "app_profile_id", clientInfo.getAppProfileId())); + + // Generate a UUID that uniquely identifies this session client instance and attach it to the + // request metadata. Since the same metadata is reused for the initial fetch and every refresh + // poll, the header is sent on all GetClientConfigurationRequests. + this.clientUuid = UUID.randomUUID().toString(); + this.metadata.put(CLIENT_UUID_KEY, clientUuid); + this.request = GetClientConfigurationRequest.newBuilder() .setInstanceName(clientInfo.getInstanceName().toString()) @@ -262,6 +286,13 @@ ClientConfiguration getDefaultConfig() { return defaultConfig; } + /** + * Returns the UUID generated for this session client, attached to every config request header. + */ + String getClientUuid() { + return clientUuid; + } + public synchronized ListenerHandle addListener( Function extractor, ConfigListener listener) { ListenerEntry entry = new ListenerEntry<>(extractor, listener); diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManagerTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManagerTest.java index e62e3ced906b..08bd82a8feaa 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManagerTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManagerTest.java @@ -53,6 +53,9 @@ import io.grpc.Metadata; import io.grpc.MethodDescriptor; import io.grpc.Server; +import io.grpc.ServerCall; +import io.grpc.ServerCallHandler; +import io.grpc.ServerInterceptor; import io.grpc.Status; import io.grpc.stub.StreamObserver; import java.io.IOException; @@ -97,6 +100,8 @@ class ClientConfigurationManagerTest { private ChannelProviders.ChannelProvider channelProvider; @Mock private ScheduledExecutorService mockExecutor; private final OutstandingRpcCounter outstandingRpcCounter = new OutstandingRpcCounter(); + // Captures the bigtable-client-config-uuid header seen by the server on the most recent request. + private final AtomicReference lastClientUuid = new AtomicReference<>(); private ClientConfigurationManager manager; private final NoopMetrics.NoopDebugTracer noopDebugTracer = NoopMetrics.NoopDebugTracer.INSTANCE; @@ -104,7 +109,17 @@ class ClientConfigurationManagerTest { void setUp() throws IOException { service = new FakeConfigService(); - server = FakeServiceBuilder.create(service).start(); + ServerInterceptor uuidCapturingInterceptor = + new ServerInterceptor() { + @Override + public ServerCall.Listener interceptCall( + ServerCall call, Metadata headers, ServerCallHandler next) { + lastClientUuid.set(headers.get(ClientConfigurationManager.CLIENT_UUID_KEY)); + return next.startCall(call, headers); + } + }; + + server = FakeServiceBuilder.create(service).intercept(uuidCapturingInterceptor).start(); channelProvider = new ForwardingChannelProvider( @@ -145,6 +160,29 @@ void tearDown() { server.shutdown(); } + @Test + void clientUuidHeaderSentOnInitialAndRefreshRequests() throws Exception { + // The header name must be prefixed with "bigtable-". + assertThat(ClientConfigurationManager.CLIENT_UUID_KEY.name()).startsWith("bigtable-"); + + // Fetch the initial config and capture the header sent with it. + manager.start().get(); + outstandingRpcCounter.waitUntilRpcsDone(); + + String initialUuid = lastClientUuid.get(); + // The header must be present and match the manager's generated UUID. + assertThat(initialUuid).isNotNull(); + assertThat(initialUuid).isEqualTo(manager.getClientUuid()); + + // Trigger a refresh poll and confirm the same UUID header rides along. + ArgumentCaptor runnableCaptor = ArgumentCaptor.forClass(Runnable.class); + verify(mockExecutor, times(1)).schedule(runnableCaptor.capture(), anyLong(), any()); + runnableCaptor.getValue().run(); + outstandingRpcCounter.waitUntilRpcsDone(); + + assertThat(lastClientUuid.get()).isEqualTo(manager.getClientUuid()); + } + @Test void initialFetchTest() throws ExecutionException, InterruptedException { // Check the initial config is correct