fix(profiling): release the EventPipe session when the SDK shuts down

src/Sentry.Profiling/ProfilingIntegration.cs

@@ -6,10 +6,14 @@ namespace Sentry.Profiling;
 /// <summary>
 /// Enables transaction performance profiling.
 /// </summary>
-public class ProfilingIntegration : ISdkIntegration
+public class ProfilingIntegration : ISdkIntegration, IDisposable
 {
     private TimeSpan _startupTimeout;
 
+    // Only set when this integration created the factory, so that Dispose() never tears down a
+    // factory that was supplied by someone else.
+    private IDisposable? _ownedFactory;
+
     /// <summary>
     /// Initializes the profiling integration.
     /// </summary>
@@ -35,7 +39,12 @@ public void Register(IHub hub, SentryOptions options)
         {
             try
             {
-                options.TransactionProfilerFactory ??= new SamplingTransactionProfilerFactory(options, _startupTimeout);
+                if (options.TransactionProfilerFactory is null)
+                {
+                    var factory = new SamplingTransactionProfilerFactory(options, _startupTimeout);
+                    options.TransactionProfilerFactory = factory;
+                    _ownedFactory = factory;
+                }
             }
             catch (Exception e)
             {
@@ -47,4 +56,14 @@ public void Register(IHub hub, SentryOptions options)
             options.LogInfo("Profiling Integration is disabled because profiling is disabled by configuration.");
         }
     }
+
+    /// <summary>
+    /// Stops the profiler session started by this integration, releasing the underlying EventPipe
+    /// session. Called by the SDK on shutdown.
+    /// </summary>
+    public void Dispose()
+    {
+        _ownedFactory?.Dispose();
+        _ownedFactory = null;
+    }
 }

src/Sentry.Profiling/SampleProfilerSession.cs

@@ -47,6 +47,10 @@ private SampleProfilerSession(SentryStopwatch stopwatch, EventPipeSession sessio
     // need a large buffer if we're connecting righ away. Leaving it too large increases app memory usage.
     internal static int CircularBufferMB = 16;
 
+    // How long Stop() waits for the event processing task to drain after the session has been stopped.
+    // Draining should be near-instant; this only bounds the worst case so shutdown can't hang.
+    internal const int ProcessingDrainTimeoutMs = 2_000;
+
     // Exposed for tests
     internal TraceLogEventSource EventSource { get; }
 
@@ -83,14 +87,17 @@ public static SampleProfilerSession StartNew(IDiagnosticLogger? logger = null)
             var eventSource = TraceLog.CreateFromEventPipeSession(session, TraceLog.EventPipeRundownConfiguration.Enable(client));
 
             // Process() blocks until the session is stopped so we need to run it on a separate thread.
+            // Note: the continuation is deliberately unconditional. A continuation whose criteria aren't
+            // met (e.g. OnlyOnFaulted when Process() returns normally) transitions to Canceled, which
+            // would make the Wait() in Stop() throw on every clean shutdown.
             var processing = Task.Factory.StartNew(eventSource.Process, TaskCreationOptions.LongRunning)
                 .ContinueWith(_ =>
                 {
                     if (_.Exception?.InnerException is { } e)
                     {
                         logger?.LogWarning(e, "Error during sampler profiler EventPipeSession processing.");
                     }
-                }, TaskContinuationOptions.OnlyOnFaulted);
+                });
 
             return new SampleProfilerSession(stopWatch, session, eventSource, processing, logger);
         }
