fix: add client id and update trace id population for StreamWriter and JsonWriter (#2389)

yirutang · gcf-owl-bot[bot] · web-flow · commit 4258af42f8ab · 2024-02-06T17:52:45.000-08:00
* fix: add a bit more message so customers are not going to be scaried by load shedding errors * . * fix:change the client id string, adding client id for stream writer * . * . * . * . * . * . * . * . * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
diff --git a/google-cloud-bigquerystorage/clirr-ignored-differences.xml b/google-cloud-bigquerystorage/clirr-ignored-differences.xml
@@ -188,5 +188,10 @@
     <className>com/google/cloud/bigquery/storage/v1/StreamWriter$Builder</className>
     <method>com.google.cloud.bigquery.storage.v1.StreamWriter$Builder setRetryFirstDelay(org.threeten.bp.Duration)</method>
   </difference>
+  <difference>
+    <differenceType>7002</differenceType>
+    <className>com/google/cloud/bigquery/storage/v1/SchemaAwareStreamWriter$Builder</className>
+    <method>com.google.cloud.bigquery.storage.v1.SchemaAwareStreamWriter$Builder setTraceIdBase(java.lang.String)</method>
+  </difference>
 </differences>
 
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/JsonStreamWriter.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/JsonStreamWriter.java
@@ -37,6 +37,7 @@
  */
 public class JsonStreamWriter implements AutoCloseable {
   private final SchemaAwareStreamWriter<Object> schemaAwareStreamWriter;
+  private static final String CLIENT_ID = "java-jsonwriter";
 
   /**
    * Constructs the JsonStreamWriter
@@ -227,8 +228,7 @@ public static final class Builder {
     private final SchemaAwareStreamWriter.Builder<Object> schemaAwareStreamWriterBuilder;
 
     private Builder(SchemaAwareStreamWriter.Builder<Object> schemaAwareStreamWriterBuilder) {
-      this.schemaAwareStreamWriterBuilder =
-          schemaAwareStreamWriterBuilder.setTraceIdBase("JsonWriter");
+      this.schemaAwareStreamWriterBuilder = schemaAwareStreamWriterBuilder.setClientId(CLIENT_ID);
     }
 
     /**
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/SchemaAwareStreamWriter.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/SchemaAwareStreamWriter.java
@@ -94,14 +94,14 @@ private SchemaAwareStreamWriter(Builder<T> builder)
         builder.executorProvider,
         builder.endpoint,
         builder.flowControlSettings,
-        builder.traceIdBase,
         builder.traceId,
         builder.compressorName,
         builder.retrySettings);
     streamWriterBuilder.setEnableConnectionPool(builder.enableConnectionPool);
     streamWriterBuilder.setLocation(builder.location);
     streamWriterBuilder.setDefaultMissingValueInterpretation(
         builder.defaultMissingValueInterpretation);
+    streamWriterBuilder.setClientId(builder.clientId);
     this.streamWriter = streamWriterBuilder.build();
     this.streamName = builder.streamName;
     this.tableSchema = builder.tableSchema;
@@ -282,7 +282,6 @@ private void setStreamWriterSettings(
       @Nullable ExecutorProvider executorProvider,
       @Nullable String endpoint,
       @Nullable FlowControlSettings flowControlSettings,
-      @Nullable String traceIdBase,
       @Nullable String traceId,
       @Nullable String compressorName,
       @Nullable RetrySettings retrySettings) {
@@ -298,18 +297,8 @@ private void setStreamWriterSettings(
     if (endpoint != null) {
       streamWriterBuilder.setEndpoint(endpoint);
     }
-    if (traceIdBase != null) {
-      if (traceId != null) {
-        streamWriterBuilder.setTraceId(traceIdBase + "_" + traceId);
-      } else {
-        streamWriterBuilder.setTraceId(traceIdBase + ":null");
-      }
-    } else {
-      if (traceId != null) {
-        streamWriterBuilder.setTraceId("SchemaAwareStreamWriter_" + traceId);
-      } else {
-        streamWriterBuilder.setTraceId("SchemaAwareStreamWriter:null");
-      }
+    if (traceId != null) {
+      streamWriterBuilder.setTraceId(traceId);
     }
     if (flowControlSettings != null) {
       if (flowControlSettings.getMaxOutstandingRequestBytes() != null) {
@@ -445,6 +434,7 @@ public static final class Builder<T> {
 
     private AppendRowsRequest.MissingValueInterpretation defaultMissingValueInterpretation =
         MissingValueInterpretation.MISSING_VALUE_INTERPRETATION_UNSPECIFIED;
+    private String clientId;
 
     private static final String streamPatternString =
         "(projects/[^/]+/datasets/[^/]+/tables/[^/]+)/streams/[^/]+";
@@ -581,14 +571,8 @@ public Builder<T> setTraceId(String traceId) {
       return this;
     }
 
-    /**
-     * Setter for a traceIdBase to help identify traffic origin.
-     *
-     * @param traceIdBase
-     * @return Builder
-     */
-    public Builder<T> setTraceIdBase(String traceIdBase) {
-      this.traceIdBase = Preconditions.checkNotNull(traceIdBase, "TraceIdBase is null.");
+    Builder<T> setClientId(String clientId) {
+      this.clientId = Preconditions.checkNotNull(clientId, "ClientId is null.");
       return this;
     }
 
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
@@ -229,7 +229,7 @@ private StreamWriter(Builder builder) throws IOException {
                   builder.maxInflightBytes,
                   builder.maxRetryDuration,
                   builder.limitExceededBehavior,
-                  builder.traceId,
+                  builder.getFullTraceId(),
                   builder.compressorName,
                   clientSettings,
                   builder.retrySettings));
@@ -295,7 +295,7 @@ private StreamWriter(Builder builder) throws IOException {
                         builder.maxInflightBytes,
                         builder.maxRetryDuration,
                         builder.limitExceededBehavior,
-                        builder.traceId,
+                        builder.getFullTraceId(),
                         builder.compressorName,
                         client.getSettings(),
                         builder.retrySettings);
@@ -376,12 +376,12 @@ static BigQueryWriteSettings getBigQueryWriteSettings(Builder builder) throws IO
   private void validateFetchedConnectonPool(StreamWriter.Builder builder) {
     String storedTraceId =
         this.singleConnectionOrConnectionPool.connectionWorkerPool().getTraceId();
-    if (!Objects.equals(storedTraceId, builder.traceId)) {
+    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.traceId));
+              storedTraceId, builder.getFullTraceId()));
     }
     FlowController.LimitExceededBehavior storedLimitExceededBehavior =
         singleConnectionOrConnectionPool.connectionWorkerPool().limitExceededBehavior();
