fix sample collection

benchmarks/Sentry.Benchmarks/ProfilingBenchmarks.cs

@@ -142,7 +142,7 @@ public void DiagnosticsSessionStartCopyStop(bool rundown, string provider)
     [Benchmark]
     public void SampleProfilerSessionStartStop()
     {
-        var session = SampleProfilerSession.StartNew();
+        using var session = SampleProfilerSession.StartNew();
         session.Stop();
     }
     #endregion

src/Sentry.Profiling/TraceLogProcessor.cs → src/Sentry.Profiling/SampleProfileBuilder.cs

@@ -1,6 +1,5 @@
 using Microsoft.Diagnostics.Tracing;
 using Microsoft.Diagnostics.Tracing.Etlx;
-using Microsoft.Diagnostics.Tracing.EventPipe;
 using Sentry.Extensibility;
 using Sentry.Internal;
 using Sentry.Protocol;
@@ -13,7 +12,7 @@ namespace Sentry.Profiling;
 /// <summary>
 /// Processes TraceLog to produce a SampleProfile.
 /// </summary>
-internal class TraceLogProcessor
+internal class SampleProfileBuilder
 {
     private readonly SentryOptions _options;
     private readonly TraceLog _traceLog;
@@ -30,16 +29,13 @@ internal class TraceLogProcessor
     // A sparse array mapping from a ThreadIndex to an index in Profile.Threads.
     private readonly SparseScalarArray<int> _threadIndexes = new(-1, 10);
 
-    private readonly double _timeShiftMs;
-
-    public TraceLogProcessor(SentryOptions options, TraceLog traceLog, double TimeShiftMs)
+    public SampleProfileBuilder(SentryOptions options, TraceLog traceLog)
     {
         _options = options;
         _traceLog = traceLog;
-        _timeShiftMs = TimeShiftMs;
     }
 
-    internal void AddSample(TraceEvent data)
+    internal void AddSample(TraceEvent data, double timestampMs)
     {
         var thread = data.Thread();
         if (thread.ThreadIndex == ThreadIndex.Invalid)
@@ -55,8 +51,6 @@ internal void AddSample(TraceEvent data)
             return;
         }
 
-        var timestampMs = data.TimeStampRelativeMSec + _timeShiftMs;
-
         var stackIndex = AddStackTrace(callStackIndex);
         if (stackIndex < 0)
         {

src/Sentry.Profiling/SampleProfilerSession.cs

@@ -3,6 +3,7 @@
 using Microsoft.Diagnostics.Tracing.Etlx;
 using Microsoft.Diagnostics.Tracing.EventPipe;
 using Microsoft.Diagnostics.Tracing.Parsers;
+using Sentry.Extensibility;
 using Sentry.Internal;
 
 namespace Sentry.Profiling;
@@ -12,25 +13,28 @@ internal class SampleProfilerSession : IDisposable
     private readonly EventPipeSession _session;
     private readonly TraceLogEventSource _eventSource;
     private readonly SampleProfilerTraceEventParser _sampleEventParser;
+    private readonly IDiagnosticLogger? _logger;
     private readonly SentryStopwatch _stopwatch = SentryStopwatch.StartNew();
     private bool _stopped;
 
-    private SampleProfilerSession(EventPipeSession session)
+    private SampleProfilerSession(EventPipeSession session, IDiagnosticLogger? logger)
     {
         _session = session;
+        _logger = logger;
         _eventSource = TraceLog.CreateFromEventPipeSession(_session);
         _sampleEventParser = new SampleProfilerTraceEventParser(_eventSource);
-        _eventSource.Process();
+        // Process() blocks until the session is stopped so we need to run it on a separate thread.
+        Task.Factory.StartNew(_eventSource.Process, TaskCreationOptions.LongRunning);
     }
 
     // Exposed only for benchmarks.
     internal static EventPipeProvider[] Providers = new[]
     {
         // Note: all events we need issued by "DotNETRuntime" provider are at "EventLevel.Informational"
         // see https://learn.microsoft.com/en-us/dotnet/fundamentals/diagnostics/runtime-events
-        new EventPipeProvider("Microsoft-Windows-DotNETRuntime", EventLevel.Informational, (long) ClrTraceEventParser.Keywords.Default),
-        new EventPipeProvider("Microsoft-DotNETCore-SampleProfiler", EventLevel.Informational),
-        // new EventPipeProvider("System.Threading.Tasks.TplEventSource", EventLevel.Informational, (long) TplEtwProviderTraceEventParser.Keywords.Default)
+        new EventPipeProvider(ClrTraceEventParser.ProviderName, EventLevel.Informational, (long) ClrTraceEventParser.Keywords.Default),
+        new EventPipeProvider(SampleProfilerTraceEventParser.ProviderName, EventLevel.Informational),
+        // new EventPipeProvider(TplEtwProviderTraceEventParser.ProviderName, EventLevel.Informational, (long) TplEtwProviderTraceEventParser.Keywords.Default)
     };
 
     // Exposed only for benchmarks.
@@ -46,21 +50,28 @@ private SampleProfilerSession(EventPipeSession session)
 
     public TraceLog TraceLog => _eventSource.TraceLog;
 
-    public static SampleProfilerSession StartNew()
+    public static SampleProfilerSession StartNew(IDiagnosticLogger? logger = null)
     {
         var client = new DiagnosticsClient(Process.GetCurrentProcess().Id);
         var session = client.StartEventPipeSession(Providers, RequestRundown, CircularBufferMB);
-        return new SampleProfilerSession(session);
+        return new SampleProfilerSession(session, logger);
     }
 
     public void Stop()
     {
         if (!_stopped)
         {
-            _stopped = true;
-            _session.Stop();
-            _session.Dispose();
-            _eventSource.Dispose();
+            try
+            {
+                _stopped = true;
+                _session.Stop();
+                _session.Dispose();
+                _eventSource.Dispose();
+            }
+            catch (Exception ex)
+            {
+                _logger?.LogWarning("Error during sampler profiler session shutdown.", ex);
+            }
         }
     }
 

src/Sentry.Profiling/SamplingTransactionProfiler.cs

@@ -1,3 +1,4 @@
+using Microsoft.Diagnostics.Tracing;
 using Sentry.Extensibility;
 using Sentry.Internal;
 using Sentry.Protocol;
@@ -8,18 +9,23 @@ internal class SamplingTransactionProfiler : ITransactionProfiler
 {
     public Action? OnFinish;
     private readonly CancellationToken _cancellationToken;
-    private TimeSpan? _duration;
+    private bool _stopped = false;
     private readonly SentryOptions _options;
-    private TraceLogProcessor _processor;
+    private SampleProfileBuilder _processor;
     private SampleProfilerSession _session;
+    private readonly double _startTimeMs;
+    private double _endTimeMs;
+    private TaskCompletionSource _completionSource = new();
 
     public SamplingTransactionProfiler(SentryOptions options, SampleProfilerSession session, int timeoutMs, CancellationToken cancellationToken)
     {
         _options = options;
         _session = session;
         _cancellationToken = cancellationToken;
-        _processor = new TraceLogProcessor(options, session.TraceLog, -session.Elapsed.TotalMilliseconds);
-        session.SampleEventParser.ThreadSample += _processor.AddSample;
+        _startTimeMs = session.Elapsed.TotalMilliseconds;
+        _endTimeMs = Double.MaxValue;
+        _processor = new SampleProfileBuilder(options, session.TraceLog);
+        session.SampleEventParser.ThreadSample += OnThreadSample;
         cancellationToken.Register(() =>
         {
             if (Stop())
@@ -29,30 +35,49 @@ public SamplingTransactionProfiler(SentryOptions options, SampleProfilerSession
         });
         Task.Delay(timeoutMs, cancellationToken).ContinueWith(_ =>
         {
-            if (Stop(TimeSpan.FromMilliseconds(timeoutMs)))
+            if (Stop(_startTimeMs + timeoutMs))
             {
                 options.LogDebug("Profiling is being cut-of after {0} ms because the transaction takes longer than that.", timeoutMs);
             }
         }, CancellationToken.None);
     }
 
-    private bool Stop(TimeSpan? duration = null)
+    private bool Stop(double? endTimeMs = null)
     {
-        if (_duration is null)
+        endTimeMs ??= _session.Elapsed.TotalMilliseconds;
+        if (!_stopped)
         {
             lock (_session)
             {
-                if (_duration is null)
+                if (!_stopped)
                 {
-                    _duration = duration ?? _session.Elapsed;
-                    _session.SampleEventParser.ThreadSample -= _processor.AddSample;
+                    _stopped = true;
+                    _endTimeMs = endTimeMs.Value;
                     return true;
                 }
             }
         }
         return false;
     }
 
+    // We need custom sampling because the TraceLog dispatches events from a queue with a delay of about 2 seconds.
+    private void OnThreadSample(TraceEvent data)
+    {
+        var timestampMs = data.TimeStampRelativeMSec;
+        if (timestampMs >= _startTimeMs)
+        {
+            if (timestampMs <= _endTimeMs)
+            {
+                _processor.AddSample(data, timestampMs - _startTimeMs);
+            }
+            else
+            {
+                _session.SampleEventParser.ThreadSample -= OnThreadSample;
+                _completionSource.TrySetResult();
+            }
+        }
+    }
+
     /// <inheritdoc />
     public void Finish()
     {
@@ -62,16 +87,19 @@ public void Finish()
         }
     }
 
-    // TODO doesn't need to be async anymore...
     /// <inheritdoc />
     public async Task<ProfileInfo> CollectAsync(Transaction transaction)
     {
-        if (_duration is null)
+        if (!_stopped)
         {
             throw new InvalidOperationException("Profiler.CollectAsync() called before Finish()");
         }
-        _cancellationToken.ThrowIfCancellationRequested();
-        return await Task.FromResult(CreateProfileInfo(transaction, _processor.Profile)).ConfigureAwait(false);
+
+        // Wait for the last sample (<= _endTimeMs), or at most 10 seconds. The timeout shouldn't happen because
+        // TraceLog.realTimeQueue should dispatch events after ~2 seconds, but if it does, send what we have.
+        await Task.WhenAny(_completionSource.Task, Task.Delay(10_000, _cancellationToken)).ConfigureAwait(false);
+
+        return CreateProfileInfo(transaction, _processor.Profile);
     }
 
     internal static ProfileInfo CreateProfileInfo(Transaction transaction, SampleProfile profile)

src/Sentry.Profiling/SamplingTransactionProfilerFactory.cs

@@ -19,12 +19,13 @@ internal class SamplingTransactionProfilerFactory : IDisposable, ITransactionPro
 
     public static SamplingTransactionProfilerFactory Create(SentryOptions options)
     {
-        return new SamplingTransactionProfilerFactory(options, SampleProfilerSession.StartNew());
+        var session = SampleProfilerSession.StartNew(options.DiagnosticLogger);
+        return new SamplingTransactionProfilerFactory(options, session);
     }
 
     public static async Task<SamplingTransactionProfilerFactory> CreateAsync(SentryOptions options)
     {
-        var session = await Task.Run(() => SampleProfilerSession.StartNew()).ConfigureAwait(false);
+        var session = await Task.Run(() => SampleProfilerSession.StartNew(options.DiagnosticLogger)).ConfigureAwait(false);
         return new SamplingTransactionProfilerFactory(options, session);
     }
 
@@ -59,6 +60,6 @@ private SamplingTransactionProfilerFactory(SentryOptions options, SampleProfiler
 
     public void Dispose()
     {
-        _session.Stop();
+        _session.Dispose();
     }
 }

test/Sentry.Profiling.Tests/SamplingTransactionProfilerTests.cs

@@ -49,10 +49,11 @@ private void ValidateProfile(SampleProfile profile, ulong maxTimestampNs)
 
     private void RunForMs(int milliseconds)
     {
-        for (int i = 0; i < milliseconds / 20; i++)
+        var clock = Stopwatch.StartNew();
+        while (clock.ElapsedMilliseconds < milliseconds)
         {
-            _testOutputLogger.LogDebug("sleeping...");
-            Thread.Sleep(20);
+            _testOutputLogger.LogDebug("Sleeping... time remaining: {0} ms", milliseconds - clock.ElapsedMilliseconds);
+            Thread.Sleep((int)Math.Min(milliseconds / 5, milliseconds - clock.ElapsedMilliseconds));
         }
     }
 
@@ -83,7 +84,7 @@ public void Profiler_AfterTimeout_Stops()
     {
         var hub = Substitute.For<IHub>();
         var options = new SentryOptions { DiagnosticLogger = _testOutputLogger };
-        using var session = SampleProfilerSession.StartNew();
+        using var session = SampleProfilerSession.StartNew(_testOutputLogger);
         var limitMs = 50;
         var sut = new SamplingTransactionProfiler(options, session, limitMs, CancellationToken.None);
         RunForMs(limitMs * 4);