Merge pull request #5699 from ctiller/accounting

Introduce byte counting to the transport API

3 files changed, 44 insertions, 0 deletions

diff --git a/src/core/lib/surface/call.c b/src/core/lib/surface/call.c
index d63a4a7401..0219d7e6d0 100644
--- a/src/core/lib/surface/call.c
+++ b/src/core/lib/surface/call.c
@@ -174,6 +174,9 @@ struct grpc_call {
   /* Received call statuses from various sources */
   received_status status[STATUS_SOURCE_COUNT];
 
+  /* Call stats: only valid after trailing metadata received */
+  grpc_transport_stream_stats stats;
+
   /* Compression algorithm for the call */
   grpc_compression_algorithm compression_algorithm;
   /* Supported encodings (compression algorithms), a bitset */
@@ -1371,6 +1374,7 @@ static grpc_call_error call_start_batch(grpc_exec_ctx *exec_ctx,
         bctl->recv_final_op = 1;
         stream_op.recv_trailing_metadata =
             &call->metadata_batch[1 /* is_receiving */][1 /* is_trailing */];
+        stream_op.collect_stats = &call->stats;
         break;
       case GRPC_OP_RECV_CLOSE_ON_SERVER:
         /* Flag validation: currently allow no flags */
@@ -1392,6 +1396,7 @@ static grpc_call_error call_start_batch(grpc_exec_ctx *exec_ctx,
         bctl->recv_final_op = 1;
         stream_op.recv_trailing_metadata =
             &call->metadata_batch[1 /* is_receiving */][1 /* is_trailing */];
+        stream_op.collect_stats = &call->stats;
         break;
     }
   }
diff --git a/src/core/lib/transport/transport.c b/src/core/lib/transport/transport.c
index 18256aae5e..411e4f140a 100644
--- a/src/core/lib/transport/transport.c
+++ b/src/core/lib/transport/transport.c
@@ -35,6 +35,7 @@
 #include <grpc/support/alloc.h>
 #include <grpc/support/atm.h>
 #include <grpc/support/log.h>
+#include <grpc/support/sync.h>
 #include "src/core/lib/transport/transport_impl.h"
 
 #ifdef GRPC_STREAM_REFCOUNT_DEBUG
@@ -76,6 +77,24 @@ void grpc_stream_ref_init(grpc_stream_refcount *refcount, int initial_refs,
   grpc_closure_init(&refcount->destroy, cb, cb_arg);
 }
 
+static void move64(uint64_t *from, uint64_t *to) {
+  *to += *from;
+  *from = 0;
+}
+
+void grpc_transport_move_one_way_stats(grpc_transport_one_way_stats *from,
+                                       grpc_transport_one_way_stats *to) {
+  move64(&from->framing_bytes, &to->framing_bytes);
+  move64(&from->data_bytes, &to->data_bytes);
+  move64(&from->header_bytes, &to->header_bytes);
+}
+
+void grpc_transport_move_stats(grpc_transport_stream_stats *from,
+                               grpc_transport_stream_stats *to) {
+  grpc_transport_move_one_way_stats(&from->incoming, &to->incoming);
+  grpc_transport_move_one_way_stats(&from->outgoing, &to->outgoing);
+}
+
 size_t grpc_transport_stream_size(grpc_transport *transport) {
   return transport->vtable->sizeof_stream;
 }
diff --git a/src/core/lib/transport/transport.h b/src/core/lib/transport/transport.h
index e98cfe9515..c253b190e7 100644
--- a/src/core/lib/transport/transport.h
+++ b/src/core/lib/transport/transport.h
@@ -78,6 +78,23 @@ void grpc_stream_unref(grpc_exec_ctx *exec_ctx, grpc_stream_refcount *refcount);
   grpc_stream_ref_init(rc, ir, cb, cb_arg)
 #endif
 
+typedef struct {
+  uint64_t framing_bytes;
+  uint64_t data_bytes;
+  uint64_t header_bytes;
+} grpc_transport_one_way_stats;
+
+typedef struct grpc_transport_stream_stats {
+  grpc_transport_one_way_stats incoming;
+  grpc_transport_one_way_stats outgoing;
+} grpc_transport_stream_stats;
+
+void grpc_transport_move_one_way_stats(grpc_transport_one_way_stats *from,
+                                       grpc_transport_one_way_stats *to);
+
+void grpc_transport_move_stats(grpc_transport_stream_stats *from,
+                               grpc_transport_stream_stats *to);
+
 /* Transport stream op: a set of operations to perform on a transport
    against a single stream */
 typedef struct grpc_transport_stream_op {
@@ -104,6 +121,9 @@ typedef struct grpc_transport_stream_op {
    */
   grpc_metadata_batch *recv_trailing_metadata;
 
+  /** Collect any stats into provided buffer, zero internal stat counters */
+  grpc_transport_stream_stats *collect_stats;
+
   /** Should be enqueued when all requested operations (excluding recv_message
       and recv_initial_metadata which have their own closures) in a given batch
       have been completed. */