Prevent ActivityId leak across threads

src/libraries/System.Diagnostics.Tracing/tests/BasicEventSourceTest/ActivityTracking.cs

@@ -0,0 +1,154 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Collections.Generic;
+using System.Diagnostics.Tracing;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+using Xunit;
+
+namespace BasicEventSourceTests
+{
+    public class ActivityTracking
+    {
+        [ConditionalFact(typeof(PlatformDetection), nameof(PlatformDetection.IsNotBrowser))]
+        public void StartStopCreatesActivity()
+        {
+            using ActivityEventListener l = new ActivityEventListener();
+            using ActivityEventSource es = new ActivityEventSource();
+
+            Assert.Equal(Guid.Empty, EventSource.CurrentThreadActivityId);
+            es.ExampleStart();
+            Assert.NotEqual(Guid.Empty, EventSource.CurrentThreadActivityId);
+            es.ExampleStop();
+            Assert.Equal(Guid.Empty, EventSource.CurrentThreadActivityId);
+        }
+
+        [ConditionalFact(typeof(PlatformDetection), nameof(PlatformDetection.IsNotBrowser))]
+        public async Task ActivityFlowsAsync()
+        {
+            using ActivityEventListener l = new ActivityEventListener();
+            using ActivityEventSource es = new ActivityEventSource();
+
+            Assert.Equal(Guid.Empty, EventSource.CurrentThreadActivityId);
+            es.ExampleStart();
+            Assert.NotEqual(Guid.Empty, EventSource.CurrentThreadActivityId);
+            await Task.Yield();
+            Assert.NotEqual(Guid.Empty, EventSource.CurrentThreadActivityId);
+            es.ExampleStop();
+            Assert.Equal(Guid.Empty, EventSource.CurrentThreadActivityId);
+        }
+
+        [ConditionalFact(typeof(PlatformDetection), nameof(PlatformDetection.IsNotBrowser))]
+        public async Task ActivityIdIsZeroedOnThreadSwitchOut()
+        {
+            using ActivityEventListener l = new ActivityEventListener();
+            using ActivityEventSource es = new ActivityEventSource();
+
+            // Run tasks on many threads. If an activity id leaks it is likely
+            // that the thread will be sheduled to run one of our other tasks
+            // and we can detect the non-zero id at the start of the task
+            List<Task> tasks = new List<Task>();
+            for (int i = 0; i < 100; i++)
+            {
+                tasks.Add(Task.Run(() => YieldTwoActivitiesDeep(es)));
+            }
+
+            await Task.WhenAll(tasks.ToArray());
+        }
+
+        private async Task YieldTwoActivitiesDeep(ActivityEventSource es)
+        {
+            Assert.Equal(Guid.Empty, EventSource.CurrentThreadActivityId);
+            es.ExampleStart();
+            es.Example2Start();
+            await Task.Yield();
+            es.Example2Stop();
+            es.ExampleStop();
+            Assert.Equal(Guid.Empty, EventSource.CurrentThreadActivityId);
+        }
+
+        // I don't know if this behavior is by-design or accidental. For now
+        // I am attempting to preserve it to lower back compat risk, but in
+        // the future we might decide it wasn't even desirable to begin with.
+        // Compare with SetCurrentActivityIdAfterEventDoesNotFlowAsync below.
+        [ConditionalFact(typeof(PlatformDetection), nameof(PlatformDetection.IsNotBrowser))]
+        public async Task SetCurrentActivityIdBeforeEventFlowsAsync()
+        {
+            using ActivityEventListener l = new ActivityEventListener();
+            using ActivityEventSource es = new ActivityEventSource();
+            try
+            {
+                Guid g = Guid.NewGuid();
+                EventSource.SetCurrentThreadActivityId(g);
+                Assert.Equal(g, EventSource.CurrentThreadActivityId);
+                es.ExampleStart();
+                await Task.Yield();
+                es.ExampleStop();
+                Assert.Equal(g, EventSource.CurrentThreadActivityId);
+            }
+            finally
+            {
+                EventSource.SetCurrentThreadActivityId(Guid.Empty);
+            }
+        }
+
+        // I don't know if this behavior is by-design or accidental. For now
+        // I am attempting to preserve it to lower back compat risk, but in
+        // the future we might decide it wasn't even desirable to begin with.
+        // Compare with SetCurrentActivityIdBeforeEventFlowsAsync above.
+        [ConditionalFact(typeof(PlatformDetection), nameof(PlatformDetection.IsNotBrowser))]
+        public async Task SetCurrentActivityIdAfterEventDoesNotFlowAsync()
+        {
+            using ActivityEventListener l = new ActivityEventListener();
+            using ActivityEventSource es = new ActivityEventSource();
+            try
+            {
+                es.ExampleStart();
+                Guid g = Guid.NewGuid();
+                EventSource.SetCurrentThreadActivityId(g);
+                Assert.Equal(g, EventSource.CurrentThreadActivityId);
+                await Task.Yield();
+                Assert.NotEqual(g, EventSource.CurrentThreadActivityId);
+                es.ExampleStop();
+            }
+            finally
+            {
+                EventSource.SetCurrentThreadActivityId(Guid.Empty);
+            }
+        }
+    }
+
+    [EventSource(Name = "ActivityEventSource")]
+    class ActivityEventSource : EventSource
+    {
+        [Event(1)]
+        public void ExampleStart() => WriteEvent(1);
+
+        [Event(2)]
+        public void ExampleStop() => WriteEvent(2);
+
+        [Event(3)]
+        public void Example2Start() => WriteEvent(3);
+
+        [Event(4)]
+        public void Example2Stop() => WriteEvent(4);
+    }
+
+    class ActivityEventListener : EventListener
+    {
+        protected override void OnEventSourceCreated(EventSource eventSource)
+        {
+            if (eventSource.Name == "System.Threading.Tasks.TplEventSource")
+            {
+                EnableEvents(eventSource, EventLevel.LogAlways, (EventKeywords)0x80);
+            }
+            else if (eventSource.Name == "ActivityEventSource")
+            {
+                EnableEvents(eventSource, EventLevel.Informational);
+            }
+        }
+    }
+}

