Randomized allocation sampling

src/coreclr/inc/eventtracebase.h

@@ -1331,17 +1331,19 @@ namespace ETW
 #define ETWLoaderStaticLoad 0 // Static reference load
 #define ETWLoaderDynamicLoad 1 // Dynamic assembly load
 
+#if defined (FEATURE_EVENT_TRACE)
+EXTERN_C DOTNET_TRACE_CONTEXT MICROSOFT_WINDOWS_DOTNETRUNTIME_PROVIDER_DOTNET_Context;
+EXTERN_C DOTNET_TRACE_CONTEXT MICROSOFT_WINDOWS_DOTNETRUNTIME_PRIVATE_PROVIDER_DOTNET_Context;
+EXTERN_C DOTNET_TRACE_CONTEXT MICROSOFT_WINDOWS_DOTNETRUNTIME_RUNDOWN_PROVIDER_DOTNET_Context;
+EXTERN_C DOTNET_TRACE_CONTEXT MICROSOFT_WINDOWS_DOTNETRUNTIME_STRESS_PROVIDER_DOTNET_Context;
+#endif // FEATURE_EVENT_TRACE
+
 #if defined(FEATURE_EVENT_TRACE) && !defined(HOST_UNIX)
 //
 // The ONE and only ONE global instantiation of this class
 //
 extern ETW::CEtwTracer *  g_pEtwTracer;
 
-EXTERN_C DOTNET_TRACE_CONTEXT MICROSOFT_WINDOWS_DOTNETRUNTIME_PROVIDER_DOTNET_Context;
-EXTERN_C DOTNET_TRACE_CONTEXT MICROSOFT_WINDOWS_DOTNETRUNTIME_PRIVATE_PROVIDER_DOTNET_Context;
-EXTERN_C DOTNET_TRACE_CONTEXT MICROSOFT_WINDOWS_DOTNETRUNTIME_RUNDOWN_PROVIDER_DOTNET_Context;
-EXTERN_C DOTNET_TRACE_CONTEXT MICROSOFT_WINDOWS_DOTNETRUNTIME_STRESS_PROVIDER_DOTNET_Context;
-
 //
 // Special Handling of Startup events
 //

src/coreclr/nativeaot/Runtime/CMakeLists.txt

@@ -49,6 +49,7 @@ set(COMMON_RUNTIME_SOURCES
     ${GC_DIR}/handletablescan.cpp
     ${GC_DIR}/objecthandle.cpp
     ${GC_DIR}/softwarewritewatch.cpp
+    ${CLR_SRC_NATIVE_DIR}/minipal/xoshiro128pp.c
 )
 
 set(SERVER_GC_SOURCES

src/coreclr/nativeaot/Runtime/GCHelpers.cpp

@@ -29,6 +29,12 @@
 
 #include "gcdesc.h"
 
+#ifdef FEATURE_EVENT_TRACE
+   #include "clretwallmain.h"
+#else // FEATURE_EVENT_TRACE
+   #include "etmdummy.h"
+#endif // FEATURE_EVENT_TRACE
+
 #define RH_LARGE_OBJECT_SIZE 85000
 
 MethodTable g_FreeObjectEEType;
@@ -471,6 +477,24 @@ EXTERN_C int64_t QCALLTYPE RhGetTotalAllocatedBytesPrecise()
     return allocated;
 }
 
