Implement Activity ID / Related Activity ID for EventSources Writing to EventPipe (#11904)

* Implement thread-local storage of activity ids.

* Plumb the activity id and related activity id into the trace file.

12 files changed, 98 insertions, 34 deletions

diff --git a/src/mscorlib/src/System/Diagnostics/Eventing/EventPipe.cs b/src/mscorlib/src/System/Diagnostics/Eventing/EventPipe.cs
index 0d66c94e8d..2f6fdf62ef 100644
--- a/src/mscorlib/src/System/Diagnostics/Eventing/EventPipe.cs
+++ b/src/mscorlib/src/System/Diagnostics/Eventing/EventPipe.cs
@@ -170,6 +170,6 @@ namespace System.Diagnostics.Tracing
 
         [DllImport(JitHelpers.QCall, CharSet = CharSet.Unicode)]
         [SuppressUnmanagedCodeSecurity]
-        internal static extern unsafe void WriteEvent(IntPtr eventHandle, uint eventID, void* pData, uint length);
+        internal static extern unsafe void WriteEvent(IntPtr eventHandle, uint eventID, void* pData, uint length, Guid* activityId, Guid* relatedActivityId);
     }
 }
diff --git a/src/mscorlib/src/System/Diagnostics/Eventing/EventPipeEventProvider.cs b/src/mscorlib/src/System/Diagnostics/Eventing/EventPipeEventProvider.cs
index cd9bd3c713..d5bc4c2889 100644
--- a/src/mscorlib/src/System/Diagnostics/Eventing/EventPipeEventProvider.cs
+++ b/src/mscorlib/src/System/Diagnostics/Eventing/EventPipeEventProvider.cs
@@ -62,7 +62,7 @@ namespace System.Diagnostics.Tracing
             {
                 if (userDataCount == 0)
                 {
-                    EventPipeInternal.WriteEvent(eventHandle, eventID, null, 0);
+                    EventPipeInternal.WriteEvent(eventHandle, eventID, null, 0, activityId, relatedActivityId);
                     return 0;
                 }
 
@@ -82,7 +82,7 @@ namespace System.Diagnostics.Tracing
                         uint singleUserDataSize = userData[i].Size;
                         WriteToBuffer(pData, length, ref offset, singleUserDataPtr, singleUserDataSize);
                     }
-                    EventPipeInternal.WriteEvent(eventHandle, eventID, pData, length);
+                    EventPipeInternal.WriteEvent(eventHandle, eventID, pData, length, activityId, relatedActivityId);
                 }
             }
             return 0;
diff --git a/src/mscorlib/src/System/Diagnostics/Eventing/EventSource_CoreCLR.cs b/src/mscorlib/src/System/Diagnostics/Eventing/EventSource_CoreCLR.cs
index a41c069699..01aac72cf6 100644
--- a/src/mscorlib/src/System/Diagnostics/Eventing/EventSource_CoreCLR.cs
+++ b/src/mscorlib/src/System/Diagnostics/Eventing/EventSource_CoreCLR.cs
@@ -11,6 +11,13 @@ 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 
@@ -33,16 +40,20 @@ namespace System.Diagnostics.Tracing
         {
             if (TplEtwProvider.Log != null)
                 TplEtwProvider.Log.SetActivityId(activityId);
-#if FEATURE_MANAGED_ETW && PLATFORM_WINDOWS
+#if FEATURE_MANAGED_ETW
 #if FEATURE_ACTIVITYSAMPLING
             Guid newId = activityId;
 #endif // FEATURE_ACTIVITYSAMPLING
             // 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;
+#elif PLATFORM_WINDOWS
             if (UnsafeNativeMethods.ManifestEtw.EventActivityIdControl(
                 UnsafeNativeMethods.ManifestEtw.ActivityControl.EVENT_ACTIVITY_CTRL_GET_SET_ID,
                 ref activityId) == 0)
+#endif // FEATURE_PERFTRACING
             {
 #if FEATURE_ACTIVITYSAMPLING
                 var activityDying = s_activityDying;
@@ -57,7 +68,7 @@ namespace System.Diagnostics.Tracing
                 }
 #endif // FEATURE_ACTIVITYSAMPLING
             }
-#endif // FEATURE_MANAGED_ETW && PLATFORM_WINDOWS
+#endif // FEATURE_MANAGED_ETW
         }
 
         /// <summary>