src/libraries/System.Diagnostics.Tracing/tests/System.Diagnostics.Tracing.Tests.csproj

@@ -17,6 +17,7 @@
     <Compile Include="BasicEventSourceTest\TestsManifestGeneration.Etw.cs" />
   </ItemGroup>
   <ItemGroup>
+    <Compile Include="BasicEventSourceTest\ActivityTracking.cs" />
     <Compile Include="BasicEventSourceTest\Harness\EventTestHarness.cs" />
     <Compile Include="BasicEventSourceTest\FuzzyTests.cs" />
     <Compile Include="BasicEventSourceTest\Harness\Listeners.cs" />

src/libraries/System.Private.CoreLib/src/System/Diagnostics/Tracing/ActivityTracker.cs

@@ -550,12 +550,36 @@ private static unsafe void WriteNibble(ref byte* ptr, byte* endPtr, uint value)
         // This callback is used to initialize the m_current AsyncLocal Variable.
         // Its job is to keep the ETW Activity ID (part of thread local storage) in sync
         // with m_current.ActivityID
+        //
+        // WARNING: When mixing manual usage of EventSource.SetCurrentThreadActivityID
+        // and Start/Stop EventSource events I can't identify a clear design how this
+        // synchronization is intended to work. For example code that changes
+        // SetCurrentThreadActivityID after a FooStart() event will not flow the
+        // explicit ID with the async work, but if FooStart() event is called after
+        // SetCurrentThreadActivityID then the explicit ID change does flow.
+        // For now I've adopted the approach:
+        // Priority 1: Make the API predictable/sensible when only Start/Stop events
+        // are in use.
+        // Priority 2: If users aren't complaining and it doesn't interfere with
+        // goal #1, try to preserve the arbitrary/buggy? existing behavior
+        // for mixed usage of SetActivityID + events.
+        //
+        // For scenarios that only use start/stop events this is what we expect:
+        // calling start -> push new ID on stack and update thread-local to match new ID
+        // calling stop -> pop ID from stack and update thread-local to match new topmost
+        //                 still active ID. If there is none, set ID to zero
+        // thread swap -> update thread-local to match topmost active ID.
+        //                 If there is none, set ID to zero.
         private void ActivityChanging(AsyncLocalValueChangedArgs<ActivityInfo?> args)
         {
             ActivityInfo? cur = args.CurrentValue;
             ActivityInfo? prev = args.PreviousValue;
 
-            // Are we popping off a value?   (we have a prev, and it creator is cur)
+            // Special case only relevant for mixed SetActivityID usage:
+            //
+            // Are we MAYBE popping off a value?   (we have a prev, and it creator is cur)
+            // We can't be certain this is a pop because a thread swapping between two
+            // ExecutionContexts can also appear the same way.
             // Then check if we should use the GUID at the time of the start event
             if (prev != null && prev.m_creator == cur)
             {
@@ -569,8 +593,7 @@ private void ActivityChanging(AsyncLocalValueChangedArgs<ActivityInfo?> args)
                 }
             }
 
-            // OK we did not have an explicit SetActivityID set.   Then we should be
-            // setting the activity to current ActivityInfo.  However that activity
+            // Set the activity to current ActivityInfo.  However that activity
             // might be dead, in which case we should skip it, so we never set
             // the ID to dead things.
             while (cur != null)
@@ -583,8 +606,10 @@ private void ActivityChanging(AsyncLocalValueChangedArgs<ActivityInfo?> args)
                 }
                 cur = cur.m_creator;
             }
+
             // we can get here if there is no information on our activity stack (everything is dead)
-            // currently we do nothing, as that seems better than setting to Guid.Emtpy.
+            // Set ActivityID to zero
+            EventSource.SetCurrentThreadActivityId(Guid.Empty);
         }
 
         /// <summary>