+void FireAllocationSampled(GC_ALLOC_FLAGS flags, size_t size, size_t samplingBudgetOffset, Object* orObject)
+{
+#ifdef FEATURE_EVENT_TRACE
+    void* typeId = GetLastAllocEEType();
+    // Note: Just as for AllocationTick, the type name cannot be retrieved
+    WCHAR* name = nullptr;
+
+    if (typeId != nullptr)
+    {
+        unsigned int allocKind =
+            (flags & GC_ALLOC_PINNED_OBJECT_HEAP) ? 2 :
+            (flags & GC_ALLOC_LARGE_OBJECT_HEAP) ? 1 :
+            0;  // SOH
+        FireEtwAllocationSampled(allocKind, GetClrInstanceId(), typeId, name, (BYTE*)orObject, size, samplingBudgetOffset);
+    }
+#endif
+}
+
 static Object* GcAllocInternal(MethodTable* pEEType, uint32_t uFlags, uintptr_t numElements, Thread* pThread)
 {
     ASSERT(!pThread->IsDoNotTriggerGcSet());
@@ -539,8 +563,47 @@ static Object* GcAllocInternal(MethodTable* pEEType, uint32_t uFlags, uintptr_t
     // Save the MethodTable for instrumentation purposes.
     tls_pLastAllocationEEType = pEEType;
 
-    Object* pObject = GCHeapUtilities::GetGCHeap()->Alloc(pThread->GetAllocContext(), cbSize, uFlags);
-    pThread->GetEEAllocContext()->UpdateCombinedLimit();
+    // check for dynamic allocation sampling
+    ee_alloc_context* pEEAllocContext = pThread->GetEEAllocContext();
+    gc_alloc_context* pAllocContext = pEEAllocContext->GetGCAllocContext();
+    bool isSampled = false;
+    size_t availableSpace = 0;
+    size_t samplingBudget = 0;
+
+    bool isRandomizedSamplingEnabled = ee_alloc_context::IsRandomizedSamplingEnabled();
+    if (isRandomizedSamplingEnabled)
+    {
+        // The number bytes we can allocate before we need to emit a sampling event.
+        // This calculation is only valid if combined_limit < alloc_limit.
+        samplingBudget = (size_t)(pEEAllocContext->combined_limit - pAllocContext->alloc_ptr);
+
+        // The number of bytes available in the current allocation context
+        availableSpace = (size_t)(pAllocContext->alloc_limit - pAllocContext->alloc_ptr);
+
+        // Check to see if the allocated object overlaps a sampled byte
+        // in this AC. This happens when both:
+        // 1) The AC contains a sampled byte (combined_limit < alloc_limit)
+        // 2) The object is large enough to overlap it (samplingBudget < aligned_size)
+        //
+        // Note that the AC could have no remaining space for allocations (alloc_ptr =
+        // alloc_limit = combined_limit). When a thread hasn't done any SOH allocations
+        // yet it also starts in an empty state where alloc_ptr = alloc_limit =
+        // combined_limit = nullptr. The (1) check handles both of these situations
+        // properly as an empty AC can not have a sampled byte inside of it.
+        isSampled =
+            (pEEAllocContext->combined_limit < pAllocContext->alloc_limit) &&
+            (samplingBudget < cbSize);
+
+        // if the object overflows the AC, we need to sample the remaining bytes
+        // the sampling budget only included at most the bytes inside the AC
+        if (cbSize > availableSpace && !isSampled)
+        {
+            samplingBudget = ee_alloc_context::ComputeGeometricRandom() + availableSpace;
+            isSampled = (samplingBudget < cbSize);
+        }
+    }
+
+    Object* pObject = GCHeapUtilities::GetGCHeap()->Alloc(pAllocContext, cbSize, uFlags);
     if (pObject == NULL)
         return NULL;
 
@@ -551,6 +614,19 @@ static Object* GcAllocInternal(MethodTable* pEEType, uint32_t uFlags, uintptr_t
         ((Array*)pObject)->InitArrayLength((uint32_t)numElements);
     }
 
+    if (isSampled)
+    {
+        FireAllocationSampled((GC_ALLOC_FLAGS)uFlags, cbSize, samplingBudget, pObject);
+    }
+
+    // There are a variety of conditions that may have invalidated the previous combined_limit value
+    // such as not allocating the object in the AC memory region (UOH allocations), moving the AC, adding
+    // extra alignment padding, allocating a new AC, or allocating an object that consumed the sampling budget.
+    // Rather than test for all the different invalidation conditions individually we conservatively always
+    // recompute it. If sampling isn't enabled this inlined function is just trivially setting
+    // combined_limit=alloc_limit.
+    pEEAllocContext->UpdateCombinedLimit(isRandomizedSamplingEnabled);
+
     if (uFlags & GC_ALLOC_USER_OLD_HEAP)
         GCHeapUtilities::GetGCHeap()->PublishObject((uint8_t*)pObject);
 

src/coreclr/nativeaot/Runtime/disabledeventtrace.cpp

@@ -12,6 +12,11 @@
 
 void EventTracing_Initialize() { }
 
+bool IsRuntimeProviderEnabled(uint8_t level, uint64_t keyword)
+{
+    return false;
+}
+
 void ETW::GCLog::FireGcStart(ETW_GC_INFO * pGcInfo) { }
 
 #ifdef FEATURE_ETW

src/coreclr/nativeaot/Runtime/eventpipe/gen-eventing-event-inc.lst

@@ -1,5 +1,6 @@
 # Native runtime events supported by aot runtime.
 
+AllocationSampled
 BGC1stConEnd
 BGC1stNonConEnd
 BGC1stSweepEnd

src/coreclr/nativeaot/Runtime/eventtrace.cpp

@@ -37,6 +37,11 @@ DOTNET_TRACE_CONTEXT MICROSOFT_WINDOWS_DOTNETRUNTIME_PRIVATE_PROVIDER_DOTNET_Con
     MICROSOFT_WINDOWS_DOTNETRUNTIME_PRIVATE_PROVIDER_EVENTPIPE_Context
 };
 
+bool IsRuntimeProviderEnabled(uint8_t level, uint64_t keyword)
+{
+    return RUNTIME_PROVIDER_CATEGORY_ENABLED(level, keyword);
+}
+
 volatile LONGLONG ETW::GCLog::s_l64LastClientSequenceNumber = 0;
 
 //---------------------------------------------------------------------------------------
@@ -245,4 +250,4 @@ void EventPipeEtwCallbackDotNETRuntimePrivate(
     _Inout_opt_ PVOID CallbackContext)
 {
     EtwCallbackCommon(DotNETRuntimePrivate, ControlCode, Level, MatchAnyKeyword, FilterData, true);
-}
+}

