From 04d4500924848ba5647b97410fdd51bf610dac42 Mon Sep 17 00:00:00 2001
From: dbolduc <dbolduc@google.com>
Date: Thu, 27 Jun 2024 19:03:09 -0400
Subject: [PATCH] fix(otel): async read stream tracing

---
 .../internal/async_streaming_read_rpc_tracing.h      | 12 ++++++++++--
 .../async_streaming_read_rpc_tracing_test.cc         | 11 ++++++-----
 2 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/google/cloud/internal/async_streaming_read_rpc_tracing.h b/google/cloud/internal/async_streaming_read_rpc_tracing.h
index 50007b7b87c50..5bd1cdc956adf 100644
--- a/google/cloud/internal/async_streaming_read_rpc_tracing.h
+++ b/google/cloud/internal/async_streaming_read_rpc_tracing.h
@@ -43,7 +43,11 @@ class AsyncStreamingReadRpcTracing : public AsyncStreamingReadRpc<Response> {
   }
 
   future<bool> Start() override {
-    auto start_span = internal::MakeSpan("Start");
+    // It is sufficient to set `span_` as the parent of `start_span`, because
+    // the lower levels do not create any spans.
+    opentelemetry::trace::StartSpanOptions options;
+    options.parent = span_->GetContext();
+    auto start_span = internal::MakeSpan("Start", options);
     return impl_->Start().then(
         [this, ss = std::move(start_span)](future<bool> f) {
           EndSpan(*ss);
@@ -66,7 +70,11 @@ class AsyncStreamingReadRpcTracing : public AsyncStreamingReadRpc<Response> {
   }
 
   future<Status> Finish() override {
-    auto finish_span = internal::MakeSpan("Finish");
+    // It is sufficient to set `span_` as the parent of `finish_span`, because
+    // the lower levels do not create any spans.
+    opentelemetry::trace::StartSpanOptions options;
+    options.parent = span_->GetContext();
+    auto finish_span = internal::MakeSpan("Finish", options);
     return impl_->Finish().then(
         [this, fs = std::move(finish_span)](future<Status> f) {
           EndSpan(*fs);
diff --git a/google/cloud/internal/async_streaming_read_rpc_tracing_test.cc b/google/cloud/internal/async_streaming_read_rpc_tracing_test.cc
index b2620e60fe3d2..9d1109e832228 100644
--- a/google/cloud/internal/async_streaming_read_rpc_tracing_test.cc
+++ b/google/cloud/internal/async_streaming_read_rpc_tracing_test.cc
@@ -35,6 +35,7 @@ using ::google::cloud::testing_util::SetServerMetadata;
 using ::google::cloud::testing_util::SpanEventAttributesAre;
 using ::google::cloud::testing_util::SpanHasAttributes;
 using ::google::cloud::testing_util::SpanNamed;
+using ::google::cloud::testing_util::SpanWithParent;
 using ::google::cloud::testing_util::SpanWithStatus;
 using ::google::cloud::testing_util::StatusIs;
 using ::testing::_;
@@ -79,7 +80,7 @@ TEST(AsyncStreamingReadRpcTracing, Cancel) {
                         SpanEventsAre(
                             EventNamed("gl-cpp.cancel"),
                             EventNamed("test-only: underlying stream cancel"))),
-                  SpanNamed("Finish")));
+                  AllOf(SpanNamed("Finish"), SpanWithParent(span))));
 }
 
 TEST(AsyncStreamingReadRpcTracing, Start) {
@@ -98,10 +99,10 @@ TEST(AsyncStreamingReadRpcTracing, Start) {
   auto spans = span_catcher->GetSpans();
   EXPECT_THAT(
       spans, UnorderedElementsAre(
-                 SpanNamed("Start"),
+                 AllOf(SpanNamed("Start"), SpanWithParent(span)),
                  AllOf(SpanNamed("span"), SpanHasAttributes(OTelAttribute<bool>(
                                               "gl-cpp.stream_started", true))),
-                 SpanNamed("Finish")));
+                 AllOf(SpanNamed("Finish"), SpanWithParent(span))));
 }
 
 TEST(AsyncStreamingReadRpcTracing, Read) {
@@ -146,7 +147,7 @@ TEST(AsyncStreamingReadRpcTracing, Read) {
                                         OTelAttribute<std::string>(
                                             "message.type", "RECEIVED"),
                                         OTelAttribute<int>("message.id", 3))))),
-          SpanNamed("Finish")));
+          AllOf(SpanNamed("Finish"), SpanWithParent(span))));
 }
 
 TEST(AsyncStreamingReadRpcTracing, Finish) {
@@ -168,7 +169,7 @@ TEST(AsyncStreamingReadRpcTracing, Finish) {
               SpanNamed("span"),
               SpanHasAttributes(OTelAttribute<std::string>("grpc.peer", _)),
               SpanWithStatus(opentelemetry::trace::StatusCode::kError, "fail")),
-          SpanNamed("Finish")));
+          AllOf(SpanNamed("Finish"), SpanWithParent(span))));
 }
 
 TEST(AsyncStreamingReadRpcTracing, GetRequestMetadata) {