diff --git a/include/grpc++/client_context.h b/include/grpc++/client_context.h index 5e10875260474ba13e33a7efa07a469ad64a61cf..4d96d862e7d1a41258f7b93d2b5c910c7df8f8a5 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 b2049b3289b8fd65f777ce7a87180accc135a92e..3fc07affc8460cb84675fdbc4a76fbba9e29f904 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 cf2353199f3dab2f124760b354c38fd1f8ae1634..0ed63469b649f78cd41722530af8cc2bdca3381f 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 f610f6ce21d0220e058b782685fa16e54ddda9a9..4637e7218e387932c733fa4bf2b5fe4b24eedff6 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 057ac78ee7248a3475feac88cff43269c6c25dcc..80165206417f2c7d6d86bba8381b1d8f546c55d3 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 fc09137b670c0644057b2946af8d950c67ec083a..a28a542c8dd5cdf339b9056e3cf9496f8baf5bc0 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 72593f877e172d0170e6e2cbc99ee95683e2f0e2..5bc6f6fd91344193471b30f20e01d60bd9a2ef8d 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);