feat: populate client_stats in AppendRowsRequest

Author: agrawal-siddharth

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

@@ -1289,6 +1289,11 @@ private void appendLoop() {
         }
         firstRequestForTableOrSchemaSwitch = false;
 
+        ClientStats clientStats = createClientStats(wrapper);
+        if (!clientStats.equals(ClientStats.getDefaultInstance())) {
+          originalRequestBuilder.setClientStats(clientStats);
+        }
+
         requestProfilerHook.startOperation(
             RequestProfiler.OperationName.RESPONSE_LATENCY, requestUniqueId);
 
@@ -1307,6 +1312,20 @@ private void appendLoop() {
     cleanupConnectionAndRequests(/* avoidBlocking= */ false);
   }
 
+  private ClientStats createClientStats(AppendRequestAndResponse requestWrapper) {
+    ClientStats.Builder clientStatsBuilder = ClientStats.newBuilder();
+    ClientStats.RequestStats.Builder requestStatsBuilder = ClientStats.RequestStats.newBuilder();
+    boolean requestStatsModified = false;
+    if (requestWrapper.requestSendTimeStamp != null) {
+      requestStatsBuilder.setSendTimeMillis(requestWrapper.requestSendTimeStamp.toEpochMilli());
+      requestStatsModified = true;
+    }
+    if (requestStatsModified) {
+      clientStatsBuilder.setRequestStats(requestStatsBuilder.build());
+    }
+    return clientStatsBuilder.build();
+  }
+
   @Nullable
   private AppendRowsSchema getSchema(AppendRowsRequest request) {
     if (request.hasProtoRows()) {

java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerTest.java

@@ -32,6 +32,7 @@
 import com.google.cloud.bigquery.storage.test.Test.FooType;
 import com.google.cloud.bigquery.storage.test.Test.InnerType;
 import com.google.cloud.bigquery.storage.v1.ConnectionWorker.Load;
+import com.google.common.collect.Iterables;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.DescriptorProtos;
 import com.google.protobuf.Int64Value;
@@ -1722,6 +1723,61 @@ void testEarliestSendTime_retryScenario() throws Exception {
 
       // Once exhausted, state should be null
       assertThat(connectionWorker.getEarliestSendTime()).isNull();
+
+      AppendRowsRequest initialRequest = testBigQueryWrite.getAppendRequests().get(0);
+      AppendRowsRequest retryRequest = testBigQueryWrite.getAppendRequests().get(1);
+      assertThat(retryRequest.getClientStats().getRequestStats().getSendTimeMillis())
+          .isAtLeast(initialRequest.getClientStats().getRequestStats().getSendTimeMillis());
+    }
+  }
+
+  @Test
+  void testClientStats_sendTimeMillisPopulated() throws Exception {
+    // Arrange
+    ProtoSchema schema1 = createProtoSchema("foo");
+    try (StreamWriter sw1 =
+            StreamWriter.newBuilder(TEST_STREAM_1, client)
+                .setLocation("us")
+                .setWriterSchema(schema1)
+                .build();
+        ConnectionWorker connectionWorker = createConnectionWorker()) {
+      testBigQueryWrite.addResponse(createAppendResponse(0));
+
+      // Act
+      long beforeSendMillis = Instant.now().toEpochMilli();
+      ApiFuture<AppendRowsResponse> future =
+          sendTestMessage(connectionWorker, sw1, createFooProtoRows(new String[] {"0"}), 0);
+      future.get();
+      long afterSendMillis = Instant.now().toEpochMilli();
+
+      // Assert
+      AppendRowsRequest sentRequest =
+          Iterables.getOnlyElement(testBigQueryWrite.getAppendRequests());
+      assertThat(sentRequest.hasClientStats()).isTrue();
+      assertThat(sentRequest.getClientStats().hasRequestStats()).isTrue();
+      assertThat(sentRequest.getClientStats().getRequestStats().hasSendTimeMillis()).isTrue();
+      long sendTimeMillis = sentRequest.getClientStats().getRequestStats().getSendTimeMillis();
+      assertThat(sendTimeMillis).isAtLeast(beforeSendMillis);
+      assertThat(sendTimeMillis).isAtMost(afterSendMillis);
+
+      // Arrange retry scenario
+      testBigQueryWrite.addResponse(
+          new DummyResponseSupplierWillFailThenSucceed(
+              new FakeBigQueryWriteImpl.Response(createAppendResponse(1)),
+              /* totalFailCount= */ 1,
+              com.google.rpc.Status.newBuilder().setCode(Code.INTERNAL.ordinal()).build()));
+
+      // Act
+      ApiFuture<AppendRowsResponse> retryFuture =
+          sendTestMessage(connectionWorker, sw1, createFooProtoRows(new String[] {"1"}), 1);
+      retryFuture.get();
+
+      // Assert
+      assertThat(testBigQueryWrite.getAppendRequests()).hasSize(3);
+      AppendRowsRequest initialRequest = testBigQueryWrite.getAppendRequests().get(1);
+      AppendRowsRequest retryRequest = testBigQueryWrite.getAppendRequests().get(2);
+      assertThat(retryRequest.getClientStats().getRequestStats().getSendTimeMillis())
+          .isAtLeast(initialRequest.getClientStats().getRequestStats().getSendTimeMillis());
     }
   }
 }