[context] Move tracing context to arena based context (#36782)

Closes #36782

COPYBARA_INTEGRATE_REVIEW=https://github.com/grpc/grpc/pull/36782 from ctiller:ctx8 ea7da93756
PiperOrigin-RevId: 639162692

BUILD

@@ -1341,7 +1341,7 @@ grpc_cc_library(
         "grpc_base",
         "grpc_public_hdrs",
         "grpc_trace",
-        "legacy_context",
+        "//src/core:arena",
     ],
 )
 

src/core/ext/filters/census/grpc_context.cc

@@ -21,7 +21,6 @@
 #include <grpc/census.h>
 #include <grpc/grpc.h>
 
-#include "src/core/lib/channel/context.h"
 #include "src/core/lib/debug/trace.h"
 #include "src/core/lib/surface/api_trace.h"
 #include "src/core/lib/surface/call.h"
@@ -30,12 +29,11 @@ void grpc_census_call_set_context(grpc_call* call, census_context* context) {
   GRPC_API_TRACE("grpc_census_call_set_context(call=%p, census_context=%p)", 2,
                  (call, context));
   if (context != nullptr) {
-    grpc_call_context_set(call, GRPC_CONTEXT_TRACING, context, nullptr);
+    grpc_call_get_arena(call)->SetContext<census_context>(context);
   }
 }
 
 census_context* grpc_census_call_get_context(grpc_call* call) {
   GRPC_API_TRACE("grpc_census_call_get_context(call=%p)", 1, (call));
-  return static_cast<census_context*>(
-      grpc_call_context_get(call, GRPC_CONTEXT_TRACING));
+  return grpc_call_get_arena(call)->GetContext<census_context>();
 }

src/core/ext/transport/chttp2/transport/chttp2_transport.h

@@ -83,7 +83,7 @@ void TestOnlyGlobalHttp2TransportDisableTransientFailureStateNotification(
 
 typedef void (*WriteTimestampsCallback)(void*, Timestamps*,
                                         grpc_error_handle error);
-typedef void* (*CopyContextFn)(void*);
+typedef void* (*CopyContextFn)(grpc_core::Arena*);
 
 void GrpcHttp2SetWriteTimestampsCallback(WriteTimestampsCallback fn);
 void GrpcHttp2SetCopyContextFn(CopyContextFn fn);

src/core/ext/transport/chttp2/transport/writing.cc

@@ -682,7 +682,7 @@ grpc_chttp2_begin_write_result grpc_chttp2_begin_write(
             grpc_core::GrpcHttp2GetCopyContextFn();
         if (copy_context_fn != nullptr &&
             grpc_core::GrpcHttp2GetWriteTimestampsCallback() != nullptr) {
-          t->context_list->emplace_back(copy_context_fn(s->context),
+          t->context_list->emplace_back(copy_context_fn(s->arena),
                                         outbuf_relative_start_pos,
                                         num_stream_bytes, s->byte_counter,
                                         s->write_counter - 1, s->tcp_tracer);

src/core/lib/channel/context.h

@@ -32,9 +32,6 @@ typedef enum {
   /// grpc_call* associated with this context.
   GRPC_CONTEXT_CALL = 0,
 
-  /// Value is a \a census_context.
-  GRPC_CONTEXT_TRACING,
-
   /// Holds a pointer to ServiceConfigCallData associated with this call.
   GRPC_CONTEXT_SERVICE_CONFIG_CALL_DATA,
 

src/core/lib/surface/call.cc

@@ -161,8 +161,8 @@ absl::Status Call::InitParent(Call* parent, uint32_t propagation_mask) {
           "Census tracing propagation requested without Census context "
           "propagation");
     }
-    ContextSet(GRPC_CONTEXT_TRACING, parent->ContextGet(GRPC_CONTEXT_TRACING),
-               nullptr);
+    arena()->SetContext<census_context>(
+        parent->arena()->GetContext<census_context>());
   } else if (propagation_mask & GRPC_PROPAGATE_CENSUS_STATS_CONTEXT) {
     return absl::UnknownError(
         "Census context propagation requested without Census tracing "

src/core/lib/surface/call.h

@@ -78,6 +78,11 @@ typedef struct grpc_call_create_args {
 
 namespace grpc_core {
 
+template <>
+struct ArenaContextType<census_context> {
+  static void Destroy(census_context*) {}
+};
+
 class Call : public CppImplOf<Call, grpc_call>,
              public grpc_event_engine::experimental::EventEngine::
                  Closure /* for deadlines */ {

src/cpp/ext/filters/census/client_filter.cc

@@ -312,8 +312,8 @@ OpenCensusCallTracer::OpenCensusCallTracer(
       method_(GetMethod(path_)),
       arena_(arena),
       tracing_enabled_(tracing_enabled) {
-  auto* parent_context = reinterpret_cast<CensusContext*>(
-      call_context_[GRPC_CONTEXT_TRACING].value);
+  auto* parent_context =
+      reinterpret_cast<CensusContext*>(arena->GetContext<census_context>());
   GenerateClientContext(tracing_enabled_ ? absl::StrCat("Sent.", method_) : "",
                         &context_,
                         (parent_context == nullptr) ? nullptr : parent_context);

src/cpp/ext/filters/census/server_call_tracer.cc

@@ -40,16 +40,16 @@
 #include "opencensus/trace/span_id.h"
 #include "opencensus/trace/trace_id.h"
 
+#include <grpc/grpc.h>
 #include <grpc/support/port_platform.h>
 #include <grpcpp/opencensus.h>
 
-#include "src/core/lib/channel/channel_stack.h"
-#include "src/core/lib/channel/context.h"
 #include "src/core/lib/iomgr/error.h"
 #include "src/core/lib/promise/context.h"
 #include "src/core/lib/resource_quota/arena.h"
 #include "src/core/lib/slice/slice.h"
 #include "src/core/lib/slice/slice_buffer.h"
+#include "src/core/lib/surface/call.h"
 #include "src/core/lib/transport/metadata_batch.h"
 #include "src/core/telemetry/call_tracer.h"
 #include "src/core/telemetry/tcp_tracer.h"
@@ -206,8 +206,8 @@ void OpenCensusServerCallTracer::RecordReceivedInitialMetadata(
       tracing_enabled ? sml.tracing_slice.as_string_view() : "",
       absl::StrCat("Recv.", method_), &context_);
   if (tracing_enabled) {
-    auto* call_context = grpc_core::GetContext<grpc_call_context_element>();
-    call_context[GRPC_CONTEXT_TRACING].value = &context_;
+    grpc_core::SetContext<census_context>(
+        reinterpret_cast<census_context*>(&context_));
   }
   if (OpenCensusStatsEnabled()) {
     std::vector<std::pair<opencensus::tags::TagKey, std::string>> tags =