diff --git a/src/Microsoft.Diagnostics.Monitoring.HostingStartup/ParameterCapturing/FunctionProbes/FunctionProbesManager.cs b/src/Microsoft.Diagnostics.Monitoring.HostingStartup/ParameterCapturing/FunctionProbes/FunctionProbesManager.cs index 685aa80381f..bf54cccf881 100644 --- a/src/Microsoft.Diagnostics.Monitoring.HostingStartup/ParameterCapturing/FunctionProbes/FunctionProbesManager.cs +++ b/src/Microsoft.Diagnostics.Monitoring.HostingStartup/ParameterCapturing/FunctionProbes/FunctionProbesManager.cs @@ -62,6 +62,8 @@ private static extern void RegisterFunctionProbeCallbacks( private TaskCompletionSource? _uninstallationTaskSource; private readonly CancellationTokenSource _disposalTokenSource = new(); + private readonly CancellationToken _disposalToken; + private long _disposedState; public event EventHandler? OnProbeFault; @@ -70,6 +72,8 @@ public FunctionProbesManager(IFunctionProbes probes) { ProfilerResolver.InitializeResolver(); + _disposalToken = _disposalTokenSource.Token; + _onRegistrationDelegate = OnRegistration; _onInstallationDelegate = OnInstallation; _onUninstallationDelegate = OnUninstallation; @@ -180,7 +184,7 @@ public async Task StopCapturingAsync(CancellationToken token) throw; } - using CancellationTokenSource cts = CancellationTokenSource.CreateLinkedTokenSource(_disposalTokenSource.Token, token); + using CancellationTokenSource cts = CancellationTokenSource.CreateLinkedTokenSource(_disposalToken, token); using IDisposable _ = cts.Token.Register(() => { _uninstallationTaskSource.TrySetCanceled(cts.Token); @@ -261,11 +265,26 @@ public async Task StartCapturingAsync(IList methods, CancellationTok throw; } - using CancellationTokenSource cts = CancellationTokenSource.CreateLinkedTokenSource(_disposalTokenSource.Token, token); + using CancellationTokenSource cts = CancellationTokenSource.CreateLinkedTokenSource(_disposalToken, token); using IDisposable _ = cts.Token.Register(() => { - // On cancellation, simply stop anyone waiting on it. - // The actual installation process isn't trivial to stop at this point. + // + // We need to uninstall the probes ourselves here if dispose has happened otherwise the probes could be left in an installed state. + // + // NOTE: It's possible that StopCapturingCore could be called twice by doing this - once by Dispose and once by us, this is OK + // as the native layer will gracefully handle multiple RequestFunctionProbeUninstallation calls. + // + if (_disposalToken.IsCancellationRequested) + { + try + { + StopCapturingCore(); + } + catch + { + } + } + _installationTaskSource.TrySetCanceled(cts.Token); });