Fix Restoring the Original Activity Parent (#48034)

Author: tarekgh

src/libraries/System.Diagnostics.DiagnosticSource/src/System/Diagnostics/Activity.cs

@@ -84,6 +84,7 @@ public partial class Activity : IDisposable
         private string? _displayName;
         private ActivityStatusCode _statusCode;
         private string? _statusDescription;
+        private Activity? _previousActiveActivity;
 
         /// <summary>
         /// Gets status code of the current activity object.
@@ -621,15 +622,15 @@ public Activity Start()
             }
             else
             {
+                _previousActiveActivity = Current;
                 if (_parentId == null && _parentSpanId is null)
                 {
-                    Activity? parent = Current;
-                    if (parent != null)
+                    if (_previousActiveActivity != null)
                     {
                         // The parent change should not form a loop.   We are actually guaranteed this because
                         // 1. Un-started activities can't be 'Current' (thus can't be 'parent'), we throw if you try.
                         // 2. All started activities have a finite parent change (by inductive reasoning).
-                        Parent = parent;
+                        Parent = _previousActiveActivity;
                     }
                 }
 
@@ -686,8 +687,7 @@ public void Stop()
                 }
 
                 Source.NotifyActivityStop(this);
-
-                SetCurrent(Parent);
+                SetCurrent(_previousActiveActivity);
             }
         }
 
@@ -1004,6 +1004,7 @@ internal static Activity CreateAndStart(ActivitySource source, string name, Acti
             activity.Source = source;
             activity.Kind = kind;
 
+            activity._previousActiveActivity = Current;
             if (parentId != null)
             {
                 activity._parentId = parentId;
@@ -1022,13 +1023,12 @@ internal static Activity CreateAndStart(ActivitySource source, string name, Acti
             }
             else
             {
-                Activity? parent = Current;
-                if (parent != null)
+                if (activity._previousActiveActivity != null)
                 {
                     // The parent change should not form a loop. We are actually guaranteed this because
                     // 1. Un-started activities can't be 'Current' (thus can't be 'parent'), we throw if you try.
                     // 2. All started activities have a finite parent change (by inductive reasoning).
-                    activity.Parent = parent;
+                    activity.Parent = activity._previousActiveActivity;
                 }
             }
 

src/libraries/System.Diagnostics.DiagnosticSource/tests/ActivitySourceTests.cs

@@ -845,6 +845,42 @@ public void TestAddSamplerAndActivityCreationTags()
             }).Dispose();
         }
 
+        [ConditionalFact(typeof(RemoteExecutor), nameof(RemoteExecutor.IsSupported))]
+        public void RestoreOriginalParentTest()
+        {
+            RemoteExecutor.Invoke(() =>
+            {
+                using ActivitySource source = new ActivitySource("OriginalParentSource");
+                using ActivityListener listener = new ActivityListener();
+
+                listener.ShouldListenTo = (activitySource) => object.ReferenceEquals(source, activitySource);
+                listener.SampleUsingParentId = (ref ActivityCreationOptions<string> activityOptions) => ActivitySamplingResult.AllData;
+                listener.Sample = (ref ActivityCreationOptions<ActivityContext> activityOptions) => ActivitySamplingResult.AllData;
+                ActivitySource.AddActivityListener(listener);
+
+                Assert.Null(Activity.Current);
+
+                using (Activity c = source.StartActivity("Root"))
+                {
+                    Assert.NotNull(Activity.Current);
+                    Assert.Equal("Root", Activity.Current.OperationName);
+
+                    // Create Activity with the parent context to not use Activity.Current as a parent
+                    using (Activity d = source.StartActivity("Child", default, new ActivityContext(ActivityTraceId.CreateRandom(), ActivitySpanId.CreateRandom(), default, default)))
+                    {
+                        Assert.NotNull(Activity.Current);
+                        Assert.Equal("Child", Activity.Current.OperationName);
+                    }
+
+                    // Now the child activity stopped. We used to restore null to the Activity.Current but now we restore
+                    // the original parent stored in Activity.Current before we started the Activity.
+                    Assert.NotNull(Activity.Current);
+                    Assert.Equal("Root", Activity.Current.OperationName);
+                }
+                Assert.Null(Activity.Current);
+            }).Dispose();
+        }
+
         public void Dispose() => Activity.Current = null;
     }
 }

src/libraries/System.Diagnostics.DiagnosticSource/tests/ActivityTests.cs

@@ -1751,6 +1751,40 @@ public void StructEnumerator_GenericLinkedList()
             Assert.True(method.ReturnType.IsValueType);
         }
 
+        [ConditionalFact(typeof(RemoteExecutor), nameof(RemoteExecutor.IsSupported))]
+        public void RestoreOriginalParentTest()
+        {
+            RemoteExecutor.Invoke(() =>
+            {
+                Assert.Null(Activity.Current);
+
+                Activity a = new Activity("Root");
+                a.Start();
+
+                Assert.NotNull(Activity.Current);
+                Assert.Equal("Root", Activity.Current.OperationName);
+
+                // Create Activity with the parent context to not use Activity.Current as a parent
+                Activity b = new Activity("Child");
+                b.SetParentId(ActivityTraceId.CreateRandom(), ActivitySpanId.CreateRandom());
+                b.Start();
+
+                Assert.NotNull(Activity.Current);
+                Assert.Equal("Child", Activity.Current.OperationName);
+
+                b.Stop();
+
+                // Now the child activity stopped. We used to restore null to the Activity.Current but now we restore
+                // the original parent stored in Activity.Current before we started the Activity.
+                Assert.NotNull(Activity.Current);
+                Assert.Equal("Root", Activity.Current.OperationName);
+
+                a.Stop();
+                Assert.Null(Activity.Current);
+
+            }).Dispose();
+        }
+
         public void Dispose()
         {
             Activity.Current = null;