7 files changed, 53 insertions, 10 deletions
diff --git a/include/grpc++/client_context.h b/include/grpc++/client_context.h
index 5e10875260..4d96d862e7 100644
--- a/include/grpc++/client_context.h
+++ b/include/grpc++/client_context.h
@@ -46,6 +46,7 @@
 
 struct grpc_call;
 struct grpc_completion_queue;
+struct census_context;
 
 namespace grpc {
 
@@ -107,6 +108,10 @@ class ClientContext {
     creds_ = creds;
   }
 
+  // Get and set census context
+  void set_census_context(census_context* ccp) { census_context_ = ccp; }
+  census_context* get_census_context() const { return census_context_; }
+
   void TryCancel();
 
  private:
@@ -154,6 +159,7 @@ class ClientContext {
   gpr_timespec deadline_;
   grpc::string authority_;
   std::shared_ptr<Credentials> creds_;
+  census_context* census_context_;
   std::multimap<grpc::string, grpc::string> send_initial_metadata_;
   std::multimap<grpc::string, grpc::string> recv_initial_metadata_;
   std::multimap<grpc::string, grpc::string> trailing_metadata_;
diff --git a/include/grpc/census.h b/include/grpc/census.h
index b2049b3289..3fc07affc8 100644
--- a/include/grpc/census.h
+++ b/include/grpc/census.h
@@ -61,6 +61,10 @@ enum census_functions {
 int census_initialize(int functions);
 void census_shutdown();
 
+/* If any census feature has been initialized, this funtion will return a
+ * non-zero value. */
+int census_available();
+
 /* Internally, Census relies on a context, which should be propagated across
  * RPC's. From the RPC subsystems viewpoint, this is an opaque data structure.
  * A context must be used as the first argument to all other census
diff --git a/src/core/census/grpc_context.c b/src/core/census/grpc_context.c
index cf2353199f..0ed63469b6 100644
--- a/src/core/census/grpc_context.c
+++ b/src/core/census/grpc_context.c
@@ -34,12 +34,28 @@
 #include <grpc/census.h>
 #include "src/core/census/grpc_context.h"
 
-void *grpc_census_context_create() {
-  census_context *context;
-  census_context_deserialize(NULL, &context);
-  return (void *)context;
+static void grpc_census_context_destroy(void *context) {
+  census_context_destroy((census_context *)context);
 }
 
-void grpc_census_context_destroy(void *context) {
-  census_context_destroy((census_context *)context);
+void grpc_census_call_set_context(grpc_call *call, census_context *context) {
+  if (!census_available()) {
+    return;
+  }
+  if (context == NULL) {
+    if (grpc_call_is_client(call)) {
+      census_context *context_ptr;
+      census_context_deserialize(NULL, &context_ptr);
+      grpc_call_context_set(call, GRPC_CONTEXT_TRACING, context_ptr,
+                            grpc_census_context_destroy);
+    } else {
+      /* TODO(aveitch): server side context code to be implemented. */
+    }
+  } else {
+    grpc_call_context_set(call, GRPC_CONTEXT_TRACING, context, NULL);
+  }
+}
+
+census_context *grpc_census_call_get_context(grpc_call *call) {
+  return (census_context *)grpc_call_context_get(call, GRPC_CONTEXT_TRACING);
 }
diff --git a/src/core/census/grpc_context.h b/src/core/census/grpc_context.h
index f610f6ce21..4637e7218e 100644
--- a/src/core/census/grpc_context.h
+++ b/src/core/census/grpc_context.h
@@ -36,7 +36,22 @@
 #ifndef CENSUS_GRPC_CONTEXT_H
 #define CENSUS_GRPC_CONTEXT_H
 
-void *grpc_census_context_create();
-void grpc_census_context_destroy(void *context);
+#include <grpc/census.h>
+#include "src/core/surface/call.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/* Set census context for the call; Must be called before first call to
+   grpc_call_start_batch(). */
+void grpc_census_call_set_context(grpc_call *call, census_context *context);
+
+/* Retrieve the calls current census context. */
+census_context *grpc_census_call_get_context(grpc_call *call);
+
+#ifdef __cplusplus
+}
+#endif
 
 #endif /* CENSUS_GRPC_CONTEXT_H */
diff --git a/src/core/census/initialize.c b/src/core/census/initialize.c
index 057ac78ee7..8016520641 100644
--- a/src/core/census/initialize.c
+++ b/src/core/census/initialize.c
@@ -48,3 +48,5 @@ int census_initialize(int functions) {
 }
 
 void census_shutdown() { census_fns_enabled = CENSUS_NONE; }
+
+int census_available() { return (census_fns_enabled != CENSUS_NONE); }
diff --git a/src/core/surface/call.c b/src/core/surface/call.c
index fc09137b67..a28a542c8d 100644
--- a/src/core/surface/call.c
+++ b/src/core/surface/call.c
@@ -298,8 +298,6 @@ grpc_call *grpc_call_create(grpc_channel *channel, grpc_completion_queue *cq,
   if (call->is_client) {
     call->request_set[GRPC_IOREQ_SEND_TRAILING_METADATA] = REQSET_DONE;
     call->request_set[GRPC_IOREQ_SEND_STATUS] = REQSET_DONE;
-    call->context[GRPC_CONTEXT_TRACING].value = grpc_census_context_create();
-    call->context[GRPC_CONTEXT_TRACING].destroy = grpc_census_context_destroy;
   }
   GPR_ASSERT(add_initial_metadata_count < MAX_SEND_INITIAL_METADATA_COUNT);
   for (i = 0; i < add_initial_metadata_count; i++) {
diff --git a/src/cpp/client/channel.cc b/src/cpp/client/channel.cc
index 72593f877e..5bc6f6fd91 100644
--- a/src/cpp/client/channel.cc
+++ b/src/cpp/client/channel.cc
@@ -39,6 +39,7 @@
 #include <grpc/support/log.h>
 #include <grpc/support/slice.h>
 
+#include "src/core/census/grpc_context.h"
 #include "src/core/profiling/timers.h"
 #include <grpc++/channel_arguments.h>
 #include <grpc++/client_context.h>
@@ -68,6 +69,7 @@ Call Channel::CreateCall(const RpcMethod& method, ClientContext* context,
                                          ? target_.c_str()
                                          : context->authority().c_str(),
                                      context->raw_deadline());
+  grpc_census_call_set_context(c_call, context->get_census_context());
   GRPC_TIMER_MARK(GRPC_PTAG_CPP_CALL_CREATED, c_call);
   context->set_call(c_call, shared_from_this());
   return Call(c_call, this, cq);