fix: allow trace id to switch within the same connection (#2852)

GaoleMeng · web-flow · commit 850eaea9eaba · 2025-02-11T11:08:59.000-08:00
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorker.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorker.java
@@ -868,9 +868,7 @@ private void appendLoop() {
           // If we are at the first request for every table switch, including the first request in
           // the connection, we will attach both stream name and table schema to the request.
           destinationSet.add(streamName);
-          if (this.traceId != null) {
-            originalRequestBuilder.setTraceId(this.traceId);
-          }
+          originalRequestBuilder.setTraceId(wrapper.streamWriter.getFullTraceId());
         } else if (!isMultiplexing) {
           // If we are not in multiplexing and not in the first request, clear the stream name.
           originalRequestBuilder.clearWriteStream();
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPool.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPool.java
@@ -95,10 +95,6 @@ public class ConnectionWorkerPool {
   /** Enable test related logic. */
   private static boolean enableTesting = false;
 
-  /*
-   * TraceId for debugging purpose.
-   */
-  private final String traceId;
   /*
    * Sets the compression to use for the calls
    */
@@ -213,7 +209,6 @@ public abstract static class Builder {
       long maxInflightBytes,
       java.time.Duration maxRetryDuration,
       FlowController.LimitExceededBehavior limitExceededBehavior,
-      String traceId,
       @Nullable String comperssorName,
       BigQueryWriteSettings clientSettings,
       RetrySettings retrySettings,
@@ -223,7 +218,6 @@ public abstract static class Builder {
     this.maxInflightBytes = maxInflightBytes;
     this.maxRetryDuration = maxRetryDuration;
     this.limitExceededBehavior = limitExceededBehavior;
-    this.traceId = traceId;
     this.compressorName = comperssorName;
     this.clientSettings = clientSettings;
     this.currentMaxConnectionCount = settings.minConnectionsPerRegion();
@@ -319,7 +313,10 @@ private ConnectionWorker createOrReuseConnectionWorker(
     if (connectionWorkerPool.size() < currentMaxConnectionCount) {
       // Always create a new connection if we haven't reached current maximum.
       return createConnectionWorker(
-          streamWriter.getStreamName(), streamWriter.getLocation(), streamWriter.getProtoSchema());
+          streamWriter.getStreamName(),
+          streamWriter.getLocation(),
+          streamWriter.getProtoSchema(),
+          streamWriter.getFullTraceId());
     } else {
       ConnectionWorker existingBestConnection =
           pickBestLoadConnection(
@@ -338,7 +335,8 @@ private ConnectionWorker createOrReuseConnectionWorker(
         return createConnectionWorker(
             streamWriter.getStreamName(),
             streamWriter.getLocation(),
-            streamWriter.getProtoSchema());
+            streamWriter.getProtoSchema(),
+            streamWriter.getFullTraceId());
       } else {
         // Stick to the original connection if all the connections are overwhelmed.
         if (existingConnectionWorker != null) {
@@ -394,7 +392,8 @@ static ConnectionWorker pickBestLoadConnection(
    * computeIfAbsent(...) which is at most once per key.
    */
   private ConnectionWorker createConnectionWorker(
-      String streamName, String location, ProtoSchema writeSchema) throws IOException {
+      String streamName, String location, ProtoSchema writeSchema, String fullTraceId)
+      throws IOException {
     if (enableTesting) {
       // Though atomic integer is super lightweight, add extra if check in case adding future logic.
       testValueCreateConnectionCount.getAndIncrement();
@@ -408,7 +407,7 @@ private ConnectionWorker createConnectionWorker(
             maxInflightBytes,
             maxRetryDuration,
             limitExceededBehavior,
-            traceId,
+            fullTraceId,
             compressorName,
             clientSettings,
             retrySettings,
@@ -494,10 +493,6 @@ int getTotalConnectionCount() {
     return connectionWorkerPool.size();
   }
 
-  String getTraceId() {
-    return traceId;
-  }
-
   FlowController.LimitExceededBehavior limitExceededBehavior() {
     return limitExceededBehavior;
   }
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java
@@ -73,6 +73,9 @@ public class StreamWriter implements AutoCloseable {
    */
   private final String streamName;
 
+  /** This is the library version may or may not include library version id. */
+  private final String fullTraceId;
+
   /** Every writer has a fixed proto schema. */
   private final ProtoSchema writerSchema;
 
@@ -233,6 +236,7 @@ private StreamWriter(Builder builder) throws IOException {
     BigQueryWriteSettings clientSettings = getBigQueryWriteSettings(builder);
     this.requestProfilerHook =
         new RequestProfiler.RequestProfilerHook(builder.enableRequestProfiler);
+    this.fullTraceId = builder.getFullTraceId();
     if (builder.enableRequestProfiler) {
       // Request profiler is enabled on singleton level, from now on a periodical flush will be
       // started
@@ -320,7 +324,6 @@ private StreamWriter(Builder builder) throws IOException {
                         builder.maxInflightBytes,
                         builder.maxRetryDuration,
                         builder.limitExceededBehavior,
-                        builder.getFullTraceId(),
                         builder.compressorName,
                         client.getSettings(),
                         builder.retrySettings,
@@ -359,6 +362,10 @@ static boolean isDefaultStream(String streamName) {
     return streamMatcher.find();
   }
 
+  String getFullTraceId() {
+    return fullTraceId;
+  }
+
   AppendRowsRequest.MissingValueInterpretation getDefaultValueInterpretation() {
     return defaultMissingValueInterpretation;
   }
@@ -401,15 +408,6 @@ static BigQueryWriteSettings getBigQueryWriteSettings(Builder builder) throws IO
 
   // Validate whether the fetched connection pool matched certain properties.
   private void validateFetchedConnectonPool(StreamWriter.Builder builder) {
-    String storedTraceId =
-        this.singleConnectionOrConnectionPool.connectionWorkerPool().getTraceId();
-    if (!Objects.equals(storedTraceId, builder.getFullTraceId())) {
-      throw new IllegalArgumentException(
-          String.format(
-              "Trace id used for the same connection pool for the same location must be the same, "
-                  + "however stored trace id is %s, and expected trace id is %s.",
-              storedTraceId, builder.getFullTraceId()));
-    }
     FlowController.LimitExceededBehavior storedLimitExceededBehavior =
         singleConnectionOrConnectionPool.connectionWorkerPool().limitExceededBehavior();
     if (!Objects.equals(storedLimitExceededBehavior, builder.limitExceededBehavior)) {
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPoolTest.java
@@ -561,7 +561,6 @@ ConnectionWorkerPool createConnectionWorkerPool(
         maxBytes,
         maxRetryDuration,
         FlowController.LimitExceededBehavior.Block,
-        TEST_TRACE_ID,
         null,
         clientSettings,
         retrySettings,
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerTest.java
@@ -86,23 +86,50 @@ public void setUp() throws Exception {
   }
 
   @Test
-  public void testMultiplexedAppendSuccess() throws Exception {
+  public void testMultiplexedAppendSuccess_NonNullTraceId() throws Exception {
+    testMultiplexedIngestion(
+        /*sw1TraceId=*/ "header_1:trailer_1",
+        /*sw2TraceId=*/ "header_2:trailer_2",
+        /*expectedSW1TraceId=*/ "java-streamwriter header_1:trailer_1",
+        /*expectedSW2TraceId=*/ "java-streamwriter header_2:trailer_2");
+  }
+
+  @Test
+  public void testMultiplexedAppendSuccess_EmptyTraceId() throws Exception {
+    testMultiplexedIngestion(
+        /*sw1TraceId=*/ "header_1:trailer_1",
+        /*sw2TraceId=*/ "",
+        /*expectedSW1TraceId=*/ "java-streamwriter header_1:trailer_1",
+        /*expectedSW2TraceId=*/ "java-streamwriter");
+  }
+
+  private void testMultiplexedIngestion(
+      String sw1TraceId, String sw2TraceId, String expectedSW1TraceId, String expectedSW2TraceId)
+      throws Exception {
     try (ConnectionWorker connectionWorker = createMultiplexedConnectionWorker()) {
       long appendCount = 20;
       for (long i = 0; i < appendCount; i++) {
         testBigQueryWrite.addResponse(createAppendResponse(i));
       }
       List<ApiFuture<AppendRowsResponse>> futures = new ArrayList<>();
-      StreamWriter sw1 =
+      // SW1 has not trace id, SW2 does
+      StreamWriter.Builder sw1Builder =
           StreamWriter.newBuilder(TEST_STREAM_1, client)
               .setWriterSchema(createProtoSchema("foo"))
-              .setLocation("us")
-              .build();
-      StreamWriter sw2 =
+              .setLocation("us");
+      StreamWriter.Builder sw2Builder =
           StreamWriter.newBuilder(TEST_STREAM_2, client)
               .setWriterSchema(createProtoSchema("complicate"))
-              .setLocation("us")
-              .build();
+              .setLocation("us");
+      if (!sw1TraceId.isEmpty()) {
+        sw1Builder.setTraceId(sw1TraceId);
+      }
+      if (!sw2TraceId.isEmpty()) {
+        sw2Builder.setTraceId(sw2TraceId);
+      }
+      StreamWriter sw1 = sw1Builder.build();
+      StreamWriter sw2 = sw2Builder.build();
+
       // We do a pattern of:
       // send to stream1, string1
       // send to stream1, string2
@@ -161,12 +188,14 @@ public void testMultiplexedAppendSuccess() throws Exception {
             assertThat(
                     serverRequest.getProtoRows().getWriterSchema().getProtoDescriptor().getName())
                 .isEqualTo("foo");
+            assertThat(serverRequest.getTraceId()).isEqualTo(expectedSW1TraceId);
             break;
           case 1:
             // The write stream is empty until we enter multiplexing.
             assertThat(serverRequest.getWriteStream()).isEqualTo(TEST_STREAM_1);
             // Schema is empty if not at the first request after table switch.
             assertThat(serverRequest.getProtoRows().hasWriterSchema()).isFalse();
+            assertThat(serverRequest.getTraceId()).isEmpty();
             break;
           case 2:
             // Stream name is always populated after multiplexing.
@@ -175,12 +204,14 @@ public void testMultiplexedAppendSuccess() throws Exception {
             assertThat(
                     serverRequest.getProtoRows().getWriterSchema().getProtoDescriptor().getName())
                 .isEqualTo("complicate");
+            assertThat(serverRequest.getTraceId()).isEqualTo(expectedSW2TraceId);
             break;
           case 3:
             // Schema is empty if not at the first request after table switch.
             assertThat(serverRequest.getProtoRows().hasWriterSchema()).isFalse();
             // Stream name is always populated after multiplexing.
             assertThat(serverRequest.getWriteStream()).isEqualTo(TEST_STREAM_2);
+            assertThat(serverRequest.getTraceId()).isEmpty();
             break;
           default: // fall out
             break;
@@ -192,6 +223,9 @@ public void testMultiplexedAppendSuccess() throws Exception {
     }
   }
 
+  @Test
+  public void testMultiplexedAppendSuccess_MixEmptyAndNonEmptyTraceId() throws Exception {}
+
   @Test
   public void testAppendInSameStream_switchSchema() throws Exception {
     try (ConnectionWorker connectionWorker = createMultiplexedConnectionWorker()) {