@@ -82,14 +93,19 @@ namespace System.Diagnostics.Tracing
         public static void SetCurrentThreadActivityId(Guid activityId, out Guid oldActivityThatWillContinue)
         {
             oldActivityThatWillContinue = activityId;
-#if FEATURE_MANAGED_ETW && PLATFORM_WINDOWS
+#if FEATURE_MANAGED_ETW
             // 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
             UnsafeNativeMethods.ManifestEtw.EventActivityIdControl(
                 UnsafeNativeMethods.ManifestEtw.ActivityControl.EVENT_ACTIVITY_CTRL_GET_SET_ID,
                     ref oldActivityThatWillContinue);
-#endif // FEATURE_MANAGED_ETW && PLATFORM_WINDOWS
+#endif // FEATURE_PERFTRACING
+#endif // FEATURE_MANAGED_ETW
 
             // We don't call the activityDying callback here because the caller has declared that
             // it is not dying.  
@@ -107,11 +123,15 @@ 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.
                 Guid retVal = new Guid();
-#if FEATURE_MANAGED_ETW && PLATFORM_WINDOWS
+#if FEATURE_MANAGED_ETW
+#if FEATURE_PERFTRACING
+                retVal = s_currentThreadActivityId;
+#elif PLATFORM_WINDOWS
                 UnsafeNativeMethods.ManifestEtw.EventActivityIdControl(
                     UnsafeNativeMethods.ManifestEtw.ActivityControl.EVENT_ACTIVITY_CTRL_GET_ID,
                     ref retVal);
-#endif // FEATURE_MANAGED_ETW && PLATFORM_WINDOWS
+#endif // FEATURE_PERFTRACING
+#endif // FEATURE_MANAGED_ETW
                 return retVal;
             }
         }
diff --git a/src/vm/eventpipe.cpp b/src/vm/eventpipe.cpp
index b73c023b64..4eae6a87c6 100644
--- a/src/vm/eventpipe.cpp
+++ b/src/vm/eventpipe.cpp
@@ -282,7 +282,7 @@ void EventPipe::DeleteProvider(EventPipeProvider *pProvider)
     }
 }
 
