feat: include writerId as part of trace id

Author: agrawal-siddharth

java-bigquerystorage/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorker.java

@@ -19,6 +19,7 @@
 import com.google.api.core.NanoClock;
 import com.google.api.core.SettableApiFuture;
 import com.google.api.gax.batching.FlowController;
+import com.google.api.gax.core.GaxProperties;
 import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.api.gax.retrying.TimedAttemptSettings;
@@ -127,11 +128,6 @@ class ConnectionWorker implements AutoCloseable {
    */
   private final FlowController.LimitExceededBehavior limitExceededBehavior;
 
-  /*
-   * TraceId for debugging purpose.
-   */
-  private final String traceId;
-
   /*
    * Enables compression on the wire.
    */
@@ -566,6 +562,7 @@ public ConnectionWorker(
       long maxInflightBytes,
       Duration maxRetryDuration,
       FlowController.LimitExceededBehavior limitExceededBehavior,
+      String clientlibId,
       String traceId,
       @Nullable String compressorName,
       BigQueryWriteSettings clientSettings,
@@ -582,6 +579,7 @@ public ConnectionWorker(
         maxInflightBytes,
         maxRetryDuration,
         limitExceededBehavior,
+        clientlibId,
         traceId,
         compressorName,
         clientSettings,
@@ -599,6 +597,7 @@ public ConnectionWorker(
       long maxInflightBytes,
       Duration maxRetryDuration,
       FlowController.LimitExceededBehavior limitExceededBehavior,
+      String clientlibId,
       String traceId,
       @Nullable String compressorName,
       BigQueryWriteSettings clientSettings,
@@ -615,6 +614,7 @@ public ConnectionWorker(
         maxInflightBytes,
         maxRetryDuration,
         limitExceededBehavior,
+        clientlibId,
         traceId,
         compressorName,
         clientSettings,
@@ -632,6 +632,7 @@ public ConnectionWorker(
       long maxInflightBytes,
       Duration maxRetryDuration,
       FlowController.LimitExceededBehavior limitExceededBehavior,
+      String clientlibId,
       String traceId,
       @Nullable String compressorName,
       BigQueryWriteSettings clientSettings,
@@ -656,14 +657,14 @@ public ConnectionWorker(
     this.maxInflightRequests = maxInflightRequests;
     this.maxInflightBytes = maxInflightBytes;
     this.limitExceededBehavior = limitExceededBehavior;
-    this.traceId = traceId;
+    String fullTraceId = getFullTraceId(clientlibId, this.writerId, traceId);
     this.waitingRequestQueue = new LinkedList<AppendRequestAndResponse>();
     this.inflightRequestQueue = new LinkedList<AppendRequestAndResponse>();
     this.compressorName = compressorName;
     this.retrySettings = retrySettings;
     this.requestProfilerHook = new RequestProfiler.RequestProfilerHook(enableRequestProfiler);
     this.telemetryMetrics =
-        new TelemetryMetrics(this, enableOpenTelemetry, getTableName(), writerId, traceId);
+        new TelemetryMetrics(this, enableOpenTelemetry, getTableName(), writerId, fullTraceId);
     this.healthCheckMetrics = new HealthCheckMetrics();
     this.isMultiplexing = isMultiplexing;
 
@@ -991,6 +992,17 @@ static long calculateSleepTimeMilli(long retryCount) {
     return (long) Math.min(Math.pow(2, retryCount) * 50, 60000);
   }
 
+  @VisibleForTesting
+  static String getFullTraceId(String clientlibId, String writerId, String traceId) {
+    String clientWithVersion =
+        clientlibId + ":" + GaxProperties.getLibraryVersion(StreamWriter.class) + ":" + writerId;
+    if (traceId == null || traceId.isEmpty()) {
+      return clientWithVersion;
+    } else {
+      return clientWithVersion + " " + traceId;
+    }
+  }
+
   @VisibleForTesting
   void setTestOnlyAppendLoopSleepTime(long testOnlyAppendLoopSleepTime) {
     this.testOnlyAppendLoopSleepTime = testOnlyAppendLoopSleepTime;
@@ -1271,7 +1283,11 @@ 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);
-          originalRequestBuilder.setTraceId(wrapper.streamWriter.getFullTraceId());
+          originalRequestBuilder.setTraceId(
+              getFullTraceId(
+                  wrapper.streamWriter.getClientId(),
+                  this.writerId,
+                  wrapper.streamWriter.getTraceId()));
         } else if (!isMultiplexing) {
           // If we are not in multiplexing and not in the first request, clear the stream name.
           originalRequestBuilder.clearWriteStream();

java-bigquerystorage/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerPool.java

@@ -317,7 +317,8 @@ private ConnectionWorker createOrReuseConnectionWorker(
           streamWriter.getStreamName(),
           streamWriter.getLocation(),
           streamWriter.getWriterSchema(),
-          streamWriter.getFullTraceId());
+          streamWriter.getClientId(),
+          streamWriter.getTraceId());
     } else {
       ConnectionWorker existingBestConnection =
           pickBestLoadConnection(
@@ -337,7 +338,8 @@ private ConnectionWorker createOrReuseConnectionWorker(
             streamWriter.getStreamName(),
             streamWriter.getLocation(),
             streamWriter.getWriterSchema(),
-            streamWriter.getFullTraceId());
+            streamWriter.getClientId(),
+            streamWriter.getTraceId());
       } else {
         // Stick to the original connection if all the connections are overwhelmed.
         if (existingConnectionWorker != null) {
@@ -396,7 +398,8 @@ private ConnectionWorker createConnectionWorker(
       String streamName,
       String location,
       AppendFormats.AppendRowsSchema writeSchema,
-      String fullTraceId)
+      String clientId,
+      String traceId)
       throws IOException {
     if (enableTesting) {
       // Though atomic integer is super lightweight, add extra if check in case adding future logic.
@@ -411,7 +414,8 @@ private ConnectionWorker createConnectionWorker(
             maxInflightBytes,
             maxRetryDuration,
             limitExceededBehavior,
-            fullTraceId,
+            clientId,
+            traceId,
             compressorName,
             clientSettings,
             retrySettings,

java-bigquerystorage/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java

@@ -19,7 +19,6 @@
 import com.google.api.gax.batching.FlowController;
 import com.google.api.gax.core.CredentialsProvider;
 import com.google.api.gax.core.ExecutorProvider;
-import com.google.api.gax.core.GaxProperties;
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.api.gax.rpc.TransportChannelProvider;
 import com.google.auth.Credentials;
@@ -93,8 +92,11 @@ private static Cache<String, String> allocateProjectLocationCache() {
    */
   private final String streamName;
 
-  /** This is the library version may or may not include library version id. */
-  private final String fullTraceId;
+  /** Client id of the writer. */
+  private final String clientId;
+
+  /** User provided trace id. */
+  private final String traceId;
 
   /** Every writer has a fixed proto schema or arrow schema. */
   private final AppendRowsSchema writerSchema;
@@ -256,7 +258,8 @@ private StreamWriter(Builder builder) throws IOException {
     BigQueryWriteSettings clientSettings = getBigQueryWriteSettings(builder);
     this.requestProfilerHook =
         new RequestProfiler.RequestProfilerHook(builder.enableRequestProfiler);
-    this.fullTraceId = builder.getFullTraceId();
+    this.clientId = builder.clientId;
+    this.traceId = builder.traceId;
     if (builder.enableRequestProfiler) {
       // Request profiler is enabled on singleton level, from now on a periodical flush will be
       // started
@@ -275,7 +278,8 @@ private StreamWriter(Builder builder) throws IOException {
                   builder.maxInflightBytes,
                   builder.maxRetryDuration,
                   builder.limitExceededBehavior,
-                  builder.getFullTraceId(),
+                  builder.clientId,
+                  builder.traceId,
                   builder.compressorName,
                   clientSettings,
                   builder.retrySettings,
@@ -395,8 +399,12 @@ static void recreateProjectLocationCache(long durationExpireMillis) {
     projectAndDatasetToLocation = allocateProjectLocationCache();
   }
 
-  String getFullTraceId() {
-    return fullTraceId;
+  String getClientId() {
+    return clientId;
+  }
+
+  String getTraceId() {
+    return traceId;
   }
 
   AppendRowsRequest.MissingValueInterpretation getDefaultValueInterpretation() {
@@ -1119,18 +1127,6 @@ public Builder setRetrySettings(RetrySettings retrySettings) {
     public StreamWriter build() throws IOException {
       return new StreamWriter(this);
     }
-
-    String getFullTraceId() {
-      String clientWithVersion =
-          GaxProperties.getLibraryVersion(StreamWriter.class).isEmpty()
-              ? clientId
-              : clientId + ":" + GaxProperties.getLibraryVersion(StreamWriter.class);
-      if (traceId == null || traceId.isEmpty()) {
-        return clientWithVersion;
-      } else {
-        return clientWithVersion + " " + traceId;
-      }
-    }
   }
 
   private String generateRequestUniqueId() {