Adding support for propagating context and retrieving span from context. (#54)

Vizerai · web-flow · commit 66ffd233e3b7 · 2018-02-26T17:03:09.000-08:00
Adding support for propagating context and retrieving span from context.
diff --git a/opencensus/plugins/grpc/internal/client_filter.cc b/opencensus/plugins/grpc/internal/client_filter.cc
@@ -83,6 +83,12 @@ void CensusClientCallData::OnDoneRecvMessageCb(void *user_data,
 void CensusClientCallData::StartTransportStreamOpBatch(
     grpc_call_element *elem, grpc::TransportStreamOpBatch *op) {
   if (op->send_initial_metadata() != nullptr) {
+    // TODO(jsking): Move GenerateClientContext to init callback after
+    // grpc has been changed to populate census context at call initialization.
+    census_context *ctxt = op->get_census_context();
+    GenerateClientContext(
+        method_, &context_,
+        (ctxt == nullptr) ? nullptr : reinterpret_cast<CensusContext *>(ctxt));
     char tracing_buf[kMaxTracingLen];
     size_t tracing_len =
         context_.TraceContextSerialize(tracing_buf, kMaxTracingLen);
@@ -129,21 +135,20 @@ void CensusClientCallData::StartTransportStreamOpBatch(
 
 grpc_error *CensusClientCallData::Init(grpc_call_element *elem,
                                        const grpc_call_element_args *args) {
-  method_ = grpc_slice_ref_internal(args->path);
+  path_ = grpc_slice_ref_internal(args->path);
   start_time_ = absl::Now();
   const char *method_str =
-      GPR_SLICE_IS_EMPTY(method_)
+      GPR_SLICE_IS_EMPTY(path_)
           ? ""
-          : reinterpret_cast<const char *>(GRPC_SLICE_START_PTR(method_));
-  method_size_ = GRPC_SLICE_IS_EMPTY(method_) ? 0 : GRPC_SLICE_LENGTH(method_);
-  absl::string_view method(method_str, method_size_);
-  GenerateClientContext(method, &context_);
+          : reinterpret_cast<const char *>(GRPC_SLICE_START_PTR(path_));
+  method_ = absl::string_view(
+      method_str, GRPC_SLICE_IS_EMPTY(path_) ? 0 : GRPC_SLICE_LENGTH(path_));
   GRPC_CLOSURE_INIT(&on_done_recv_message_, OnDoneRecvMessageCb, elem,
                     grpc_schedule_on_exec_ctx);
   GRPC_CLOSURE_INIT(&on_done_recv_trailing_metadata_,
                     OnDoneRecvTrailingMetadataCb, elem,
                     grpc_schedule_on_exec_ctx);
-  stats::Record({{RpcClientStartedCount(), 1}}, {{kMethodTagKey, method}});
+  stats::Record({{RpcClientStartedCount(), 1}}, {{kMethodTagKey, method_}});
   return GRPC_ERROR_NONE;
 }
 
@@ -164,11 +169,10 @@ void CensusClientCallData::Destroy(grpc_call_element *elem,
        {RpcClientFinishedCount(), 1},
        {RpcClientRequestCount(), sent_message_count_},
        {RpcClientResponseCount(), recv_message_count_}},
-      {{kMethodTagKey, absl::string_view(reinterpret_cast<char *>(
-                                             GRPC_SLICE_START_PTR(method_)),
-                                         method_size_)},
+      {{kMethodTagKey, method_},
        {kStatusTagKey, StatusCodeToString(final_info->final_status)}});
-  grpc_slice_unref_internal(method_);
+  grpc_slice_unref_internal(path_);
+  context_.EndSpan();
 }
 
 }  // namespace opencensus
diff --git a/opencensus/plugins/grpc/internal/client_filter.h b/opencensus/plugins/grpc/internal/client_filter.h
@@ -29,8 +29,7 @@ namespace opencensus {
 class CensusClientCallData : public grpc::CallData {
  public:
   CensusClientCallData()
-      : method_size_(0),
-        recv_trailing_metadata_(nullptr),
+      : recv_trailing_metadata_(nullptr),
         initial_on_done_recv_trailing_metadata_(nullptr),
         initial_on_done_recv_message_(nullptr),
         elapsed_time_(0),
@@ -39,7 +38,7 @@ class CensusClientCallData : public grpc::CallData {
         sent_message_count_(0) {
     memset(&stats_bin_, 0, sizeof(grpc_linked_mdelem));
     memset(&tracing_bin_, 0, sizeof(grpc_linked_mdelem));
-    memset(&method_, 0, sizeof(grpc_slice));
+    memset(&path_, 0, sizeof(grpc_slice));
     memset(&on_done_recv_trailing_metadata_, 0, sizeof(grpc_closure));
     memset(&on_done_recv_message_, 0, sizeof(grpc_closure));
   }
@@ -65,9 +64,8 @@ class CensusClientCallData : public grpc::CallData {
   grpc_linked_mdelem stats_bin_;
   grpc_linked_mdelem tracing_bin_;
   // Client method.
-  grpc_slice method_;
-  // Client method size.
-  size_t method_size_;
+  absl::string_view method_;
+  grpc_slice path_;
   // The recv trailing metadata callbacks.
   grpc_metadata_batch *recv_trailing_metadata_;
   grpc_closure *initial_on_done_recv_trailing_metadata_;
diff --git a/opencensus/plugins/grpc/internal/filter.cc b/opencensus/plugins/grpc/internal/filter.cc
@@ -26,9 +26,17 @@ void GenerateServerContext(absl::string_view tracing, absl::string_view stats,
   new (context) CensusContext(method, parent_ctx);
 }
 
-void GenerateClientContext(absl::string_view method, CensusContext *context) {
-  // TODO: Change this so that it can be set to root or non-root span.
-  new (context) CensusContext(method);
+void GenerateClientContext(absl::string_view method, CensusContext *ctxt,
+                           CensusContext *parent_ctxt) {
+  if (parent_ctxt != nullptr) {
+    trace::SpanContext span_ctxt = parent_ctxt->Context();
+    trace::Span span = parent_ctxt->Span();
+    if (span_ctxt.IsValid()) {
+      new (ctxt) CensusContext(method, &span);
+      return;
+    }
+  }
+  new (ctxt) CensusContext(method);
 }
 
 size_t ServerStatsSerialize(uint64_t server_elapsed_time, char *buf,
@@ -54,6 +62,10 @@ trace::SpanContext SpanContextFromCensusContext(const census_context *ctxt) {
   return reinterpret_cast<const CensusContext *>(ctxt)->Context();
 }
 
+trace::Span SpanFromCensusContext(const census_context *ctxt) {
+  return reinterpret_cast<const CensusContext *>(ctxt)->Span();
+}
+
 absl::string_view StatusCodeToString(grpc_status_code code) {
   switch (code) {
     case GRPC_STATUS_OK:
diff --git a/opencensus/plugins/grpc/internal/filter.h b/opencensus/plugins/grpc/internal/filter.h
@@ -38,8 +38,11 @@ class CensusContext {
   explicit CensusContext(absl::string_view name)
       : span_(trace::Span::StartSpan(name)) {}
 
-  CensusContext(absl::string_view name, const trace::SpanContext &parent_ctx)
-      : span_(trace::Span::StartSpanWithRemoteParent(name, parent_ctx)) {}
+  CensusContext(absl::string_view name, const trace::Span *parent)
+      : span_(trace::Span::StartSpan(name, parent)) {}
+
+  CensusContext(absl::string_view name, const trace::SpanContext &parent_ctxt)
+      : span_(trace::Span::StartSpanWithRemoteParent(name, parent_ctxt)) {}
 
   // Serializes the outgoing trace context. Field IDs are 1 byte followed by
   // field data. A 1 byte version ID is always encoded first.
@@ -52,11 +55,13 @@ class CensusContext {
   // Serializes the outgoing stats context.  Field IDs are 1 byte followed by
   // field data. A 1 byte version ID is always encoded first.
   size_t StatsContextSerialize(char *stats_buf, size_t stats_buf_size) {
-    // TODO: Add implementation.
+    // TODO: Add implementation. Waiting on stats tagging to be added.
     return 0;
   }
 
   trace::SpanContext Context() const { return span_.context(); }
+  trace::Span Span() const { return span_; }
+  void EndSpan() { span_.End(); }
 
  private:
   trace::Span span_;
@@ -81,18 +86,20 @@ void GenerateServerContext(absl::string_view tracing, absl::string_view stats,
 // If the current context is the default context then the newly created
 // span automatically becomes a root span. This should only be called with a
 // blank CensusContext as it overwrites it.
-void GenerateClientContext(absl::string_view method, CensusContext *context);
+void GenerateClientContext(absl::string_view method, CensusContext *ctxt,
+                           CensusContext *parent_ctx);
 
 // Returns the incoming data size from the grpc call final info.
 uint64_t GetIncomingDataSize(const grpc_call_final_info *final_info);
 
 // Returns the outgoing data size from the grpc call final info.
 uint64_t GetOutgoingDataSize(const grpc_call_final_info *final_info);
 
-// This is a helper function which will return the SpanContext associated with
-// the census_context* stored by grpc. The user will need to call this for
-// manual propagation of tracing data.
+// These helper functions return the SpanContext and Span, respectively
+// associated with the census_context* stored by grpc. The user will need to
+// call this for manual propagation of tracing data.
 trace::SpanContext SpanContextFromCensusContext(const census_context *ctxt);
+trace::Span SpanFromCensusContext(const census_context *ctxt);
 
 // Returns a string representation of the StatusCode enum.
 absl::string_view StatusCodeToString(grpc_status_code code);
diff --git a/opencensus/plugins/grpc/internal/server_filter.cc b/opencensus/plugins/grpc/internal/server_filter.cc
@@ -91,13 +91,14 @@ void CensusServerCallData::OnDoneRecvInitialMetadataCb(void *user_data,
     sml.tracing_slice = grpc_empty_slice();
     sml.census_proto = grpc_empty_slice();
     FilterInitialMetadata(initial_metadata, &sml);
-    calld->method_ = grpc_slice_ref_internal(sml.path);
-    const char *method_str = GRPC_SLICE_IS_EMPTY(calld->method_)
+    calld->path_ = grpc_slice_ref_internal(sml.path);
+    const char *method_str = GRPC_SLICE_IS_EMPTY(calld->path_)
                                  ? ""
                                  : reinterpret_cast<const char *>(
-                                       GRPC_SLICE_START_PTR(calld->method_));
-    calld->method_size_ =
-        GRPC_SLICE_IS_EMPTY(sml.path) ? 0 : GRPC_SLICE_LENGTH(sml.path);
+                                       GRPC_SLICE_START_PTR(calld->path_));
+    calld->method_ = absl::string_view(
+        method_str,
+        GRPC_SLICE_IS_EMPTY(sml.path) ? 0 : GRPC_SLICE_LENGTH(sml.path));
     const char *tracing_str =
         GRPC_SLICE_IS_EMPTY(sml.tracing_slice)
             ? ""
@@ -114,14 +115,12 @@ void CensusServerCallData::OnDoneRecvInitialMetadataCb(void *user_data,
                                 ? 0
                                 : GRPC_SLICE_LENGTH(sml.census_proto);
 
-    // TODO: Setup GenerateServerContext().
-    GenerateServerContext(
-        absl::string_view(tracing_str, tracing_str_len),
-        absl::string_view(census_str, census_str_len), absl::string_view("", 0),
-        absl::string_view(method_str, calld->method_size_), &calld->context_);
-
-    absl::string_view method(method_str, calld->method_size_);
-    stats::Record({{RpcServerStartedCount(), 1}}, {{kMethodTagKey, method}});
+    GenerateServerContext(absl::string_view(tracing_str, tracing_str_len),
+                          absl::string_view(census_str, census_str_len),
+                          /*primary_role*/ "", calld->method_,
+                          &calld->context_);
+    stats::Record({{RpcServerStartedCount(), 1}},
+                  {{kMethodTagKey, calld->method_}});
 
     grpc_slice_unref_internal(sml.tracing_slice);
     grpc_slice_unref_internal(sml.census_proto);
@@ -193,7 +192,6 @@ void CensusServerCallData::Destroy(grpc_call_element *elem,
   const uint64_t response_size = GetIncomingDataSize(final_info);
   double elapsed_time_ms = absl::ToDoubleMilliseconds(elapsed_time_);
   grpc_auth_context_release(auth_context_);
-  grpc_slice_unref_internal(method_);
   stats::Record(
       {{RpcServerErrorCount(),
         final_info->final_status == GRPC_STATUS_OK ? 0 : 1},
@@ -203,10 +201,10 @@ void CensusServerCallData::Destroy(grpc_call_element *elem,
        {RpcServerRequestCount(), sent_message_count_},
        {RpcServerFinishedCount(), 1},
        {RpcServerResponseCount(), recv_message_count_}},
-      {{kMethodTagKey, absl::string_view(reinterpret_cast<char *>(
-                                             GRPC_SLICE_START_PTR(method_)),
-                                         method_size_)},
+      {{kMethodTagKey, method_},
        {kStatusTagKey, StatusCodeToString(final_info->final_status)}});
+  grpc_slice_unref_internal(path_);
+  context_.EndSpan();
 }
 
 }  // namespace opencensus
diff --git a/opencensus/plugins/grpc/internal/server_filter.h b/opencensus/plugins/grpc/internal/server_filter.h
@@ -31,8 +31,7 @@ namespace opencensus {
 class CensusServerCallData : public grpc::CallData {
  public:
   CensusServerCallData()
-      : method_size_(0),
-        gc_(nullptr),
+      : gc_(nullptr),
         auth_context_(nullptr),
         recv_initial_metadata_(nullptr),
         initial_on_done_recv_initial_metadata_(nullptr),
@@ -41,7 +40,7 @@ class CensusServerCallData : public grpc::CallData {
         recv_message_count_(0),
         sent_message_count_(0) {
     memset(&census_bin_, 0, sizeof(grpc_linked_mdelem));
-    memset(&method_, 0, sizeof(grpc_slice));
+    memset(&path_, 0, sizeof(grpc_slice));
     memset(&on_done_recv_initial_metadata_, 0, sizeof(grpc_closure));
     memset(&on_done_recv_message_, 0, sizeof(grpc_closure));
   }
@@ -62,9 +61,8 @@ class CensusServerCallData : public grpc::CallData {
  private:
   CensusContext context_;
   // server method
-  grpc_slice method_;
-  // server method size
-  size_t method_size_;
+  absl::string_view method_;
+  grpc_slice path_;
   // Pointer to the grpc_call element
   grpc_call *gc_;
   // Authorization context for the call.
