From 41714ae283854a31955cb31fb5b68326806cd075 Mon Sep 17 00:00:00 2001
From: James Crosswell <jamescrosswell@users.noreply.github.com>
Date: Tue, 10 Oct 2023 22:57:47 +1300
Subject: [PATCH] Apply Scope to exceptions captured via
 Activity.RecordException (#2712)

---
 src/Sentry.AspNetCore/ScopeExtensions.cs      |  4 ++-
 .../SentrySpanProcessor.cs                    | 25 +++++++++++++-----
 .../SentrySpanProcessorTests.cs               | 26 +++++++++++++++++--
 3 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/src/Sentry.AspNetCore/ScopeExtensions.cs b/src/Sentry.AspNetCore/ScopeExtensions.cs
index 44eecf786c..10249aa492 100644
--- a/src/Sentry.AspNetCore/ScopeExtensions.cs
+++ b/src/Sentry.AspNetCore/ScopeExtensions.cs
@@ -169,7 +169,9 @@ private static void SetEnv(Scope scope, HttpContext context, SentryAspNetCoreOpt
 
     private static void SetBody(Scope scope, HttpContext context, SentryAspNetCoreOptions options)
     {
-        var extractors = context.RequestServices.GetService<IEnumerable<IRequestPayloadExtractor>>();
+        // Resharper says this can't be null, but actually it can if SetBody is called multiple times in a single request
+        // ReSharper disable once ConditionalAccessQualifierIsNonNullableAccordingToAPIContract
+        var extractors = context.RequestServices?.GetService<IEnumerable<IRequestPayloadExtractor>>();
         if (extractors == null)
         {
             return;
diff --git a/src/Sentry.OpenTelemetry/SentrySpanProcessor.cs b/src/Sentry.OpenTelemetry/SentrySpanProcessor.cs
index 217ff88120..dcdff5fe05 100644
--- a/src/Sentry.OpenTelemetry/SentrySpanProcessor.cs
+++ b/src/Sentry.OpenTelemetry/SentrySpanProcessor.cs
@@ -157,12 +157,6 @@ public override void OnEnd(Activity data)
 
             // Transactions set otel attributes (and resource attributes) as context.
             transaction.Contexts["otel"] = GetOtelContext(attributes);
-            // Events are received/processed in a different AsyncLocal context. Restoring the scope that started it.
-            var activityScope = data.GetFused<Scope>();
-            if (activityScope is { } savedScope && _hub is Hub hub)
-            {
-                hub.RestoreScope(savedScope);
-            }
         }
         else
         {
@@ -175,6 +169,12 @@ public override void OnEnd(Activity data)
             span.SetExtra("otel.kind", data.Kind);
         }
 
+        // Events are received/processed in a different AsyncLocal context. Restoring the scope that started it.
+        var activityScope = GetSavedScope(data);
+        if (activityScope is { } savedScope && _hub is Hub hub)
+        {
+            hub.RestoreScope(savedScope);
+        }
         GenerateSentryErrorsFromOtelSpan(data, attributes);
 
         var status = GetSpanStatus(data.Status, attributes);
@@ -187,6 +187,19 @@ public override void OnEnd(Activity data)
         _map.TryRemove(data.SpanId, out _);
     }
 
+    private static Scope? GetSavedScope(Activity? activity)
+    {
+        while (activity is not null)
+        {
+            if (activity.GetFused<Scope>() is {} savedScope)
+            {
+                return savedScope;
+            }
+            activity = activity.Parent;
+        }
+        return null;
+    }
+
     internal static SpanStatus GetSpanStatus(ActivityStatusCode status, IDictionary<string, object?> attributes)
     {
         // See https://github.com/open-telemetry/opentelemetry-dotnet/discussions/4703
diff --git a/test/Sentry.OpenTelemetry.Tests/SentrySpanProcessorTests.cs b/test/Sentry.OpenTelemetry.Tests/SentrySpanProcessorTests.cs
index dc0de2b27d..4a7f87e757 100644
--- a/test/Sentry.OpenTelemetry.Tests/SentrySpanProcessorTests.cs
+++ b/test/Sentry.OpenTelemetry.Tests/SentrySpanProcessorTests.cs
@@ -264,7 +264,7 @@ public void OnEnd_FinishesSpan()
     }
 
     [Fact]
-    public void OnEnd_RestoresSavedScope()
+    public void OnEnd_Transaction_RestoresSavedScope()
     {
         // Arrange
         _fixture.Options.Instrumenter = Instrumenter.OpenTelemetry;
@@ -276,7 +276,29 @@ public void OnEnd_RestoresSavedScope()
         data.SetFused(scope);
         sut.OnStart(data);
 
-        sut._map.TryGetValue(data.SpanId, out var span);
+        // Act
+        sut.OnEnd(data);
+
+        // Assert
+        _fixture.ScopeManager.Received(1).RestoreScope(scope);
+    }
+
+    [Fact]
+    public void OnEnd_Span_RestoresSavedScope()
+    {
+        // Arrange
+        _fixture.Options.Instrumenter = Instrumenter.OpenTelemetry;
+        _fixture.ScopeManager = Substitute.For<IInternalScopeManager>();
+        var sut = _fixture.GetSut();
+
+        var scope = new Scope();
+        var parent = Tracer.StartActivity("transaction")!;
+        parent.SetFused(scope);
+        sut.OnStart(parent);
+
+        var data = Tracer.StartActivity("test operation")!;
+        data.DisplayName = "test display name";
+        sut.OnStart(data);
 
         // Act
         sut.OnEnd(data);