-void EventPipe::WriteEvent(EventPipeEvent &event, BYTE *pData, unsigned int length)
+void EventPipe::WriteEvent(EventPipeEvent &event, BYTE *pData, unsigned int length, LPCGUID pActivityId, LPCGUID pRelatedActivityId)
 {
     CONTRACTL
     {
@@ -309,7 +309,7 @@ void EventPipe::WriteEvent(EventPipeEvent &event, BYTE *pData, unsigned int leng
 
     if(!s_pConfig->RundownEnabled() && s_pBufferManager != NULL)
     {
-        if(!s_pBufferManager->WriteEvent(pThread, event, pData, length))
+        if(!s_pBufferManager->WriteEvent(pThread, event, pData, length, pActivityId, pRelatedActivityId))
         {
             // This is used in DEBUG to make sure that we don't log an event synchronously that we didn't log to the buffer.
             return;
@@ -323,7 +323,9 @@ void EventPipe::WriteEvent(EventPipeEvent &event, BYTE *pData, unsigned int leng
             event,
             pThread->GetOSThreadId(),
             pData,
-            length);
+            length,
+            pActivityId,
+            pRelatedActivityId);
 
         if(s_pFile != NULL)
         {
@@ -340,7 +342,9 @@ void EventPipe::WriteEvent(EventPipeEvent &event, BYTE *pData, unsigned int leng
             event,
             pThread->GetOSThreadId(),
             pData,
-            length);
+            length,
+            pActivityId,
+            pRelatedActivityId);
 
         // Write to the EventPipeFile if it exists.
         if(s_pSyncFile != NULL)
@@ -372,7 +376,7 @@ void EventPipe::WriteSampleProfileEvent(Thread *pSamplingThread, EventPipeEvent
     {
         // Specify the sampling thread as the "current thread", so that we select the right buffer.
         // Specify the target thread so that the event gets properly attributed.
-        if(!s_pBufferManager->WriteEvent(pSamplingThread, *pEvent, pData, length, pTargetThread, &stackContents))
+        if(!s_pBufferManager->WriteEvent(pSamplingThread, *pEvent, pData, length, NULL /* pActivityId */, NULL /* pRelatedActivityId */, pTargetThread, &stackContents))
         {
             // This is used in DEBUG to make sure that we don't log an event synchronously that we didn't log to the buffer.
             return;
@@ -577,14 +581,16 @@ void QCALLTYPE EventPipeInternal::WriteEvent(
     INT_PTR eventHandle,
     unsigned int eventID,
     void *pData,
-    unsigned int length)
+    unsigned int length,
+    LPCGUID pActivityId,
+    LPCGUID pRelatedActivityId)
 {
     QCALL_CONTRACT;
     BEGIN_QCALL;
 
     _ASSERTE(eventHandle != NULL);
     EventPipeEvent *pEvent = reinterpret_cast<EventPipeEvent *>(eventHandle);
-    EventPipe::WriteEvent(*pEvent, (BYTE *)pData, length);
+    EventPipe::WriteEvent(*pEvent, (BYTE *)pData, length, pActivityId, pRelatedActivityId);
 
     END_QCALL;
 }
diff --git a/src/vm/eventpipe.h b/src/vm/eventpipe.h
index 0860080d48..a69e1ba840 100644
--- a/src/vm/eventpipe.h
+++ b/src/vm/eventpipe.h
@@ -188,7 +188,7 @@ class EventPipe
 
         // Write out an event.
         // Data is written as a serialized blob matching the ETW serialization conventions.
-        static void WriteEvent(EventPipeEvent &event, BYTE *pData, unsigned int length);
+        static void WriteEvent(EventPipeEvent &event, BYTE *pData, unsigned int length, LPCGUID pActivityId = NULL, LPCGUID pRelatedActivityId = NULL);
 
         // Write out a sample profile event.
         static void WriteSampleProfileEvent(Thread *pSamplingThread, EventPipeEvent *pEvent, Thread *pTargetThread, StackContents &stackContents, BYTE *pData = NULL, unsigned int length = 0);
@@ -305,7 +305,8 @@ public:
         INT_PTR eventHandle,
         unsigned int eventID,
         void *pData,
-        unsigned int length);
+        unsigned int length,
+        LPCGUID pActivityId, LPCGUID pRelatedActivityId);
 };
 
 #endif // FEATURE_PERFTRACING
diff --git a/src/vm/eventpipebuffer.cpp b/src/vm/eventpipebuffer.cpp
index ed1c5479bb..00652c9fac 100644
--- a/src/vm/eventpipebuffer.cpp
+++ b/src/vm/eventpipebuffer.cpp
@@ -46,7 +46,7 @@ EventPipeBuffer::~EventPipeBuffer()
     }
 }
 
-bool EventPipeBuffer::WriteEvent(Thread *pThread, EventPipeEvent &event, BYTE *pData, unsigned int dataLength, StackContents *pStack)
+bool EventPipeBuffer::WriteEvent(Thread *pThread, EventPipeEvent &event, BYTE *pData, unsigned int dataLength, LPCGUID pActivityId, LPCGUID pRelatedActivityId, StackContents *pStack)
 {
     CONTRACTL
     {
@@ -77,7 +77,9 @@ bool EventPipeBuffer::WriteEvent(Thread *pThread, EventPipeEvent &event, BYTE *p
             event,
             pThread->GetOSThreadId(),
             pDataDest,
-            dataLength);
+            dataLength,
+            pActivityId,
+            pRelatedActivityId);
 
         // Copy the stack if a separate stack trace was provided.
         if(pStack != NULL)
diff --git a/src/vm/eventpipebuffer.h b/src/vm/eventpipebuffer.h
index 97b858d018..f279a2865c 100644
--- a/src/vm/eventpipebuffer.h
+++ b/src/vm/eventpipebuffer.h
@@ -81,7 +81,7 @@ public:
     // Returns:
     //  - true: The write succeeded.
     //  - false: The write failed.  In this case, the buffer should be considered full.
-    bool WriteEvent(Thread *pThread, EventPipeEvent &event, BYTE *pData, unsigned int dataLength, StackContents *pStack = NULL);
+    bool WriteEvent(Thread *pThread, EventPipeEvent &event, BYTE *pData, unsigned int dataLength, LPCGUID pActivityId, LPCGUID pRelatedActivityId, StackContents *pStack = NULL);
 
     // Get the timestamp of the most recent event in the buffer.
     LARGE_INTEGER GetMostRecentTimeStamp() const;
diff --git a/src/vm/eventpipebuffermanager.cpp b/src/vm/eventpipebuffermanager.cpp
index e271fcef88..5edc462542 100644
--- a/src/vm/eventpipebuffermanager.cpp
+++ b/src/vm/eventpipebuffermanager.cpp
@@ -217,7 +217,7 @@ void EventPipeBufferManager::DeAllocateBuffer(EventPipeBuffer *pBuffer)
     }
 }
 