@@ -119,19 +126,39 @@ public async Task WaitForFirstEventAsync(CancellationToken cancellationToken = d
 
     public void Stop()
     {
-        if (!_stopped)
+        if (_stopped)
+        {
+            return;
+        }
+
+        _stopped = true;
+        try
+        {
+            _session.Stop();
+
+            // Let the processing task drain the events that are still in flight, but don't hold up
+            // shutdown indefinitely if it doesn't get there.
+            if (!_processing.Wait(ProcessingDrainTimeoutMs))
+            {
+                _logger?.LogWarning("Sampler profiler event processing didn't finish within {0} ms of stopping the session.", ProcessingDrainTimeoutMs);
+            }
+        }
+        catch (Exception ex)
+        {
+            _logger?.LogWarning(ex, "Error during sampler profiler session shutdown.");
+        }
+        finally
         {
+            // These need to happen even if stopping the session or draining the events failed, otherwise
+            // the EventPipe connection to the runtime is left open.
             try
             {
-                _stopped = true;
-                _session.Stop();
-                _processing.Wait();
                 _session.Dispose();
                 EventSource.Dispose();
             }
             catch (Exception ex)
             {
-                _logger?.LogWarning(ex, "Error during sampler profiler session shutdown.");
+                _logger?.LogWarning(ex, "Error disposing the sampler profiler session.");
             }
         }
     }

src/Sentry.Profiling/SamplingTransactionProfilerFactory.cs

@@ -14,10 +14,26 @@ internal class SamplingTransactionProfilerFactory : IDisposable, ITransactionPro
     // Stop profiling after the given number of milliseconds.
     private const int TIME_LIMIT_MS = 30_000;
 
+    // How long Dispose() waits for an in-flight session startup to complete before giving up on it.
+    private const int SHUTDOWN_TIMEOUT_MS = 2_000;
+
     private readonly SentryOptions _options;
 
     internal Task<SampleProfilerSession> _sessionTask;
 
+    // Cancels the wait for the first event so that Dispose() doesn't have to wait for a session that
+    // may never receive one.
+    private readonly CancellationTokenSource _shutdownCts = new();
+
+    // Assigned as soon as the session exists, which is earlier than _sessionTask completing. Dispose()
+    // uses this so it can also stop a session that never saw its first event.
+    private SampleProfilerSession? _session;
+
+    private int _disposed;
+
+    // Exposed for tests.
+    internal bool IsDisposed => Volatile.Read(ref _disposed) != 0;
+
     private bool _errorLogged = false;
 
     public SamplingTransactionProfilerFactory(SentryOptions options, TimeSpan startupTimeout)
@@ -28,9 +44,10 @@ public SamplingTransactionProfilerFactory(SentryOptions options, TimeSpan startu
         {
             // This can block up to 30 seconds. The timeout is out of our hands.
             var session = SampleProfilerSession.StartNew(options.DiagnosticLogger);
+            _session = session;
 
-            // This can block indefinitely.
-            await session.WaitForFirstEventAsync().ConfigureAwait(false);
+            // This can block indefinitely, so it's cancelled when the factory is disposed.
+            await session.WaitForFirstEventAsync(_shutdownCts.Token).ConfigureAwait(false);
 
             return session;
         });
@@ -83,6 +100,34 @@ public SamplingTransactionProfilerFactory(SentryOptions options, TimeSpan startu
 
     public void Dispose()
     {
-        _sessionTask.ContinueWith(session => session.Dispose());
+        if (Interlocked.Exchange(ref _disposed, 1) != 0)
+        {
+            return;
+        }
+
+        // Unblocks the startup task if it's still waiting for the first event to arrive.
+        _shutdownCts.Cancel();
+
+        try
+        {
+            // Gives an in-flight startup a chance to finish, and observes the exception if it failed
+            // or was cancelled above.
+            _sessionTask.Wait(SHUTDOWN_TIMEOUT_MS);
+        }
+        catch (Exception e)
+        {
+            _options.LogDebug("Profiler session didn't start up cleanly before shutdown: {0}", e.Message);
+        }
+
+        try
+        {
+            _session?.Dispose();
+        }
+        catch (Exception e)
+        {
+            _options.LogWarning(e, "Failed to stop the profiler session.");
+        }
+
+        _shutdownCts.Dispose();
     }
 }

test/Sentry.Profiling.Tests/ProfilingSentryOptionsExtensionsTests.cs

@@ -61,6 +61,40 @@ public void DisableProfilingIntegration_RemovesProfilingIntegration()
         Assert.DoesNotContain(integrations, i => i is ProfilingIntegration);
     }
 
+    [Fact]
+    public void HubDispose_DisposesTheProfilerFactoryItCreated()
+    {
+        _options.TracesSampleRate = 1.0;
+        _options.ProfilesSampleRate = 1.0;
+
+        var hub = GetSut();
+        var factory = (SamplingTransactionProfilerFactory)_options.TransactionProfilerFactory!;
+        Assert.False(factory.IsDisposed);
+
+        hub.Dispose();
+
+        // ProfilingIntegration must be IDisposable for the Hub to register it for cleanup - otherwise
+        // the factory (and the EventPipe session it owns) is never disposed on SDK shutdown.
+        Assert.True(factory.IsDisposed);
+    }
+
+    [Fact]
+    public void HubDispose_DoesNotDisposeAProfilerFactoryItDidNotCreate()
+    {
+        _options.TracesSampleRate = 1.0;
+        _options.ProfilesSampleRate = 1.0;
+
+        var externalFactory = Substitute.For<ITransactionProfilerFactory, IDisposable>();
+        _options.TransactionProfilerFactory = externalFactory;
+
+        using (var hub = GetSut())
+        {
+            Assert.Same(externalFactory, _options.TransactionProfilerFactory);
+        }
+
+        ((IDisposable)externalFactory).DidNotReceive().Dispose();
+    }
+
     [Fact]
     public void AddProfilingIntegration_DoesntDuplicate()
     {

test/Sentry.Profiling.Tests/SamplingTransactionProfilerTests.cs

@@ -228,6 +228,27 @@ private static long MethodToBeLoaded(int n)
         return -n;
     }
 
+    [SkippableFact]
+    public async Task Session_Stop_ShutsDownWithoutError()
+    {
+        Skip.If(TestEnvironment.IsGitHubActions, "Flaky in CI");
+
+        SampleProfilerSession? session = null;
+        SkipIfFailsInCI(() => session = SampleProfilerSession.StartNew(_testOutputLogger));
+        await session!.WaitForFirstEventAsync(CancellationToken.None);
+
+        session.Stop();
+
+        // The event processing task used to be an OnlyOnFaulted continuation, which transitions to
+        // Canceled when processing completes normally. Waiting on it therefore threw on every clean
+        // shutdown, and the EventPipeSession and TraceLogEventSource were left undisposed.
+        _testOutputLogger.Entries.Select(e => e.Message).Should().NotContain(
+            m => m.StartsWith("Error during sampler profiler session shutdown"));
+
+        // Stopping again must stay a no-op.
+        session.Stop();
+    }
+
     [SkippableTheory]
     [InlineData(true)]
     [InlineData(false)]