@@ -629,6 +629,8 @@ public static final class Builder {
 
     private String traceId = null;
 
+    private String clientId = "java-streamwriter";
+
     private TableSchema updatedTableSchema = null;
 
     private String location = null;
@@ -730,6 +732,15 @@ public Builder setTraceId(String traceId) {
       return this;
     }
 
+    /**
+     * Sets the client id of the writer, for example, JsonStreamWriter has the client id of
+     * "java-jsonwriter".
+     */
+    Builder setClientId(String clientId) {
+      this.clientId = clientId;
+      return this;
+    }
+
     /** Location of the table this stream writer is targeting. */
     public Builder setLocation(String location) {
       this.location = location;
@@ -811,5 +822,13 @@ public Builder setRetrySettings(RetrySettings retrySettings) {
     public StreamWriter build() throws IOException {
       return new StreamWriter(this);
     }
+
+    String getFullTraceId() {
+      if (traceId == null) {
+        return clientId;
+      } else {
+        return clientId + " " + traceId;
+      }
+    }
   }
 }
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/JsonStreamWriterTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/JsonStreamWriterTest.java
@@ -237,7 +237,7 @@ public void testSingleAppendSimpleJson() throws Exception {
               .getSerializedRows(0),
           expectedProto.toByteString());
       assertEquals(
-          testBigQueryWrite.getAppendRequests().get(0).getTraceId(), "JsonWriter_test:empty");
+          "java-jsonwriter test:empty", testBigQueryWrite.getAppendRequests().get(0).getTraceId());
     }
   }
 
@@ -257,7 +257,7 @@ public void testFlexibleColumnAppend() throws Exception {
     jsonArr.put(flexible);
 
     try (JsonStreamWriter writer =
-        getTestJsonStreamWriterBuilder(TEST_STREAM, tableSchema).setTraceId("test:empty").build()) {
+        getTestJsonStreamWriterBuilder(TEST_STREAM, tableSchema).build()) {
 
       testBigQueryWrite.addResponse(
           AppendRowsResponse.newBuilder()
@@ -284,8 +284,7 @@ public void testFlexibleColumnAppend() throws Exception {
               .getRows()
               .getSerializedRows(0),
           expectedProto.toByteString());
-      assertEquals(
-          testBigQueryWrite.getAppendRequests().get(0).getTraceId(), "JsonWriter_test:empty");
+      assertEquals("java-jsonwriter", testBigQueryWrite.getAppendRequests().get(0).getTraceId());
     }
   }
 
@@ -507,7 +506,7 @@ public void testSingleAppendMultipleSimpleJson() throws Exception {
               .getProtoRows()
               .getRows()
               .getSerializedRowsCount());
-      assertEquals(testBigQueryWrite.getAppendRequests().get(0).getTraceId(), "JsonWriter:null");
+      assertEquals("java-jsonwriter", testBigQueryWrite.getAppendRequests().get(0).getTraceId());
       for (int i = 0; i < 4; i++) {
         assertEquals(
             testBigQueryWrite
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java
@@ -286,13 +286,13 @@ private void verifyAppendRequests(long appendCount) {
       if (i == 0) {
         // First request received by server should have schema and stream name.
         assertTrue(serverRequest.getProtoRows().hasWriterSchema());
-        assertEquals(serverRequest.getWriteStream(), TEST_STREAM_1);
-        assertEquals(serverRequest.getTraceId(), TEST_TRACE_ID);
+        assertEquals(TEST_STREAM_1, serverRequest.getWriteStream());
+        assertEquals("java-streamwriter " + TEST_TRACE_ID, serverRequest.getTraceId());
       } else {
         // Following request should not have schema and stream name.
         assertFalse(serverRequest.getProtoRows().hasWriterSchema());
-        assertEquals(serverRequest.getWriteStream(), "");
-        assertEquals(serverRequest.getTraceId(), "");
+        assertEquals("", serverRequest.getWriteStream());
+        assertEquals("", serverRequest.getTraceId());
       }
     }
   }
