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.

Author: mutianf

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<String> CLIENT_UUID_KEY =
+      Metadata.Key.of(CLIENT_UUID_HEADER, Metadata.ASCII_STRING_MARSHALLER);
+
   public interface ConfigListener<T> {
     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 <T> ListenerHandle addListener(
       Function<ClientConfiguration, T> extractor, ConfigListener<T> listener) {
     ListenerEntry<T> entry = new ListenerEntry<>(extractor, listener);

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,14 +100,26 @@ 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<String> lastClientUuid = new AtomicReference<>();
   private ClientConfigurationManager manager;
   private final NoopMetrics.NoopDebugTracer noopDebugTracer = NoopMetrics.NoopDebugTracer.INSTANCE;
 
   @BeforeEach
   void setUp() throws IOException {
     service = new FakeConfigService();
 
-    server = FakeServiceBuilder.create(service).start();
+    ServerInterceptor uuidCapturingInterceptor =
+        new ServerInterceptor() {
+          @Override
+          public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
+              ServerCall<ReqT, RespT> call, Metadata headers, ServerCallHandler<ReqT, RespT> 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<Runnable> 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