From 3dddc28d63088f164350b0b380fc8afdb99a3eac Mon Sep 17 00:00:00 2001
From: Darren Bolduc <dbolduc@google.com>
Date: Tue, 17 Oct 2023 10:43:03 -0400
Subject: [PATCH] impl(otel): detach context from EndSpan(..., future<>)
 (#12902)

---
 google/cloud/internal/grpc_opentelemetry.h    |  7 +++--
 .../cloud/internal/grpc_opentelemetry_test.cc | 27 +++++++++++++++++++
 google/cloud/internal/opentelemetry.h         |  9 +++++--
 google/cloud/internal/opentelemetry_test.cc   | 23 ++++++++++++++++
 4 files changed, 62 insertions(+), 4 deletions(-)

diff --git a/google/cloud/internal/grpc_opentelemetry.h b/google/cloud/internal/grpc_opentelemetry.h
index c5e957953594d..eb8f078688dca 100644
--- a/google/cloud/internal/grpc_opentelemetry.h
+++ b/google/cloud/internal/grpc_opentelemetry.h
@@ -90,9 +90,12 @@ future<T> EndSpan(
     std::shared_ptr<grpc::ClientContext> context,
     opentelemetry::nostd::shared_ptr<opentelemetry::trace::Span> span,
     future<T> fut) {
-  return fut.then([c = std::move(context), s = std::move(span)](auto f) {
+  return fut.then([oc = opentelemetry::context::RuntimeContext::GetCurrent(),
+                   c = std::move(context), s = std::move(span)](auto f) {
+    auto t = f.get();
     ExtractAttributes(*c, *s);
-    return EndSpan(*s, f.get());
+    DetachOTelContext(oc);
+    return EndSpan(*s, std::move(t));
   });
 }
 
diff --git a/google/cloud/internal/grpc_opentelemetry_test.cc b/google/cloud/internal/grpc_opentelemetry_test.cc
index ca1f17fdba2db..88cfb41272e9d 100644
--- a/google/cloud/internal/grpc_opentelemetry_test.cc
+++ b/google/cloud/internal/grpc_opentelemetry_test.cc
@@ -140,6 +140,33 @@ TEST(OpenTelemetry, EndSpanFuture) {
           SpanHasAttributes(OTelAttribute<std::string>("grpc.peer", _)))));
 }
 
+TEST(OpenTelemetry, EndSpanFutureDetachesContext) {
+  auto span_catcher = InstallSpanCatcher();
+
+  // Set the `OTelContext` like we do in `AsyncOperation`s
+  auto c = opentelemetry::context::Context("key", true);
+  ScopedOTelContext scope({c});
+  EXPECT_EQ(c, opentelemetry::context::RuntimeContext::GetCurrent());
+
+  promise<Status> p;
+  auto f =
+      EndSpan(std::make_shared<grpc::ClientContext>(),
+              MakeSpanGrpc("google.cloud.foo.v1.Foo", "GetBar"), p.get_future())
+          .then([](auto f) {
+            auto s = f.get();
+            // The `OTelContext` should be cleared by the time we exit
+            // `EndSpan()`.
+            EXPECT_EQ(opentelemetry::context::Context{},
+                      opentelemetry::context::RuntimeContext::GetCurrent());
+            return s;
+          });
+
+  p.set_value(Status());
+  EXPECT_STATUS_OK(f.get());
+  EXPECT_THAT(span_catcher->GetSpans(),
+              ElementsAre(SpanNamed("google.cloud.foo.v1.Foo/GetBar")));
+}
+
 TEST(OpenTelemetry, TracedAsyncBackoffEnabled) {
   auto span_catcher = InstallSpanCatcher();
 
diff --git a/google/cloud/internal/opentelemetry.h b/google/cloud/internal/opentelemetry.h
index 270d79785e239..01f91bc804eae 100644
--- a/google/cloud/internal/opentelemetry.h
+++ b/google/cloud/internal/opentelemetry.h
@@ -16,6 +16,7 @@
 #define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_INTERNAL_OPENTELEMETRY_H
 
 #include "google/cloud/future.h"
+#include "google/cloud/internal/opentelemetry_context.h"
 #include "google/cloud/options.h"
 #include "google/cloud/status.h"
 #include "google/cloud/status_or.h"
@@ -178,8 +179,12 @@ template <typename T>
 future<T> EndSpan(
     opentelemetry::nostd::shared_ptr<opentelemetry::trace::Span> span,
     future<T> fut) {
-  return fut.then(
-      [s = std::move(span)](auto f) { return EndSpan(*s, f.get()); });
+  return fut.then([oc = opentelemetry::context::RuntimeContext::GetCurrent(),
+                   s = std::move(span)](auto f) {
+    auto t = f.get();
+    DetachOTelContext(oc);
+    return EndSpan(*s, std::move(t));
+  });
 }
 
 /**
diff --git a/google/cloud/internal/opentelemetry_test.cc b/google/cloud/internal/opentelemetry_test.cc
index c7859919ef6cb..7f178a56a2e51 100644
--- a/google/cloud/internal/opentelemetry_test.cc
+++ b/google/cloud/internal/opentelemetry_test.cc
@@ -16,6 +16,7 @@
 #include "google/cloud/internal/make_status.h"
 #include "google/cloud/opentelemetry_options.h"
 #include "google/cloud/testing_util/opentelemetry_matchers.h"
+#include "google/cloud/testing_util/status_matchers.h"
 #include <gmock/gmock.h>
 #ifdef GOOGLE_CLOUD_CPP_HAVE_OPENTELEMETRY
 #include <opentelemetry/trace/default_span.h>
@@ -269,6 +270,28 @@ TEST(OpenTelemetry, EndSpanFutureStatusOr) {
                   SpanWithStatus(opentelemetry::trace::StatusCode::kError)));
 }
 
+TEST(OpenTelemetry, EndSpanFutureDetachesContext) {
+  auto span_catcher = InstallSpanCatcher();
+
+  // Set the `OTelContext` like we do in `AsyncOperation`s
+  auto c = opentelemetry::context::Context("key", true);
+  ScopedOTelContext scope({c});
+  EXPECT_EQ(c, opentelemetry::context::RuntimeContext::GetCurrent());
+
+  promise<Status> p;
+  auto f = EndSpan(MakeSpan("span"), p.get_future()).then([](auto f) {
+    auto s = f.get();
+    // The `OTelContext` should be cleared by the time we exit `EndSpan()`.
+    EXPECT_EQ(opentelemetry::context::Context{},
+              opentelemetry::context::RuntimeContext::GetCurrent());
+    return s;
+  });
+
+  p.set_value(Status());
+  EXPECT_STATUS_OK(f.get());
+  EXPECT_THAT(span_catcher->GetSpans(), ElementsAre(SpanNamed("span")));
+}
+
 TEST(OpenTelemetry, EndSpanVoid) {
   auto span_catcher = InstallSpanCatcher();