From 916fe8f3033a74904e78528e08de2d36484ff704 Mon Sep 17 00:00:00 2001
From: Anna Levenberg <alevenb@google.com>
Date: Tue, 19 Dec 2023 16:35:13 -0500
Subject: [PATCH] impl(pubsub): start the publish span as a root span (#13285)

---
 .../pubsub/internal/tracing_batch_sink.cc     |  7 ++++-
 .../internal/tracing_batch_sink_test.cc       | 27 +++++++++++++++++++
 .../testing_util/opentelemetry_matchers.h     |  6 +++++
 3 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/google/cloud/pubsub/internal/tracing_batch_sink.cc b/google/cloud/pubsub/internal/tracing_batch_sink.cc
index 4fb685b111541..31fd113744ba5 100644
--- a/google/cloud/pubsub/internal/tracing_batch_sink.cc
+++ b/google/cloud/pubsub/internal/tracing_batch_sink.cc
@@ -60,6 +60,12 @@ auto MakeParent(Links const& links, Spans const& message_spans,
                 pubsub::Topic const& topic) {
   namespace sc = ::opentelemetry::trace::SemanticConventions;
   opentelemetry::trace::StartSpanOptions options;
+  opentelemetry::context::Context root_context;
+  // TODO(#13287): Use the constant instead of the string.
+  // Setting a span as a root span was added in OTel v1.13+. It is a no-op for
+  // earlier versions.
+  options.parent = root_context.SetValue(
+      /*opentelemetry::trace::kIsRootSpanKey=*/"is_root_span", true);
   options.kind = opentelemetry::trace::SpanKind::kClient;
   auto batch_sink_parent = internal::MakeSpan(
       topic.topic_id() + " publish",
@@ -74,7 +80,6 @@ auto MakeParent(Links const& links, Spans const& message_spans,
        {sc::kMessagingDestinationTemplate, topic.FullName()}},
       /*links*/ std::move(links), options);
 
-  // Add metadata to the message spans about the batch sink span.
   auto context = batch_sink_parent->GetContext();
   auto trace_id = internal::ToString(context.trace_id());
   auto span_id = internal::ToString(context.span_id());
diff --git a/google/cloud/pubsub/internal/tracing_batch_sink_test.cc b/google/cloud/pubsub/internal/tracing_batch_sink_test.cc
index b14a8fbecac22..7843a0193eeb8 100644
--- a/google/cloud/pubsub/internal/tracing_batch_sink_test.cc
+++ b/google/cloud/pubsub/internal/tracing_batch_sink_test.cc
@@ -43,6 +43,7 @@ using ::google::cloud::testing_util::SpanEventAttributesAre;
 using ::google::cloud::testing_util::SpanHasAttributes;
 using ::google::cloud::testing_util::SpanHasEvents;
 using ::google::cloud::testing_util::SpanHasInstrumentationScope;
+using ::google::cloud::testing_util::SpanIsRoot;
 using ::google::cloud::testing_util::SpanKindIsClient;
 using ::google::cloud::testing_util::SpanLinksSizeIs;
 using ::google::cloud::testing_util::SpanNamed;
@@ -216,6 +217,32 @@ TEST(TracingBatchSink, PublishSpanHasAttributes) {
                                  TestTopic().FullName())))));
 }
 
+#if OPENTELEMETRY_VERSION_MAJOR >= 2 || \
+    (OPENTELEMETRY_VERSION_MAJOR == 1 && OPENTELEMETRY_VERSION_MINOR >= 13)
+TEST(TracingBatchSink, PublishSpanIsRoot) {
+  auto span_catcher = InstallSpanCatcher();
+  auto message_span = MakeSpan("test span");
+  auto scope = opentelemetry::trace::Scope(message_span);
+  auto mock = std::make_unique<pubsub_testing::MockBatchSink>();
+  EXPECT_CALL(*mock, AddMessage(_));
+  EXPECT_CALL(*mock, AsyncPublish)
+      .WillOnce([](google::pubsub::v1::PublishRequest const& request) {
+        EXPECT_THAT(request, IsProtoEqual(MakeRequest(1)));
+        return make_ready_future(make_status_or(MakeResponse(request)));
+      });
+  auto batch_sink = MakeTestBatchSink(std::move(mock));
+  auto initial_spans = {message_span};
+  AddMessages(initial_spans, batch_sink);
+
+  auto response = batch_sink->AsyncPublish(MakeRequest(1)).get();
+  EXPECT_THAT(response, IsOk());
+
+  auto spans = span_catcher->GetSpans();
+  EXPECT_THAT(spans,
+              Contains(AllOf(SpanNamed("test-topic publish"), SpanIsRoot())));
+}
+#endif
+
 TEST(TracingBatchSink, AsyncPublishOnlyIncludeSampledLink) {
   namespace sc = ::opentelemetry::trace::SemanticConventions;
   // Create span before the span catcher so it is not sampled.
diff --git a/google/cloud/testing_util/opentelemetry_matchers.h b/google/cloud/testing_util/opentelemetry_matchers.h
index 5094258a704b8..b21efb175ae4b 100644
--- a/google/cloud/testing_util/opentelemetry_matchers.h
+++ b/google/cloud/testing_util/opentelemetry_matchers.h
@@ -150,6 +150,12 @@ MATCHER_P(SpanWithParent, span,
   return actual == span->GetContext().span_id();
 }
 
+MATCHER(SpanIsRoot, "is root span") {
+  auto const actual = arg->GetParentSpanId() == opentelemetry::trace::SpanId();
+  *result_listener << "is root span: " << (actual ? "true" : "false");
+  return actual;
+}
+
 MATCHER_P(SpanNamed, name, "has name: " + std::string{name}) {
   auto const& actual = arg->GetName();
   *result_listener << "has name: " << actual;