1 files changed, 64 insertions, 44 deletions

diff --git a/src/core/channel/census_filter.c b/src/core/channel/census_filter.c
index 9c0c20af22..7e393a01a6 100644
--- a/src/core/channel/census_filter.c
+++ b/src/core/channel/census_filter.c
@@ -49,6 +49,11 @@ typedef struct call_data {
   census_op_id op_id;
   census_rpc_stats stats;
   gpr_timespec start_ts;
+
+  /* recv callback */
+  grpc_stream_op_buffer* recv_ops;
+  void (*on_done_recv)(void* user_data, int success);
+  void* recv_user_data;
 } call_data;
 
 typedef struct channel_data {
@@ -60,57 +65,68 @@ static void init_rpc_stats(census_rpc_stats* stats) {
   stats->cnt = 1;
 }
 
-static void extract_and_annotate_method_tag(grpc_call_op* op, call_data* calld,
+static void extract_and_annotate_method_tag(grpc_stream_op_buffer* sopb,
+                                            call_data* calld,
                                             channel_data* chand) {
   grpc_linked_mdelem* m;
-  for (m = op->data.metadata.list.head; m != NULL; m = m->next) {
-    if (m->md->key == chand->path_str) {
-      gpr_log(GPR_DEBUG, "%s", (const char*)GPR_SLICE_START_PTR(m->md->value->slice));
-      census_add_method_tag(
-          calld->op_id, (const char*)GPR_SLICE_START_PTR(m->md->value->slice));
+  size_t i;
+  for (i = 0; i < sopb->nops; i++) {
+    grpc_stream_op* op = &sopb->ops[i];
+    if (op->type != GRPC_OP_METADATA) continue;
+    for (m = op->data.metadata.list.head; m != NULL; m = m->next) {
+      if (m->md->key == chand->path_str) {
+        gpr_log(GPR_DEBUG, "%s",
+                (const char*)GPR_SLICE_START_PTR(m->md->value->slice));
+        census_add_method_tag(calld->op_id, (const char*)GPR_SLICE_START_PTR(
+                                                m->md->value->slice));
+      }
     }
   }
 }
 
-static void client_call_op(grpc_call_element* elem,
-                           grpc_call_element* from_elem, grpc_call_op* op) {
+static void client_mutate_op(grpc_call_element* elem, grpc_transport_op* op) {
   call_data* calld = elem->call_data;
   channel_data* chand = elem->channel_data;
-  GPR_ASSERT(calld != NULL);
-  GPR_ASSERT(chand != NULL);
-  GPR_ASSERT((calld->op_id.upper != 0) || (calld->op_id.lower != 0));
-  switch (op->type) {
-    case GRPC_SEND_METADATA:
-      extract_and_annotate_method_tag(op, calld, chand);
-      break;
-    case GRPC_RECV_FINISH:
-      /* Should we stop timing the rpc here? */
-      break;
-    default:
-      break;
+  if (op->send_ops) {
+    extract_and_annotate_method_tag(op->send_ops, calld, chand);
   }
-  /* Always pass control up or down the stack depending on op->dir */
+}
+
+static void client_start_transport_op(grpc_call_element* elem,
+                                      grpc_transport_op* op) {
+  call_data* calld = elem->call_data;
+  GPR_ASSERT((calld->op_id.upper != 0) || (calld->op_id.lower != 0));
+  client_mutate_op(elem, op);
   grpc_call_next_op(elem, op);
 }
 
-static void server_call_op(grpc_call_element* elem,
-                           grpc_call_element* from_elem, grpc_call_op* op) {
+static void server_on_done_recv(void* ptr, int success) {
+  grpc_call_element* elem = ptr;
   call_data* calld = elem->call_data;
   channel_data* chand = elem->channel_data;
-  GPR_ASSERT(calld != NULL);
-  GPR_ASSERT(chand != NULL);
-  GPR_ASSERT((calld->op_id.upper != 0) || (calld->op_id.lower != 0));
-  switch (op->type) {
-    case GRPC_RECV_METADATA:
-      extract_and_annotate_method_tag(op, calld, chand);
-      break;
-    case GRPC_SEND_FINISH:
-      /* Should we stop timing the rpc here? */
-      break;
-    default:
-      break;
+  if (success) {
+    extract_and_annotate_method_tag(calld->recv_ops, calld, chand);
   }
-  /* Always pass control up or down the stack depending on op->dir */
+  calld->on_done_recv(calld->recv_user_data, success);
+}
+
+static void server_mutate_op(grpc_call_element* elem, grpc_transport_op* op) {
+  call_data* calld = elem->call_data;
+  if (op->recv_ops) {
+    /* substitute our callback for the op callback */
+    calld->recv_ops = op->recv_ops;
+    calld->on_done_recv = op->on_done_recv;
+    calld->recv_user_data = op->recv_user_data;
+    op->on_done_recv = server_on_done_recv;
+    op->recv_user_data = elem;
+  }
+}
+
+static void server_start_transport_op(grpc_call_element* elem,
+                                      grpc_transport_op* op) {
+  call_data* calld = elem->call_data;
+  GPR_ASSERT((calld->op_id.upper != 0) || (calld->op_id.lower != 0));
+  server_mutate_op(elem, op);
   grpc_call_next_op(elem, op);
 }
 
@@ -128,12 +144,14 @@ static void channel_op(grpc_channel_element* elem,
 }
 
 static void client_init_call_elem(grpc_call_element* elem,
-                                  const void* server_transport_data) {
+                                  const void* server_transport_data,
+                                  grpc_transport_op* initial_op) {
   call_data* d = elem->call_data;
   GPR_ASSERT(d != NULL);
   init_rpc_stats(&d->stats);
   d->start_ts = gpr_now();
   d->op_id = census_tracing_start_op();
+  if (initial_op) client_mutate_op(elem, initial_op);
 }
 
 static void client_destroy_call_elem(grpc_call_element* elem) {
@@ -144,12 +162,14 @@ static void client_destroy_call_elem(grpc_call_element* elem) {
 }
 
 static void server_init_call_elem(grpc_call_element* elem,
-                                  const void* server_transport_data) {
+                                  const void* server_transport_data,
+                                  grpc_transport_op* initial_op) {
   call_data* d = elem->call_data;
   GPR_ASSERT(d != NULL);
   init_rpc_stats(&d->stats);
   d->start_ts = gpr_now();
   d->op_id = census_tracing_start_op();
+  if (initial_op) server_mutate_op(elem, initial_op);
 }
 
 static void server_destroy_call_elem(grpc_call_element* elem) {
@@ -180,11 +200,11 @@ static void destroy_channel_elem(grpc_channel_element* elem) {
 }
 
 const grpc_channel_filter grpc_client_census_filter = {
-    client_call_op, channel_op, sizeof(call_data), client_init_call_elem,
-    client_destroy_call_elem, sizeof(channel_data), init_channel_elem,
-    destroy_channel_elem, "census-client"};
+    client_start_transport_op, channel_op, sizeof(call_data),
+    client_init_call_elem, client_destroy_call_elem, sizeof(channel_data),
+    init_channel_elem, destroy_channel_elem, "census-client"};
 
 const grpc_channel_filter grpc_server_census_filter = {
-    server_call_op, channel_op, sizeof(call_data), server_init_call_elem,
-    server_destroy_call_elem, sizeof(channel_data), init_channel_elem,
-    destroy_channel_elem, "census-server"};
+    server_start_transport_op, channel_op, sizeof(call_data),
+    server_init_call_elem, server_destroy_call_elem, sizeof(channel_data),
+    init_channel_elem, destroy_channel_elem, "census-server"};