C++ opencensus filter: Fix point of creating context for overall call (#27221)

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

@@ -194,14 +194,10 @@ void OpenCensusCallTracer::OpenCensusCallAttemptTracer::RecordEnd(
 //
 
 OpenCensusCallTracer::OpenCensusCallTracer(const grpc_call_element_args* args)
-    : path_(grpc_slice_ref_internal(args->path)),
+    : call_context_(args->context),
+      path_(grpc_slice_ref_internal(args->path)),
       method_(GetMethod(&path_)),
-      arena_(args->arena) {
-  auto* parent_context = reinterpret_cast<CensusContext*>(
-      args->context[GRPC_CONTEXT_TRACING].value);
-  GenerateClientContext(absl::StrCat("Sent.", method_), &context_,
-                        (parent_context == nullptr) ? nullptr : parent_context);
-}
+      arena_(args->arena) {}
 
 OpenCensusCallTracer::~OpenCensusCallTracer() {
   std::vector<std::pair<opencensus::tags::TagKey, std::string>> tags =
@@ -239,6 +235,16 @@ 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/open_census_call_tracer.h

@@ -87,6 +87,7 @@ class OpenCensusCallTracer : public grpc_core::CallTracer {
       bool is_transparent_retry) override;
 
  private:
+  const grpc_call_context_element* call_context_;
   // Client method.
   grpc_slice path_;
   absl::string_view method_;

test/cpp/ext/filters/census/stats_plugin_end2end_test.cc

@@ -31,6 +31,7 @@
 #include "opencensus/stats/testing/test_utils.h"
 #include "opencensus/tags/tag_map.h"
 #include "opencensus/tags/with_tag_map.h"
+#include "src/cpp/ext/filters/census/context.h"
 #include "src/cpp/ext/filters/census/grpc_plugin.h"
 #include "src/proto/grpc/testing/echo.grpc.pb.h"
 #include "test/core/util/test_config.h"
@@ -49,11 +50,23 @@ using ::opencensus::tags::WithTagMap;
 
 static const auto TEST_TAG_KEY = TagKey::Register("my_key");
 static const auto TEST_TAG_VALUE = "my_value";
+const char* kExpectedTraceIdKey = "expected_trace_id";
 
 class EchoServer final : public EchoTestService::Service {
-  ::grpc::Status Echo(::grpc::ServerContext* /*context*/,
+  ::grpc::Status Echo(::grpc::ServerContext* context,
                       const EchoRequest* request,
                       EchoResponse* response) override {
+    for (const auto& metadata : context->client_metadata()) {
+      if (metadata.first == kExpectedTraceIdKey) {
+        EXPECT_EQ(metadata.second, reinterpret_cast<const grpc::CensusContext*>(
+                                       context->census_context())
+                                       ->Span()
+                                       .context()
+                                       .trace_id()
+                                       .ToHex());
+        break;
+      }
+    }
     if (request->param().expected_error().code() == 0) {
       response->set_message(request->message());
       return ::grpc::Status::OK;
@@ -515,6 +528,24 @@ TEST_F(StatsPluginEnd2EndTest, TestRetryStatsWithAdditionalRetries) {
   }
 }
 
+// Test that CensusContext object set by application is used.
+TEST_F(StatsPluginEnd2EndTest, TestApplicationCensusContextFlows) {
+  auto channel = CreateChannel(server_address_, InsecureChannelCredentials());
+  ResetStub(channel);
+  EchoRequest request;
+  request.set_message("foo");
+  EchoResponse response;
+  ::grpc::ClientContext context;
+  ::grpc::CensusContext app_census_context("root",
+                                           ::opencensus::tags::TagMap{});
+  context.set_census_context(
+      reinterpret_cast<census_context*>(&app_census_context));
+  context.AddMetadata(kExpectedTraceIdKey,
+                      app_census_context.Span().context().trace_id().ToHex());
+  ::grpc::Status status = stub_->Echo(&context, request, &response);
+  EXPECT_TRUE(status.ok());
+}
+
 }  // namespace
 }  // namespace testing
 }  // namespace grpc