impl(pubsub): inject trace context into message (#12868)

Author: alevenberg

google/cloud/pubsub/internal/message_propagator.cc

@@ -27,10 +27,12 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 
 void InjectTraceContext(
     pubsub::Message& message,
-    opentelemetry::context::propagation::TextMapPropagator& propagator) {
+    std::shared_ptr<
+        opentelemetry::context::propagation::TextMapPropagator> const&
+        propagator) {
   auto current = opentelemetry::context::RuntimeContext::GetCurrent();
   MessageCarrier carrier(message);
-  propagator.Inject(carrier, current);
+  propagator->Inject(carrier, current);
 }
 
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END

google/cloud/pubsub/internal/message_propagator.h

@@ -32,7 +32,9 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
  */
 void InjectTraceContext(
     pubsub::Message& message,
-    opentelemetry::context::propagation::TextMapPropagator& propagator);
+    std::shared_ptr<
+        opentelemetry::context::propagation::TextMapPropagator> const&
+        propagator);
 
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace pubsub_internal

google/cloud/pubsub/internal/message_propagator_test.cc

@@ -23,6 +23,7 @@
 #include "opentelemetry/context/propagation/text_map_propagator.h"
 #include "opentelemetry/trace/scope.h"
 #include <gmock/gmock.h>
+#include <opentelemetry/trace/propagation/http_trace_context.h>
 
 namespace google {
 namespace cloud {
@@ -32,6 +33,8 @@ namespace {
 
 using ::google::cloud::testing_util::InstallSpanCatcher;
 using ::testing::_;
+using ::testing::Contains;
+using ::testing::Pair;
 using ::testing::StartsWith;
 
 opentelemetry::nostd::shared_ptr<opentelemetry::trace::Span> MakeTestSpan() {
@@ -43,9 +46,10 @@ TEST(MessagePropagatorTest, InjectTraceContext) {
   auto span_catcher = InstallSpanCatcher();
   opentelemetry::trace::Scope scope(MakeTestSpan());
   auto message = pubsub::MessageBuilder().Build();
-  auto propagator = internal::MakePropagator();
+  auto propagator =
+      std::make_shared<opentelemetry::trace::propagation::HttpTraceContext>();
 
-  InjectTraceContext(message, *propagator);
+  InjectTraceContext(message, propagator);
 
   EXPECT_THAT(message.attributes(),
               Contains(Pair(StartsWith("googclient_"), _)));

google/cloud/pubsub/internal/publisher_tracing_connection.cc

@@ -22,9 +22,12 @@
 #include <string>
 #include <utility>
 #ifdef GOOGLE_CLOUD_CPP_HAVE_OPENTELEMETRY
+#include "google/cloud/pubsub/internal/message_propagator.h"
 #include "google/cloud/internal/opentelemetry.h"
 #include "google/cloud/status_or.h"
+#include <opentelemetry/context/propagation/text_map_propagator.h>
 #include <opentelemetry/nostd/shared_ptr.h>
+#include <opentelemetry/trace/propagation/http_trace_context.h>
 #include <opentelemetry/trace/scope.h>
 #include <opentelemetry/trace/semantic_conventions.h>
 #include <opentelemetry/trace/span_metadata.h>
@@ -78,13 +81,19 @@ class PublisherTracingConnection : public pubsub::PublisherConnection {
  public:
   explicit PublisherTracingConnection(
       pubsub::Topic topic, std::shared_ptr<pubsub::PublisherConnection> child)
-      : topic_(std::move(topic)), child_(std::move(child)) {}
+      : topic_(std::move(topic)),
+        child_(std::move(child)),
+        propagator_(std::make_shared<
+                    opentelemetry::trace::propagation::HttpTraceContext>()) {}
 
   ~PublisherTracingConnection() override = default;
 
   future<StatusOr<std::string>> Publish(PublishParams p) override {
     auto span = StartPublishSpan(topic_.FullName(), p.message);
     auto scope = opentelemetry::trace::Scope(span);
+
+    InjectTraceContext(p.message, propagator_);
+
     return EndPublishSpan(std::move(span), child_->Publish(std::move(p)));
   };
 
@@ -105,6 +114,8 @@ class PublisherTracingConnection : public pubsub::PublisherConnection {
  private:
   pubsub::Topic const topic_;
   std::shared_ptr<pubsub::PublisherConnection> child_;
+  std::shared_ptr<opentelemetry::context::propagation::TextMapPropagator>
+      propagator_;
 };
 
 }  // namespace

google/cloud/pubsub/internal/publisher_tracing_connection_test.cc

@@ -51,8 +51,11 @@ using ::google::cloud::testing_util::StatusIs;
 using ::google::cloud::testing_util::ThereIsAnActiveSpan;
 using ::testing::_;
 using ::testing::AllOf;
+using ::testing::Contains;
 using ::testing::ElementsAre;
 using ::testing::Not;
+using ::testing::Pair;
+using ::testing::StartsWith;
 
 TEST(PublisherTracingConnectionTest, PublishSpanOnSuccess) {
   namespace sc = ::opentelemetry::trace::SemanticConventions;
@@ -139,6 +142,29 @@ TEST(PublisherTracingConnectionTest, PublishSpanOnError) {
                         "messaging.message.total_size_bytes", 45)))));
 }
 
+TEST(PublisherTracingConnectionTest, PublishInjectsTraceContext) {
+  auto mock = std::make_shared<MockPublisherConnection>();
+  // Need to install the span catcher so ThereIsAnActiveSpan() detects a span.
+  auto span_catcher = InstallSpanCatcher();
+  EXPECT_CALL(*mock, Publish)
+      .WillOnce([&](pubsub::PublisherConnection::PublishParams const& p) {
+        EXPECT_TRUE(ThereIsAnActiveSpan());
+        // We need to test the trace context has been injected here, since the
+        // connection moves the message to the child connection.
+        EXPECT_THAT(p.message.attributes(),
+                    Contains(Pair(StartsWith("googclient_"), _)));
+        return make_ready_future(StatusOr<std::string>("test-id-0"));
+      });
+  auto connection = MakePublisherTracingConnection(
+      Topic("test-project", "test-topic"), std::move(mock));
+
+  auto message = pubsub::MessageBuilder{}
+                     .SetData("test-data-0")
+                     .SetOrderingKey("ordering-key-0")
+                     .Build();
+  auto response = connection->Publish({message}).get();
+}
+
 TEST(PublisherTracingConnectionTest, PublishSpanOmitsOrderingKey) {
   auto span_catcher = InstallSpanCatcher();
   auto mock = std::make_shared<MockPublisherConnection>();