Fix python observability

ctiller · May 31, 2024 · 904b036 · 904b036
1 parent ef2439b
commit 904b036
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 9 deletions.
diff --git a/src/core/lib/surface/call.cc b/src/core/lib/surface/call.cc
@@ -3760,6 +3760,19 @@ void* grpc_call_context_get(grpc_call* call, grpc_context_index elem) {
   return grpc_core::Call::FromC(call)->ContextGet(elem);
 }
 
+void grpc_call_tracer_set(grpc_call* call, grpc_core::ClientCallTracer* tracer,
+                          void (*destroy)(void* value)) {
+  grpc_core::Arena* arena = grpc_call_get_arena(call);
+  return arena->SetContext<grpc_core::CallTracerAnnotationInterface>(tracer);
+}
+
+void* grpc_call_tracer_get(grpc_call* call) {
+  grpc_core::Arena* arena = grpc_call_get_arena(call);
+  auto* call_tracer =
+      arena->GetContext<grpc_core::CallTracerAnnotationInterface>();
+  return call_tracer;
+}
+
 uint8_t grpc_call_is_client(grpc_call* call) {
   return grpc_core::Call::FromC(call)->is_client();
 }

diff --git a/src/core/lib/surface/call.h b/src/core/lib/surface/call.h
@@ -338,6 +338,11 @@ void grpc_call_context_set(grpc_call* call, grpc_context_index elem,
 // Get a context pointer.
 void* grpc_call_context_get(grpc_call* call, grpc_context_index elem);
 
+void grpc_call_tracer_set(grpc_call* call, grpc_core::ClientCallTracer* tracer,
+                          void (*destroy)(void* value));
+
+void* grpc_call_tracer_get(grpc_call* call);
+
 #define GRPC_CALL_LOG_BATCH(sev, ops, nops)        \
   do {                                             \
     if (GRPC_TRACE_FLAG_ENABLED(grpc_api_trace)) { \

diff --git a/src/python/grpcio/grpc/_cython/_cygrpc/grpc.pxi b/src/python/grpcio/grpc/_cython/_cygrpc/grpc.pxi
@@ -63,6 +63,9 @@ cdef extern from "src/core/telemetry/call_tracer.h" namespace "grpc_core":
     cdef cppclass ClientCallTracer:
         pass
 
+    cdef cppclass CallTracerAnnotationInterface:
+        pass
+
     cdef cppclass ServerCallTracer:
         string TraceId() nogil
         string SpanId() nogil
@@ -72,14 +75,11 @@ cdef extern from "src/core/telemetry/call_tracer.h" namespace "grpc_core":
         @staticmethod
         void RegisterGlobal(ServerCallTracerFactory* factory) nogil
 
-cdef extern from "src/core/lib/channel/context.h":
-  ctypedef enum grpc_context_index:
-    GRPC_CONTEXT_CALL_TRACER_ANNOTATION_INTERFACE
-
 cdef extern from "src/core/lib/surface/call.h":
-  void grpc_call_context_set(grpc_call* call, grpc_context_index elem,
-                             void* value, void (*destroy)(void* value)) nogil
-  void *grpc_call_context_get(grpc_call* call, grpc_context_index elem) nogil
+  void grpc_call_tracer_set(grpc_call* call, void* value,
+                                void (*destroy)(void* value)) nogil
+
+  void* grpc_call_tracer_get(grpc_call* call) nogil
 
 cdef extern from "grpc/support/alloc.h":
 

diff --git a/src/python/grpcio/grpc/_cython/_cygrpc/observability.pyx.pxi b/src/python/grpcio/grpc/_cython/_cygrpc/observability.pyx.pxi
@@ -48,11 +48,11 @@ def maybe_save_server_trace_context(RequestCallEvent event) -> None:
 
 cdef void _set_call_tracer(grpc_call* call, void* capsule_ptr):
   cdef ClientCallTracer* call_tracer = <ClientCallTracer*>capsule_ptr
-  grpc_call_context_set(call, GRPC_CONTEXT_CALL_TRACER_ANNOTATION_INTERFACE, call_tracer, NULL)
+  grpc_call_tracer_set(call, call_tracer, NULL)
 
 
 cdef void* _get_call_tracer(grpc_call* call):
-  cdef void* call_tracer = grpc_call_context_get(call, GRPC_CONTEXT_CALL_TRACER_ANNOTATION_INTERFACE)
+  cdef void* call_tracer = grpc_call_tracer_get(call)
   return call_tracer