feat(storage): Link traces between Open and ReadRange traces.

Author: v-pratap

google/cloud/storage/internal/async/object_descriptor_connection_tracing.cc

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "google/cloud/storage/internal/async/object_descriptor_connection_tracing.h"
+#include "google/cloud/storage/internal/async/reader_connection_tracing.h"
 #include "google/cloud/storage/async/reader_connection.h"
 #include "google/cloud/internal/opentelemetry.h"
 #include "google/cloud/version.h"
@@ -58,7 +59,9 @@ class AsyncObjectDescriptorConnectionTracing
                     {{sc::kThreadId, internal::CurrentThreadId()},
                      {"read-start", p.start},
                      {"read-length", p.length}});
-    return result;
+    return MakeTracingReaderConnection(
+        span_,
+        std::move(result));
   }
 
  private:

google/cloud/storage/internal/async/object_descriptor_connection_tracing_test.cc

@@ -36,16 +36,39 @@ using ReadResponse =
 using ::google::cloud::storage_experimental::ObjectDescriptorConnection;
 using ::google::cloud::storage_mocks::MockAsyncObjectDescriptorConnection;
 using ::google::cloud::storage_mocks::MockAsyncReaderConnection;
+using ::google::cloud::storage_experimental::ReadPayload;
 using ::google::cloud::testing_util::EventNamed;
 using ::google::cloud::testing_util::InstallSpanCatcher;
 using ::google::cloud::testing_util::OTelAttribute;
+using ::google::cloud::testing_util::OTelContextCaptured;
+using ::google::cloud::testing_util::ThereIsAnActiveSpan;
+using ::google::cloud::testing_util::PromiseWithOTelContext;
 using ::google::cloud::testing_util::SpanEventAttributesAre;
 using ::google::cloud::testing_util::SpanHasInstrumentationScope;
 using ::google::cloud::testing_util::SpanKindIsClient;
 using ::google::cloud::testing_util::SpanNamed;
 using ::google::cloud::testing_util::SpanWithStatus;
 using ::testing::_;
 
+// A helper to set expectations on a mock async reader. It captures the OTel
+// context and returns a future that can be controlled by the test.
+auto expect_context = [](auto& p) {
+  return [&p] {
+    EXPECT_TRUE(ThereIsAnActiveSpan());
+    EXPECT_TRUE(OTelContextCaptured());
+    return p.get_future();
+  };
+};
+
+// A helper to be used in a `.then()` clause. It verifies the OTel context
+// has been detached before the user receives the result.
+auto expect_no_context = [](auto f) {
+  auto t = f.get();
+  EXPECT_FALSE(ThereIsAnActiveSpan());
+  EXPECT_FALSE(OTelContextCaptured());
+  return t;
+};
+
 TEST(ObjectDescriptorConnectionTracing, Read) {
   namespace sc = ::opentelemetry::trace::SemanticConventions;
   auto span_catcher = InstallSpanCatcher();
@@ -76,6 +99,56 @@ TEST(ObjectDescriptorConnectionTracing, Read) {
                           OTelAttribute<std::string>(sc::kThreadId, _)))))));
 }
 
+TEST(ObjectDescriptorConnectionTracing, ReadThenRead) {
+  namespace sc = ::opentelemetry::trace::SemanticConventions;
+  auto span_catcher = InstallSpanCatcher();
+
+  auto mock_connection = std::make_shared<MockAsyncObjectDescriptorConnection>();
+  auto* mock_reader_ptr = new MockAsyncReaderConnection;
+  PromiseWithOTelContext<ReadResponse> p;
+  EXPECT_CALL(*mock_reader_ptr, Read).WillOnce(expect_context(p));
+
+  EXPECT_CALL(*mock_connection, Read)
+      .WillOnce([&](ObjectDescriptorConnection::ReadParams) {
+        return std::unique_ptr<storage_experimental::AsyncReaderConnection>(
+            mock_reader_ptr);
+      });
+
+  auto connection = MakeTracingObjectDescriptorConnection(
+      internal::MakeSpan("test-span"), std::move(mock_connection));
+
+  auto reader = connection->Read({});
+  auto f = reader->Read().then(expect_no_context);
+  p.set_value(ReadPayload("test-payload").set_offset(123));
+  (void)f.get();
+
+  connection.reset();  // End the span
+
+  auto spans = span_catcher->GetSpans();
+  EXPECT_THAT(
+      spans,
+      ElementsAre(AllOf(
+          SpanNamed("test-span"),
+          SpanWithStatus(opentelemetry::trace::StatusCode::kOk),
+          SpanHasInstrumentationScope(), SpanKindIsClient(),
+          SpanHasEvents(
+              AllOf(EventNamed("gl-cpp.open.read"),
+                    SpanEventAttributesAre(
+                        OTelAttribute<std::int64_t>("read-length", 0),
+                        OTelAttribute<std::int64_t>("read-start", 0),
+                        OTelAttribute<std::string>(sc::kThreadId, _))),
+              AllOf(
+                  EventNamed("gl-cpp.read"),
+                  SpanEventAttributesAre(
+                      OTelAttribute<std::int64_t>("message.starting_offset",
+                                                  123),
+                      OTelAttribute<std::string>(sc::kThreadId, _),
+                      OTelAttribute<std::int64_t>("rpc.message.id", 1),
+                      // THIS WAS THE MISSING ATTRIBUTE:
+                      OTelAttribute<std::string>("rpc.message.type",
+                                                 "RECEIVED")))))));
+}
+
 }  // namespace
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace storage_internal

google/cloud/storage/internal/async/object_descriptor_reader_tracing.cc

@@ -41,30 +41,15 @@ class ObjectDescriptorReaderTracing : public ObjectDescriptorReader {
   ~ObjectDescriptorReaderTracing() override = default;
 
   future<ObjectDescriptorReader::ReadResponse> Read() override {
-    auto span = internal::MakeSpan("storage::AsyncConnection::ReadObjectRange");
-    internal::OTelScope scope(span);
     return ObjectDescriptorReader::Read().then(
-        [span = std::move(span),
-         oc = opentelemetry::context::RuntimeContext::GetCurrent()](
+        [](
             auto f) -> ReadResponse {
           auto result = f.get();
-          internal::DetachOTelContext(oc);
+          // internal::DetachOTelContext(oc);
           if (!absl::holds_alternative<Status>(result)) {
             auto const& payload =
                 absl::get<storage_experimental::ReadPayload>(result);
 
-            span->AddEvent(
-                "gl-cpp.read-range",
-                {{/*sc::kRpcMessageType=*/"rpc.message.type", "RECEIVED"},
-                 {sc::kThreadId, internal::CurrentThreadId()},
-                 {"message.size", static_cast<std::uint32_t>(payload.size())}});
-          } else {
-            span->AddEvent(
-                "gl-cpp.read-range",
-                {{/*sc::kRpcMessageType=*/"rpc.message.type", "RECEIVED"},
-                 {sc::kThreadId, internal::CurrentThreadId()}});
-            return internal::EndSpan(*span,
-                                     absl::get<Status>(std::move(result)));
           }
           return result;
         });