fix(otel): async read write stream tracing (#14379)

googleapis · Jun 27, 2024 · d0e00f9 · d0e00f9
1 parent 52e16d5
commit d0e00f9
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 11 deletions.
diff --git a/google/cloud/internal/async_read_write_stream_tracing.h b/google/cloud/internal/async_read_write_stream_tracing.h
@@ -44,7 +44,11 @@ class AsyncStreamingReadWriteRpcTracing
   }
 
   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);
@@ -89,7 +93,12 @@ class AsyncStreamingReadWriteRpcTracing
   }
 
   future<Status> Finish() override {
-    auto finish_span = internal::MakeSpan("Finish");
+    internal::OTelScope scope(span_);
+    // 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_read_write_stream_tracing_test.cc b/google/cloud/internal/async_read_write_stream_tracing_test.cc
@@ -34,6 +34,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::_;
@@ -78,12 +79,12 @@ TEST(AsyncStreamingReadWriteRpcTracing, Cancel) {
   auto spans = span_catcher->GetSpans();
   EXPECT_THAT(spans,
               UnorderedElementsAre(
-                  SpanNamed("Start"),
+                  AllOf(SpanNamed("Start"), SpanWithParent(span)),
                   AllOf(SpanNamed("span"),
                         SpanEventsAre(
                             EventNamed("gl-cpp.cancel"),
                             EventNamed("test-only: underlying stream cancel"))),
-                  SpanNamed("Finish")));
+                  AllOf(SpanNamed("Finish"), SpanWithParent(span))));
 }
 
 TEST(AsyncStreamingReadWriteRpcTracing, Start) {
@@ -102,10 +103,10 @@ TEST(AsyncStreamingReadWriteRpcTracing, 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(AsyncStreamingReadWriteRpcTracing, Read) {
@@ -150,7 +151,7 @@ TEST(AsyncStreamingReadWriteRpcTracing, Read) {
                                         OTelAttribute<std::string>(
                                             "message.type", "RECEIVED"),
                                         OTelAttribute<int>("message.id", 3))))),
-          SpanNamed("Finish")));
+          AllOf(SpanNamed("Finish"), SpanWithParent(span))));
 }
 
 TEST(AsyncStreamingReadWriteRpcTracing, Write) {
@@ -197,7 +198,7 @@ TEST(AsyncStreamingReadWriteRpcTracing, Write) {
                             OTelAttribute<int>("message.id", 3),
                             OTelAttribute<bool>("message.is_last", true),
                             OTelAttribute<bool>("message.success", true))))),
-          SpanNamed("Finish")));
+          AllOf(SpanNamed("Finish"), SpanWithParent(span))));
 }
 
 TEST(AsyncStreamingReadWriteRpcTracing, SeparateCountersForReadAndWrite) {
@@ -237,7 +238,7 @@ TEST(AsyncStreamingReadWriteRpcTracing, SeparateCountersForReadAndWrite) {
                             OTelAttribute<std::string>("message.type",
                                                        "RECEIVED"),
                             OTelAttribute<int>("message.id", 1))))),
-          SpanNamed("Finish")));
+          AllOf(SpanNamed("Finish"), SpanWithParent(span))));
 }
 
 TEST(AsyncStreamingReadWriteRpcTracing, WritesDone) {
@@ -257,7 +258,7 @@ TEST(AsyncStreamingReadWriteRpcTracing, WritesDone) {
   EXPECT_THAT(spans, UnorderedElementsAre(
                          AllOf(SpanNamed("span"),
                                SpanEventsAre(EventNamed("gl-cpp.writes_done"))),
-                         SpanNamed("Finish")));
+                         AllOf(SpanNamed("Finish"), SpanWithParent(span))));
 }
 
 TEST(AsyncStreamingReadWriteRpcTracing, Finish) {
@@ -279,7 +280,7 @@ TEST(AsyncStreamingReadWriteRpcTracing, Finish) {
               SpanNamed("span"),
               SpanHasAttributes(OTelAttribute<std::string>("grpc.peer", _)),
               SpanWithStatus(opentelemetry::trace::StatusCode::kError, "fail")),
-          SpanNamed("Finish")));
+          AllOf(SpanNamed("Finish"), SpanWithParent(span))));
 }
 
 TEST(AsyncStreamingReadWriteRpcTracing, GetRequestMetadata) {