Add ActivityId Support to EventPipe (#16055)

9 files changed, 139 insertions, 21 deletions

diff --git a/src/mscorlib/System.Private.CoreLib.csproj b/src/mscorlib/System.Private.CoreLib.csproj
index a85e2988c7..df394edccb 100644
--- a/src/mscorlib/System.Private.CoreLib.csproj
+++ b/src/mscorlib/System.Private.CoreLib.csproj
@@ -518,7 +518,7 @@
     <Compile Condition="'$(FeatureXplatEventSource)' == 'true'" Include="$(BclSourcesRoot)\System\Diagnostics\Eventing\XplatEventLogger.cs" />
     <Compile Include="$(BclSourcesRoot)\System\Diagnostics\Eventing\FrameworkEventSource.cs" />
     <Compile Include="$(BclSourcesRoot)\System\Diagnostics\Eventing\EventPipe.cs" />
-    <Compile Condition="'$(FeaturePerfTracing)' == 'true'" Include="$(BclSourcesRoot)\System\Diagnostics\Eventing\EventPipeEventProvider.cs" />
+    <Compile Include="$(BclSourcesRoot)\System\Diagnostics\Eventing\EventPipeEventProvider.cs" />
   </ItemGroup>
   <ItemGroup>
     <Compile Include="$(BclSourcesRoot)\System\Diagnostics\Contracts\Contracts.cs" />
diff --git a/src/mscorlib/src/System/Diagnostics/Eventing/EventPipe.cs b/src/mscorlib/src/System/Diagnostics/Eventing/EventPipe.cs
index 5363b2aece..57c3b8f7c5 100644
--- a/src/mscorlib/src/System/Diagnostics/Eventing/EventPipe.cs
+++ b/src/mscorlib/src/System/Diagnostics/Eventing/EventPipe.cs
@@ -169,6 +169,9 @@ namespace System.Diagnostics.Tracing
         internal static extern void DeleteProvider(IntPtr provHandle);
 
         [DllImport(JitHelpers.QCall, CharSet = CharSet.Unicode)]
+        internal static extern int EventActivityIdControl(uint controlCode, ref Guid activityId);
+
+        [DllImport(JitHelpers.QCall, CharSet = CharSet.Unicode)]
         internal static extern unsafe void WriteEvent(IntPtr eventHandle, uint eventID, void* pData, uint length, Guid* activityId, Guid* relatedActivityId);
 
         [DllImport(JitHelpers.QCall, CharSet = CharSet.Unicode)]
diff --git a/src/mscorlib/src/System/Diagnostics/Eventing/EventPipeEventProvider.cs b/src/mscorlib/src/System/Diagnostics/Eventing/EventPipeEventProvider.cs
index a8789f5692..0d99ff44ec 100644
--- a/src/mscorlib/src/System/Diagnostics/Eventing/EventPipeEventProvider.cs
+++ b/src/mscorlib/src/System/Diagnostics/Eventing/EventPipeEventProvider.cs
@@ -74,7 +74,7 @@ namespace System.Diagnostics.Tracing
         // Get or set the per-thread activity ID.
         int IEventProvider.EventActivityIdControl(UnsafeNativeMethods.ManifestEtw.ActivityControl ControlCode, ref Guid ActivityId)
         {
-            return 0;
+            return EventPipeInternal.EventActivityIdControl((uint)ControlCode, ref ActivityId);
         }
 
         // Define an EventPipeEvent handle.
diff --git a/src/mscorlib/src/System/Diagnostics/Eventing/EventSource_CoreCLR.cs b/src/mscorlib/src/System/Diagnostics/Eventing/EventSource_CoreCLR.cs
index c5e2b20a80..98556a7aab 100644
--- a/src/mscorlib/src/System/Diagnostics/Eventing/EventSource_CoreCLR.cs
+++ b/src/mscorlib/src/System/Diagnostics/Eventing/EventSource_CoreCLR.cs
@@ -11,13 +11,6 @@ namespace System.Diagnostics.Tracing
 {
     public partial class EventSource
     {
-#if FEATURE_MANAGED_ETW && FEATURE_PERFTRACING
-        // For non-Windows, we use a thread-local variable to hold the activity ID.
-        // On Windows, ETW has it's own thread-local variable and we participate in its use.
-        [ThreadStatic]
-        private static Guid s_currentThreadActivityId;
-#endif // FEATURE_MANAGED_ETW && FEATURE_PERFTRACING
-
         // ActivityID support (see also WriteEventWithRelatedActivityIdCore)
         /// <summary>
         /// When a thread starts work that is on behalf of 'something else' (typically another 
@@ -47,13 +40,25 @@ namespace System.Diagnostics.Tracing
             // We ignore errors to keep with the convention that EventSources do not throw errors.
             // Note we can't access m_throwOnWrites because this is a static method.  
 
-#if FEATURE_PERFTRACING
-            s_currentThreadActivityId = activityId;
+#if FEATURE_PERFTRACING && PLATFORM_WINDOWS
+            // Set the activity id via EventPipe.
+            EventPipeInternal.EventActivityIdControl(
+                (uint)UnsafeNativeMethods.ManifestEtw.ActivityControl.EVENT_ACTIVITY_CTRL_SET_ID,
+                ref activityId);
+
+            // Set the activity id via ETW and fetch the previous id.
+            if (UnsafeNativeMethods.ManifestEtw.EventActivityIdControl(
+                UnsafeNativeMethods.ManifestEtw.ActivityControl.EVENT_ACTIVITY_CTRL_GET_SET_ID,
+                ref activityId) == 0)
+#elif FEATURE_PERFTRACING
+            if (EventPipeInternal.EventActivityIdControl(
+                (uint)UnsafeNativeMethods.ManifestEtw.ActivityControl.EVENT_ACTIVITY_CTRL_GET_SET_ID,
+                ref activityId) == 0)
 #elif PLATFORM_WINDOWS
             if (UnsafeNativeMethods.ManifestEtw.EventActivityIdControl(
                 UnsafeNativeMethods.ManifestEtw.ActivityControl.EVENT_ACTIVITY_CTRL_GET_SET_ID,
                 ref activityId) == 0)
-#endif // FEATURE_PERFTRACING
+#endif // FEATURE_PERFTRACING && PLATFORM_WINDOWS
             {
 #if FEATURE_ACTIVITYSAMPLING
                 var activityDying = s_activityDying;
@@ -97,14 +102,21 @@ namespace System.Diagnostics.Tracing
             // We ignore errors to keep with the convention that EventSources do not throw errors.
             // Note we can't access m_throwOnWrites because this is a static method.  
 
-#if FEATURE_PERFTRACING
-            oldActivityThatWillContinue = s_currentThreadActivityId;
-            s_currentThreadActivityId = activityId;
-#elif PLATFORM_WINDOWS
+#if FEATURE_PERFTRACING && PLATFORM_WINDOWS
+            EventPipeInternal.EventActivityIdControl(
+                (uint)UnsafeNativeMethods.ManifestEtw.ActivityControl.EVENT_ACTIVITY_CTRL_SET_ID,
+                    ref oldActivityThatWillContinue);
+#elif FEATURE_PERFTRACING
+            EventPipeInternal.EventActivityIdControl(
+                (uint)UnsafeNativeMethods.ManifestEtw.ActivityControl.EVENT_ACTIVITY_CTRL_GET_SET_ID,
+                    ref oldActivityThatWillContinue);
+#endif // FEATURE_PERFTRACING && PLATFORM_WINDOWS
+
+#if PLATFORM_WINDOWS
             UnsafeNativeMethods.ManifestEtw.EventActivityIdControl(
                 UnsafeNativeMethods.ManifestEtw.ActivityControl.EVENT_ACTIVITY_CTRL_GET_SET_ID,
                     ref oldActivityThatWillContinue);
-#endif // FEATURE_PERFTRACING
+#endif // PLATFORM_WINDOWS
 #endif // FEATURE_MANAGED_ETW
 
             // We don't call the activityDying callback here because the caller has declared that
@@ -124,13 +136,15 @@ namespace System.Diagnostics.Tracing
                 // errors. Note we can't access m_throwOnWrites because this is a static method.
                 Guid retVal = new Guid();
 #if FEATURE_MANAGED_ETW
-#if FEATURE_PERFTRACING
-                retVal = s_currentThreadActivityId;
-#elif PLATFORM_WINDOWS
+#if PLATFORM_WINDOWS
                 UnsafeNativeMethods.ManifestEtw.EventActivityIdControl(
                     UnsafeNativeMethods.ManifestEtw.ActivityControl.EVENT_ACTIVITY_CTRL_GET_ID,
                     ref retVal);
-#endif // FEATURE_PERFTRACING
+#elif FEATURE_PERFTRACING
+                EventPipeInternal.EventActivityIdControl(
+                    (uint)UnsafeNativeMethods.ManifestEtw.ActivityControl.EVENT_ACTIVITY_CTRL_GET_ID,
+                    ref retVal);
+#endif // PLATFORM_WINDOWS
 #endif // FEATURE_MANAGED_ETW
                 return retVal;
             }
diff --git a/src/vm/ecalllist.h b/src/vm/ecalllist.h
index 8fbfd206e5..b9706bb755 100644
--- a/src/vm/ecalllist.h
+++ b/src/vm/ecalllist.h
@@ -1203,6 +1203,7 @@ FCFuncStart(gEventPipeInternalFuncs)
     QCFuncElement("CreateProvider", EventPipeInternal::CreateProvider)
     QCFuncElement("DefineEvent", EventPipeInternal::DefineEvent)
     QCFuncElement("DeleteProvider", EventPipeInternal::DeleteProvider)
+    QCFuncElement("EventActivityIdControl", EventPipeInternal::EventActivityIdControl)
     QCFuncElement("WriteEvent", EventPipeInternal::WriteEvent)
     QCFuncElement("WriteEventData", EventPipeInternal::WriteEventData)
 FCFuncEnd()
diff --git a/src/vm/eventpipe.cpp b/src/vm/eventpipe.cpp
index 41af3efa03..2d7e3027bf 100644
--- a/src/vm/eventpipe.cpp
+++ b/src/vm/eventpipe.cpp
@@ -545,6 +545,12 @@ void EventPipe::WriteEventInternal(EventPipeEvent &event, EventPipeEventPayload
         return;
     }
 
+    // If the activity id isn't specified, pull it from the current thread.
+    if(pActivityId == NULL)
+    {
+        pActivityId = pThread->GetActivityId();
+    }
+
     if(!s_pConfig->RundownEnabled() && s_pBufferManager != NULL)
     {
         if(!s_pBufferManager->WriteEvent(pThread, *s_pSession, event, payload, pActivityId, pRelatedActivityId))
@@ -964,6 +970,67 @@ void QCALLTYPE EventPipeInternal::DeleteProvider(
     END_QCALL;
 }
 
+int QCALLTYPE EventPipeInternal::EventActivityIdControl(
+    uint controlCode,
+    GUID *pActivityId)
+{
+
+    QCALL_CONTRACT;
+
+    int retVal = 0;
+
+    BEGIN_QCALL;
+
+    Thread *pThread = GetThread();
+    if(pThread == NULL || pActivityId == NULL)
+    {
+        retVal = 1;
+    }
+    else
+    {
+        ActivityControlCode activityControlCode = (ActivityControlCode)controlCode;
+        GUID currentActivityId;
+        switch(activityControlCode)
+        {
+            case ActivityControlCode::EVENT_ACTIVITY_CONTROL_GET_ID:
+
+                *pActivityId = *pThread->GetActivityId();
+                break;
+
+            case ActivityControlCode::EVENT_ACTIVITY_CONTROL_SET_ID:
+
+                pThread->SetActivityId(pActivityId);
+                break;
+
+            case ActivityControlCode::EVENT_ACTIVITY_CONTROL_CREATE_ID:
+
+                CoCreateGuid(pActivityId);
+                break;
+
+            case ActivityControlCode::EVENT_ACTIVITY_CONTROL_GET_SET_ID:
+
+                currentActivityId = *pThread->GetActivityId();
+                pThread->SetActivityId(pActivityId);
+                *pActivityId = currentActivityId;
+
+                break;
+
+            case ActivityControlCode::EVENT_ACTIVITY_CONTROL_CREATE_SET_ID:
+
+                *pActivityId = *pThread->GetActivityId();
+                CoCreateGuid(&currentActivityId);
+                pThread->SetActivityId(&currentActivityId);
+                break;
+
+            default:
+                retVal = 1;
+        };
+    }
+
+    END_QCALL;
+    return retVal;
+}
+
 void QCALLTYPE EventPipeInternal::WriteEvent(
     INT_PTR eventHandle,
     UINT32 eventID,
diff --git a/src/vm/eventpipe.h b/src/vm/eventpipe.h
index d1f7d60f8e..907e0d7778 100644
--- a/src/vm/eventpipe.h
+++ b/src/vm/eventpipe.h
@@ -356,6 +356,16 @@ public:
 
 class EventPipeInternal
 {
+private:
+
+    enum class ActivityControlCode
+    {
+        EVENT_ACTIVITY_CONTROL_GET_ID = 1,
+        EVENT_ACTIVITY_CONTROL_SET_ID = 2,
+        EVENT_ACTIVITY_CONTROL_CREATE_ID = 3,
+        EVENT_ACTIVITY_CONTROL_GET_SET_ID = 4,
+        EVENT_ACTIVITY_CONTROL_CREATE_SET_ID = 5
+    };
 
 public:
 
@@ -384,6 +394,10 @@ public:
     static void QCALLTYPE DeleteProvider(
         INT_PTR provHandle);
 
+    static int QCALLTYPE EventActivityIdControl(
+        uint controlCode,
+        GUID *pActivityId);
+
     static void QCALLTYPE WriteEvent(
         INT_PTR eventHandle,
         UINT32 eventID,
diff --git a/src/vm/threads.cpp b/src/vm/threads.cpp
index 54821022b7..422789b23d 100644
--- a/src/vm/threads.cpp
+++ b/src/vm/threads.cpp
@@ -1687,6 +1687,7 @@ Thread::Thread()
 #ifdef FEATURE_PERFTRACING
     m_pEventPipeBufferList = NULL;
     m_eventWriteInProgress = false;
+    memset(&m_activityId, 0, sizeof(m_activityId));
 #endif // FEATURE_PERFTRACING
     m_HijackReturnKind = RT_Illegal;
 }
diff --git a/src/vm/threads.h b/src/vm/threads.h
index 57cb872d28..d55d1f8323 100644
--- a/src/vm/threads.h
+++ b/src/vm/threads.h
@@ -5256,6 +5256,10 @@ private:
     // True if the thread was in cooperative mode.  False if it was in preemptive when the suspension started.
     Volatile<ULONG> m_gcModeOnSuspension;
 
+    // The activity ID for the current thread.
+    // An activity ID of zero means the thread is not executing in the context of an activity.
+    GUID m_activityId;
+
 public:
     EventPipeBufferList* GetEventPipeBufferList()
     {
@@ -5297,6 +5301,20 @@ public:
     {
         m_gcModeOnSuspension = 0;
     }
+
+    LPCGUID GetActivityId() const
+    {
+        LIMITED_METHOD_CONTRACT;
+        return &m_activityId;
+    }
+
+    void SetActivityId(LPCGUID pActivityId)
+    {
+        LIMITED_METHOD_CONTRACT;
+        _ASSERTE(pActivityId != NULL);
+
+        m_activityId = *pActivityId;
+    }
 #endif // FEATURE_PERFTRACING
 
 #ifdef FEATURE_HIJACK