-bool EventPipeBufferManager::WriteEvent(Thread *pThread, EventPipeEvent &event, BYTE *pData, unsigned int length, Thread *pEventThread, StackContents *pStack)
+bool EventPipeBufferManager::WriteEvent(Thread *pThread, EventPipeEvent &event, BYTE *pData, unsigned int length, LPCGUID pActivityId, LPCGUID pRelatedActivityId, Thread *pEventThread, StackContents *pStack)
 {
     CONTRACTL
     {
@@ -276,7 +276,7 @@ bool EventPipeBufferManager::WriteEvent(Thread *pThread, EventPipeEvent &event,
         else
         {
             // Attempt to write the event to the buffer.  If this fails, we should allocate a new buffer.
-            allocNewBuffer = !pBuffer->WriteEvent(pEventThread, event, pData, length, pStack);
+            allocNewBuffer = !pBuffer->WriteEvent(pEventThread, event, pData, length, pActivityId, pRelatedActivityId, pStack);
         }
     }
 
@@ -294,7 +294,7 @@ bool EventPipeBufferManager::WriteEvent(Thread *pThread, EventPipeEvent &event,
     // This is the second time if this thread did have one or more buffers, but they were full.
     if(allocNewBuffer && pBuffer != NULL)
     {
-        allocNewBuffer = !pBuffer->WriteEvent(pEventThread, event, pData, length, pStack);
+        allocNewBuffer = !pBuffer->WriteEvent(pEventThread, event, pData, length, pActivityId, pRelatedActivityId, pStack);
     }
 
     // Mark that the thread is no longer writing an event.
diff --git a/src/vm/eventpipebuffermanager.h b/src/vm/eventpipebuffermanager.h
index 74783d26e8..a53721b7b8 100644
--- a/src/vm/eventpipebuffermanager.h
+++ b/src/vm/eventpipebuffermanager.h
@@ -67,7 +67,7 @@ public:
     // This is because the thread that writes the events is not the same as the "event thread".
     // An optional stack trace can be provided for sample profiler events.
     // Otherwise, if a stack trace is needed, one will be automatically collected.
-    bool WriteEvent(Thread *pThread, EventPipeEvent &event, BYTE *pData, unsigned int length, Thread *pEventThread = NULL, StackContents *pStack = NULL);
+    bool WriteEvent(Thread *pThread, EventPipeEvent &event, BYTE *pData, unsigned int length, LPCGUID pActivityId, LPCGUID pRelatedActivityId, Thread *pEventThread = NULL, StackContents *pStack = NULL);
 
     // Write the contents of the managed buffers to the specified file.
     // The stopTimeStamp is used to determine when tracing was stopped to ensure that we
diff --git a/src/vm/eventpipeconfiguration.cpp b/src/vm/eventpipeconfiguration.cpp
index 73379fed1b..42f9daf528 100644
--- a/src/vm/eventpipeconfiguration.cpp
+++ b/src/vm/eventpipeconfiguration.cpp
@@ -370,7 +370,9 @@ EventPipeEventInstance* EventPipeConfiguration::BuildEventMetadataEvent(EventPip
         *m_pMetadataEvent,
         GetCurrentThreadId(),
         pInstancePayload,
-        instancePayloadSize);
+        instancePayloadSize,
+        NULL /* pActivityId */,
+        NULL /* pRelatedActivityId */);
 
     // Set the timestamp to match the source event, because the metadata event
     // will be emitted right before the source event.
diff --git a/src/vm/eventpipeeventinstance.cpp b/src/vm/eventpipeeventinstance.cpp
index 9372cece65..afde2c0547 100644
--- a/src/vm/eventpipeeventinstance.cpp
+++ b/src/vm/eventpipeeventinstance.cpp
@@ -14,7 +14,9 @@ EventPipeEventInstance::EventPipeEventInstance(
     EventPipeEvent &event,
     DWORD threadID,
     BYTE *pData,
-    unsigned int length)
+    unsigned int length,
+    LPCGUID pActivityId,
+    LPCGUID pRelatedActivityId)
 {
     CONTRACTL
     {
@@ -30,6 +32,23 @@ EventPipeEventInstance::EventPipeEventInstance(
 #endif // _DEBUG
     m_pEvent = &event;
     m_threadID = threadID;
+    if(pActivityId != NULL)
+    {
+        m_activityId = *pActivityId;
+    }
+    else
+    {
+        m_activityId = {0};
+    }
+    if(pRelatedActivityId != NULL)
+    {
+        m_relatedActivityId = *pRelatedActivityId;
+    }
+    else
+    {
+        m_relatedActivityId = {0};
+    }
+
     m_pData = pData;
     m_dataLength = length;
     QueryPerformanceCounter(&m_timeStamp);
@@ -98,11 +117,14 @@ void EventPipeEventInstance::FastSerialize(FastSerializer *pSerializer, StreamLa
     // Calculate the size of the total payload so that it can be written to the file.
     unsigned int payloadLength =
         sizeof(metadataLabel) +
-        sizeof(m_threadID) +        // Thread ID
-        sizeof(m_timeStamp) +       // TimeStamp
-        m_dataLength +              // Event payload data length
-        sizeof(unsigned int) +      // Prepended stack payload size in bytes
-        m_stackContents.GetSize();  // Stack payload size
+        sizeof(m_threadID) +            // Thread ID
+        sizeof(m_timeStamp) +           // TimeStamp
+        sizeof(m_activityId) +          // Activity ID
+        sizeof(m_relatedActivityId) +   // Related Activity ID
+        sizeof(m_dataLength) +          // Data payload length
+        m_dataLength +                  // Event payload data
+        sizeof(unsigned int) +          // Prepended stack payload size in bytes
+        m_stackContents.GetSize();      // Stack payload size
 
     // Write the size of the event to the file.
     pSerializer->WriteBuffer((BYTE*)&payloadLength, sizeof(payloadLength));
@@ -116,6 +138,15 @@ void EventPipeEventInstance::FastSerialize(FastSerializer *pSerializer, StreamLa
     // Write the timestamp.
     pSerializer->WriteBuffer((BYTE*)&m_timeStamp, sizeof(m_timeStamp));
 
+    // Write the activity id.
+    pSerializer->WriteBuffer((BYTE*)&m_activityId, sizeof(m_activityId));
+
+    // Write the related activity id.
+    pSerializer->WriteBuffer((BYTE*)&m_relatedActivityId, sizeof(m_relatedActivityId));
+
+    // Write the data payload size.
+    pSerializer->WriteBuffer((BYTE*)&m_dataLength, sizeof(m_dataLength));
+
     // Write the event data payload.
     if(m_dataLength > 0)
     {
@@ -199,7 +230,7 @@ bool EventPipeEventInstance::EnsureConsistency()
 #endif // _DEBUG
 
 SampleProfilerEventInstance::SampleProfilerEventInstance(EventPipeEvent &event, Thread *pThread, BYTE *pData, unsigned int length)
-    :EventPipeEventInstance(event, pThread->GetOSThreadId(), pData, length)
+    :EventPipeEventInstance(event, pThread->GetOSThreadId(), pData, length, NULL /* pActivityId */, NULL /* pRelatedActivityId */)
 {
     LIMITED_METHOD_CONTRACT;
 }
diff --git a/src/vm/eventpipeeventinstance.h b/src/vm/eventpipeeventinstance.h
index f54a894066..4fcf95c73e 100644
--- a/src/vm/eventpipeeventinstance.h
+++ b/src/vm/eventpipeeventinstance.h
@@ -19,7 +19,7 @@ class EventPipeEventInstance
 
 public:
 
-    EventPipeEventInstance(EventPipeEvent &event, DWORD threadID, BYTE *pData, unsigned int length);
+    EventPipeEventInstance(EventPipeEvent &event, DWORD threadID, BYTE *pData, unsigned int length, LPCGUID pActivityId, LPCGUID pRelatedActivityId);
 
     // Get the event associated with this instance.
     EventPipeEvent* GetEvent() const;
@@ -55,6 +55,8 @@ protected:
     EventPipeEvent *m_pEvent;
     DWORD m_threadID;
     LARGE_INTEGER m_timeStamp;
+    GUID m_activityId;
+    GUID m_relatedActivityId;
 
     BYTE *m_pData;
     unsigned int m_dataLength;