3 files changed, 41 insertions, 8 deletions
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); }