src/coreclr/nativeaot/Runtime/eventtracebase.h

@@ -30,6 +30,8 @@ void InitializeEventTracing();
 
 #ifdef FEATURE_EVENT_TRACE
 
+bool IsRuntimeProviderEnabled(uint8_t level, uint64_t keyword);
+
 // !!!!!!! NOTE !!!!!!!!
 // The flags must match those in the ETW manifest exactly
 // !!!!!!! NOTE !!!!!!!!
@@ -102,6 +104,7 @@ struct ProfilingScanContext;
 #define CLR_GCHEAPSURVIVALANDMOVEMENT_KEYWORD 0x400000
 #define CLR_MANAGEDHEAPCOLLECT_KEYWORD 0x800000
 #define CLR_GCHEAPANDTYPENAMES_KEYWORD 0x1000000
+#define CLR_ALLOCATIONSAMPLING_KEYWORD 0x80000000000
 
 //
 // Using KEYWORDZERO means when checking the events category ignore the keyword

src/coreclr/nativeaot/Runtime/gctoclreventsink.cpp

@@ -4,6 +4,7 @@
 #include "common.h"
 #include "gctoclreventsink.h"
 #include "thread.h"
+#include "eventtracebase.h"
 
 GCToCLREventSink g_gcToClrEventSink;
 
