fix(otel): async write stream tracing (#14394)

googleapis · Jun 28, 2024 · c883817 · c883817
1 parent f384d1a
commit c883817
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 8 deletions.
diff --git a/google/cloud/internal/async_streaming_write_rpc_tracing.h b/google/cloud/internal/async_streaming_write_rpc_tracing.h
@@ -48,7 +48,11 @@ class AsyncStreamingWriteRpcTracing
   }
 
   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);
@@ -82,7 +86,11 @@ class AsyncStreamingWriteRpcTracing
   }
 
   future<StatusOr<Response>> 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<StatusOr<Response>> f) {
           EndSpan(*fs);

diff --git a/google/cloud/internal/async_streaming_write_rpc_tracing_test.cc b/google/cloud/internal/async_streaming_write_rpc_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,7 +79,7 @@ TEST(AsyncStreamingWriteRpcTracing, Cancel) {
                         SpanEventsAre(
                             EventNamed("gl-cpp.cancel"),
                             EventNamed("test-only: underlying stream cancel"))),
-                  SpanNamed("Finish")));
+                  AllOf(SpanNamed("Finish"), SpanWithParent(span))));
 }
 
 TEST(AsyncStreamingWriteRpcTracing, Start) {
@@ -98,10 +99,10 @@ TEST(AsyncStreamingWriteRpcTracing, 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(AsyncStreamingWriteRpcTracing, Write) {
@@ -149,7 +150,7 @@ TEST(AsyncStreamingWriteRpcTracing, 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(AsyncStreamingWriteRpcTracing, WritesDone) {
@@ -171,7 +172,7 @@ TEST(AsyncStreamingWriteRpcTracing, WritesDone) {
                          AllOf(SpanNamed("span"),
                                SpanEventsAre(EventNamed("gl-cpp.first-write"),
                                              EventNamed("gl-cpp.writes_done"))),
-                         SpanNamed("Finish")));
+                         AllOf(SpanNamed("Finish"), SpanWithParent(span))));
 }
 
 TEST(AsyncStreamingWriteRpcTracing, Finish) {
@@ -194,7 +195,7 @@ TEST(AsyncStreamingWriteRpcTracing, Finish) {
               SpanNamed("span"),
               SpanHasAttributes(OTelAttribute<std::string>("grpc.peer", _)),
               SpanWithStatus(opentelemetry::trace::StatusCode::kError, "fail")),
-          SpanNamed("Finish")));
+          AllOf(SpanNamed("Finish"), SpanWithParent(span))));
 }
 
 TEST(AsyncStreamingWriteRpcTracing, GetRequestMetadata) {