From 822dab21d9995c5cf942476b35ca12a1aa9d2737 Mon Sep 17 00:00:00 2001
From: Craig Tiller <ctiller@google.com>
Date: Wed, 8 Mar 2023 16:15:48 -0800
Subject: [PATCH] [promises] Support marking calls as traced (#32355)

<!--

If you know who should review your pull request, please assign it to
that
person, otherwise the pull request would get assigned randomly.

If your pull request is for a specific language, please add the
appropriate
lang label.

-->
---
 src/core/lib/channel/connected_channel.cc    | 1 +
 src/core/lib/channel/promise_based_filter.cc | 6 ++++++
 src/core/lib/channel/promise_based_filter.h  | 9 +++++++--
 src/core/lib/surface/call.h                  | 5 +++++
 4 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/src/core/lib/channel/connected_channel.cc b/src/core/lib/channel/connected_channel.cc
index b45501d0d51..40582f99e44 100644
--- a/src/core/lib/channel/connected_channel.cc
+++ b/src/core/lib/channel/connected_channel.cc
@@ -327,6 +327,7 @@ class ConnectedChannelStream : public Orphanable {
 
   void SchedulePush(grpc_transport_stream_op_batch* batch)
       ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+    batch->is_traced = GetContext<CallContext>()->traced();
     if (grpc_call_trace.enabled()) {
       gpr_log(GPR_DEBUG, "%s[connected] Push batch to transport: %s",
               Activity::current()->DebugTag().c_str(),
diff --git a/src/core/lib/channel/promise_based_filter.cc b/src/core/lib/channel/promise_based_filter.cc
index 1a713623d31..d44a9920216 100644
--- a/src/core/lib/channel/promise_based_filter.cc
+++ b/src/core/lib/channel/promise_based_filter.cc
@@ -91,6 +91,9 @@ BaseCallData::BaseCallData(
       arena_(args->arena),
       call_combiner_(args->call_combiner),
       deadline_(args->deadline),
+      call_context_(flags & kFilterExaminesCallContext
+                        ? arena_->New<CallContext>(nullptr)
+                        : nullptr),
       context_(args->context),
       server_initial_metadata_pipe_(
           flags & kFilterExaminesServerInitialMetadata
@@ -281,6 +284,9 @@ BaseCallData::Flusher::~Flusher() {
   };
   for (size_t i = 1; i < release_.size(); i++) {
     auto* batch = release_[i];
+    if (call_->call_context_ != nullptr && call_->call_context_->traced()) {
+      batch->is_traced = true;
+    }
     if (grpc_trace_channel.enabled()) {
       gpr_log(
           GPR_INFO, "FLUSHER:queue batch to forward in closure: %s",
diff --git a/src/core/lib/channel/promise_based_filter.h b/src/core/lib/channel/promise_based_filter.h
index 5597349337e..e5ec0b2aed8 100644
--- a/src/core/lib/channel/promise_based_filter.h
+++ b/src/core/lib/channel/promise_based_filter.h
@@ -62,6 +62,7 @@
 #include "src/core/lib/promise/poll.h"
 #include "src/core/lib/resource_quota/arena.h"
 #include "src/core/lib/slice/slice_buffer.h"
+#include "src/core/lib/surface/call.h"
 #include "src/core/lib/transport/error_utils.h"
 #include "src/core/lib/transport/metadata_batch.h"
 #include "src/core/lib/transport/transport.h"
@@ -134,6 +135,7 @@ static constexpr uint8_t kFilterExaminesServerInitialMetadata = 1;
 static constexpr uint8_t kFilterIsLast = 2;
 static constexpr uint8_t kFilterExaminesOutboundMessages = 4;
 static constexpr uint8_t kFilterExaminesInboundMessages = 8;
+static constexpr uint8_t kFilterExaminesCallContext = 16;
 
 namespace promise_filter_detail {
 
@@ -197,7 +199,8 @@ class BaseCallData : public Activity, private Wakeable {
         public promise_detail::Context<grpc_polling_entity>,
         public promise_detail::Context<CallFinalization>,
         public promise_detail::Context<
-            grpc_event_engine::experimental::EventEngine> {
+            grpc_event_engine::experimental::EventEngine>,
+        public promise_detail::Context<CallContext> {
    public:
     explicit ScopedContext(BaseCallData* call_data)
         : promise_detail::Context<Arena>(call_data->arena_),
@@ -207,7 +210,8 @@ class BaseCallData : public Activity, private Wakeable {
               call_data->pollent_.load(std::memory_order_acquire)),
           promise_detail::Context<CallFinalization>(&call_data->finalization_),
           promise_detail::Context<grpc_event_engine::experimental::EventEngine>(
-              call_data->event_engine_) {}
+              call_data->event_engine_),
+          promise_detail::Context<CallContext>(call_data->call_context_) {}
   };
 
   class Flusher {
@@ -549,6 +553,7 @@ class BaseCallData : public Activity, private Wakeable {
   CallCombiner* const call_combiner_;
   const Timestamp deadline_;
   CallFinalization finalization_;
+  CallContext* call_context_ = nullptr;
   grpc_call_context_element* const context_;
   std::atomic<grpc_polling_entity*> pollent_{nullptr};
   Pipe<ServerMetadataHandle>* const server_initial_metadata_pipe_;
diff --git a/src/core/lib/surface/call.h b/src/core/lib/surface/call.h
index 0a77709946d..ea4d3937157 100644
--- a/src/core/lib/surface/call.h
+++ b/src/core/lib/surface/call.h
@@ -125,6 +125,9 @@ class CallContext {
 
   ServerCallContext* server_call_context();
 
+  void set_traced(bool traced) { traced_ = traced; }
+  bool traced() const { return traced_; }
+
  private:
   friend class PromiseBasedCall;
   // Call final info.
@@ -135,6 +138,8 @@ class CallContext {
   // TODO(ctiller): remove this once transport APIs are promise based and we
   // don't need refcounting here.
   PromiseBasedCall* const call_;
+  // Is this call traced?
+  bool traced_ = false;
 };
 
 template <>