@@ -174,6 +175,14 @@ void GCToCLREventSink::FireGCAllocationTick_V4(uint64_t allocationAmount,
 {
     LIMITED_METHOD_CONTRACT;
 
+#ifdef FEATURE_EVENT_TRACE
+    if (IsRuntimeProviderEnabled(TRACE_LEVEL_INFORMATION, CLR_ALLOCATIONSAMPLING_KEYWORD))
+    {
+        // skip AllocationTick if AllocationSampled is emitted
+        return;
+    }
+#endif // FEATURE_EVENT_TRACE
+
     void * typeId = GetLastAllocEEType();
     WCHAR * name = nullptr;
 

src/coreclr/nativeaot/Runtime/thread.cpp

@@ -35,6 +35,13 @@ static Thread* g_RuntimeInitializingThread;
 
 #endif //!DACCESS_COMPILE
 
+ee_alloc_context::PerThreadRandom::PerThreadRandom()
+{
+    minipal_xoshiro128pp_init(&random_state, (uint32_t)PalGetTickCount64());
+}
+
+thread_local ee_alloc_context::PerThreadRandom ee_alloc_context::t_random = PerThreadRandom();
+
 PInvokeTransitionFrame* Thread::GetTransitionFrame()
 {
     if (ThreadStore::GetSuspendingThread() == this)

src/coreclr/nativeaot/Runtime/thread.h

@@ -6,6 +6,7 @@
 
 #include "StackFrameIterator.h"
 #include "slist.h" // DefaultSListTraits
+#include <minipal/xoshiro128pp.h>
 
 struct gc_alloc_context;
 class RuntimeInstance;
@@ -113,7 +114,19 @@ struct ee_alloc_context
 
     gc_alloc_context* GetGCAllocContext();
     uint8_t* GetCombinedLimit();
-    void UpdateCombinedLimit();
+    void UpdateCombinedLimit(bool samplingEnabled);
+    static bool IsRandomizedSamplingEnabled();
+    static uint32_t ComputeGeometricRandom();
+
+    struct PerThreadRandom
+    {
+        minipal_xoshiro128pp random_state;
+
+        PerThreadRandom();
+        double NextDouble();
+    };
+
+    static thread_local PerThreadRandom t_random;
 };
 
 

src/coreclr/nativeaot/Runtime/thread.inl

@@ -3,7 +3,9 @@
 
 #ifndef DACCESS_COMPILE
 
+#include "eventtracebase.h"
 
+const uint32_t SamplingDistributionMean = (100 * 1024);
 
 inline gc_alloc_context* ee_alloc_context::GetGCAllocContext()
 {
@@ -22,11 +24,49 @@ struct _thread_inl_gc_alloc_context
     uint8_t* alloc_limit;
 };
 
-inline void ee_alloc_context::UpdateCombinedLimit()
+
+inline bool ee_alloc_context::IsRandomizedSamplingEnabled()
+{
+#ifdef FEATURE_EVENT_TRACE
+    return IsRuntimeProviderEnabled(TRACE_LEVEL_INFORMATION, CLR_ALLOCATIONSAMPLING_KEYWORD);
+#else
+    return false;
+#endif // FEATURE_EVENT_TRACE
+}
+
+inline void ee_alloc_context::UpdateCombinedLimit(bool samplingEnabled)
+{
+    _thread_inl_gc_alloc_context* gc_alloc_context = (_thread_inl_gc_alloc_context*)GetGCAllocContext();
+    if (!samplingEnabled)
+    {
+        combined_limit = gc_alloc_context->alloc_limit;
+    }
+    else
+    {
+        // compute the next sampling budget based on a geometric distribution
+        size_t samplingBudget = ComputeGeometricRandom();
+
+        // if the sampling limit is larger than the allocation context, no sampling will occur in this AC
+        // We do Min() prior to adding to alloc_ptr to ensure alloc_ptr+samplingBudget doesn't cause an overflow.
+
+        size_t size = gc_alloc_context->alloc_limit - gc_alloc_context->alloc_ptr;
+        combined_limit = gc_alloc_context->alloc_ptr + min(samplingBudget, size);
+    }
+}
+
+inline uint32_t ee_alloc_context::ComputeGeometricRandom()
+{
+    // compute a random sample from the Geometric distribution.
+    double probability = t_random.NextDouble();
+    uint32_t threshold = (uint32_t)(-log(1 - probability) * SamplingDistributionMean);
+    return threshold;
+}
+
+// Returns a random double in the range [0, 1).
+inline double ee_alloc_context::PerThreadRandom::NextDouble()
 {
-    // The randomized allocation sampling feature is being submitted in stages. For now sampling is never enabled so
-    // combined_limit is always the same as alloc_limit.
-    combined_limit = ((_thread_inl_gc_alloc_context*)GetGCAllocContext())->alloc_limit;
+    uint32_t value = minipal_xoshiro128pp_next(&random_state);
+    return value * (1.0/(UINT32_MAX+1.0));
 }
 
 // Set the m_pDeferredTransitionFrame field for GC allocation helpers that setup transition frame