OpenCensusCallTracer: Move context generation to StartTransportStreamOpBatch (#27523)

* OpenCensusCallTracer: Move context generation to StartTransportStreamOpBatch

* Reviewer comments

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

@@ -44,15 +44,28 @@ constexpr uint32_t
 
 grpc_error_handle CensusClientCallData::Init(
     grpc_call_element* /* elem */, const grpc_call_element_args* args) {
-  auto tracer = args->arena->New<OpenCensusCallTracer>(args);
+  tracer_ = args->arena->New<OpenCensusCallTracer>(args);
   GPR_DEBUG_ASSERT(args->context[GRPC_CONTEXT_CALL_TRACER].value == nullptr);
-  args->context[GRPC_CONTEXT_CALL_TRACER].value = tracer;
+  args->context[GRPC_CONTEXT_CALL_TRACER].value = tracer_;
   args->context[GRPC_CONTEXT_CALL_TRACER].destroy = [](void* tracer) {
     (static_cast<OpenCensusCallTracer*>(tracer))->~OpenCensusCallTracer();
   };
   return GRPC_ERROR_NONE;
 }
 
+void CensusClientCallData::StartTransportStreamOpBatch(
+    grpc_call_element* elem, TransportStreamOpBatch* op) {
+  // Note that we are generating the overall call context here instead of in
+  // the constructor of `OpenCensusCallTracer` due to the semantics of
+  // `grpc_census_call_set_context` which allows the application to set the
+  // census context for a call anytime before the first call to
+  // `grpc_call_start_batch`.
+  if (op->op()->send_initial_metadata) {
+    tracer_->GenerateContext();
+  }
+  grpc_call_next_op(elem, op->op());
+}
+
 //
 // OpenCensusCallTracer::OpenCensusCallAttemptTracer
 //
@@ -213,6 +226,13 @@ OpenCensusCallTracer::~OpenCensusCallTracer() {
   grpc_slice_unref_internal(path_);
 }
 
+void OpenCensusCallTracer::GenerateContext() {
+  auto* parent_context = reinterpret_cast<CensusContext*>(
+      call_context_[GRPC_CONTEXT_TRACING].value);
+  GenerateClientContext(absl::StrCat("Sent.", method_), &context_,
+                        (parent_context == nullptr) ? nullptr : parent_context);
+}
+
 OpenCensusCallTracer::OpenCensusCallAttemptTracer*
 OpenCensusCallTracer::StartNewAttempt(bool is_transparent_retry) {
   // We allocate the first attempt on the arena and all subsequent attempts on
@@ -237,16 +257,6 @@ OpenCensusCallTracer::StartNewAttempt(bool is_transparent_retry) {
     ++num_active_rpcs_;
   }
   if (is_first_attempt) {
-    // Note that we are generating the overall call context here instead of in
-    // the constructor of `OpenCensusCallTracer` due to the semantics of
-    // `grpc_census_call_set_context` which allows the application to set the
-    // census context for a call anytime before the first call to
-    // `grpc_call_start_batch`.
-    auto* parent_context = reinterpret_cast<CensusContext*>(
-        call_context_[GRPC_CONTEXT_TRACING].value);
-    GenerateClientContext(
-        absl::StrCat("Sent.", method_), &context_,
-        (parent_context == nullptr) ? nullptr : parent_context);
     return arena_->New<OpenCensusCallAttemptTracer>(
         this, attempt_num, is_transparent_retry, true /* arena_allocated */);
   }

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

@@ -37,6 +37,11 @@ class CensusClientCallData : public CallData {
  public:
   grpc_error_handle Init(grpc_call_element* /* elem */,
                          const grpc_call_element_args* args) override;
+  void StartTransportStreamOpBatch(grpc_call_element* elem,
+                                   TransportStreamOpBatch* op) override;
+
+ private:
+  OpenCensusCallTracer* tracer_ = nullptr;
 };
 
 }  // namespace grpc

src/cpp/ext/filters/census/open_census_call_tracer.h

@@ -83,6 +83,7 @@ class OpenCensusCallTracer : public grpc_core::CallTracer {
   explicit OpenCensusCallTracer(const grpc_call_element_args* args);
   ~OpenCensusCallTracer() override;
 
+  void GenerateContext();
   OpenCensusCallAttemptTracer* StartNewAttempt(
       bool is_transparent_retry) override;