Merge pull request #17331 from yashykt/timestamplength

Add the byte offset for the RPC that is traced in Timestamps

4 files changed, 9 insertions, 3 deletions

diff --git a/src/core/ext/transport/chttp2/transport/context_list.cc b/src/core/ext/transport/chttp2/transport/context_list.cc
index 4acd0c9583..f30d41c332 100644
--- a/src/core/ext/transport/chttp2/transport/context_list.cc
+++ b/src/core/ext/transport/chttp2/transport/context_list.cc
@@ -32,6 +32,7 @@ void ContextList::Execute(void* arg, grpc_core::Timestamps* ts,
   while (head != nullptr) {
     if (error == GRPC_ERROR_NONE && ts != nullptr) {
       if (write_timestamps_callback_g) {
+        ts->byte_offset = static_cast<uint32_t>(head->byte_offset_);
         write_timestamps_callback_g(head->s_->context, ts);
       }
     }
diff --git a/src/core/ext/transport/chttp2/transport/context_list.h b/src/core/ext/transport/chttp2/transport/context_list.h
index 68d11e94d8..d870107749 100644
--- a/src/core/ext/transport/chttp2/transport/context_list.h
+++ b/src/core/ext/transport/chttp2/transport/context_list.h
@@ -50,6 +50,7 @@ class ContextList {
     /* Create a new element in the list and add it at the front */
     ContextList* elem = grpc_core::New<ContextList>();
     elem->s_ = s;
+    elem->byte_offset_ = s->byte_counter;
     elem->next_ = *head;
     *head = elem;
   }
@@ -61,6 +62,7 @@ class ContextList {
  private:
   grpc_chttp2_stream* s_ = nullptr;
   ContextList* next_ = nullptr;
+  size_t byte_offset_ = 0;
 };
 
 void grpc_http2_set_write_timestamps_callback(
diff --git a/src/core/ext/transport/chttp2/transport/internal.h b/src/core/ext/transport/chttp2/transport/internal.h
index aeaa4935ad..6aa68f5d4a 100644
--- a/src/core/ext/transport/chttp2/transport/internal.h
+++ b/src/core/ext/transport/chttp2/transport/internal.h
@@ -646,6 +646,8 @@ struct grpc_chttp2_stream {
   bool traced = false;
   /** gRPC header bytes that are already decompressed */
   size_t decompressed_header_bytes = 0;
+  /** Byte counter for number of bytes written */
+  size_t byte_counter = 0;
 };
 
 /** Transport writing call flow:
diff --git a/src/core/ext/transport/chttp2/transport/writing.cc b/src/core/ext/transport/chttp2/transport/writing.cc
index 77320b496f..265d3365d3 100644
--- a/src/core/ext/transport/chttp2/transport/writing.cc
+++ b/src/core/ext/transport/chttp2/transport/writing.cc
@@ -363,6 +363,7 @@ class DataSendContext {
     grpc_chttp2_encode_data(s_->id, &s_->compressed_data_buffer, send_bytes,
                             is_last_frame_, &s_->stats.outgoing, &t_->outbuf);
     s_->flow_control->SentData(send_bytes);
+    s_->byte_counter += send_bytes;
     if (s_->compressed_data_buffer.length == 0) {
       s_->sending_bytes += s_->uncompressed_data_size;
     }
@@ -488,9 +489,6 @@ class StreamWriteContext {
       return;  // early out: nothing to do
     }
 
-    if (s_->traced && grpc_endpoint_can_track_err(t_->ep)) {
-      grpc_core::ContextList::Append(&t_->cl, s_);
-    }
     while ((s_->flow_controlled_buffer.length > 0 ||
             s_->compressed_data_buffer.length > 0) &&
            data_send_context.max_outgoing() > 0) {
@@ -500,6 +498,9 @@ class StreamWriteContext {
         data_send_context.CompressMoreBytes();
       }
     }
+    if (s_->traced && grpc_endpoint_can_track_err(t_->ep)) {
+      grpc_core::ContextList::Append(&t_->cl, s_);
+    }
     write_context_->ResetPingClock();
     if (data_send_context.is_last_frame()) {
       SentLastFrame();