[Parameter Capturing] Make probe management APIs asynchronous

cspell.json

@@ -99,6 +99,7 @@
         "trce",
         "tstr",
         "ukwn",
+        "uninitialize",
         "Uninstallation",
         "uniquifier",
         "Unlocalized",

src/Microsoft.Diagnostics.Monitoring.HostingStartup/ParameterCapturing/FunctionProbes/FunctionProbesManager.cs

@@ -6,8 +6,11 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
+using System.Globalization;
 using System.Reflection;
 using System.Runtime.InteropServices;
+using System.Threading;
+using System.Threading.Tasks;
 
 namespace Microsoft.Diagnostics.Monitoring.HostingStartup.ParameterCapturing.FunctionProbes
 {
@@ -26,49 +29,219 @@ private static extern void RequestFunctionProbeInstallation(
             [MarshalAs(UnmanagedType.LPArray)] uint[] boxingTokens,
             [MarshalAs(UnmanagedType.LPArray)] uint[] boxingTokenCounts);
 
-        private readonly object _requestLocker = new();
+        private delegate void FunctionProbeRegistrationCallback(int hresult);
+        private delegate void FunctionProbeInstallationCallback(int hresult);
+        private delegate void FunctionProbeUninstallationCallback(int hresult);
+        private delegate void FunctionProbeFaultCallback(ulong uniquifier);
+
+        [DllImport(ProfilerIdentifiers.MutatingProfiler.LibraryRootFileName, CallingConvention = CallingConvention.StdCall, PreserveSig = false)]
+        private static extern void RegisterFunctionProbeCallbacks(
+            IntPtr onRegistration,
+            IntPtr onInstallation,
+            IntPtr onUninstallation,
+            IntPtr onFault);
+
+        [DllImport(ProfilerIdentifiers.MutatingProfiler.LibraryRootFileName, CallingConvention = CallingConvention.StdCall, PreserveSig = false)]
+        private static extern void UnregisterFunctionProbeCallbacks();
+
+        private readonly FunctionProbeRegistrationCallback _onRegistrationDelegate;
+        private readonly FunctionProbeInstallationCallback _onInstallationDelegate;
+        private readonly FunctionProbeUninstallationCallback _onUninstallationDelegate;
+        private readonly FunctionProbeFaultCallback _onFaultDelegate;
+
+        private long _probeState;
+        private const long ProbeStateUninitialized = default(long);
+        private const long ProbeStateUninstalled = 1;
+        private const long ProbeStateUninstalling = 2;
+        private const long ProbeStateInstalling = 3;
+        private const long ProbeStateInstalled = 4;
+        private const long ProbeStateUnrecoverable = 5;
+
+        private readonly TaskCompletionSource _probeRegistrationTaskSource = new(TaskCreationOptions.RunContinuationsAsynchronously);
+        private TaskCompletionSource? _installationTaskSource;
+        private TaskCompletionSource? _uninstallationTaskSource;
+
+        private readonly CancellationTokenSource _disposalTokenSource = new();
+        private readonly CancellationToken _disposalToken;
+
         private long _disposedState;
-        private bool _isCapturing;
+
+        public event EventHandler<InstrumentedMethod>? OnProbeFault;
 
         public FunctionProbesManager(IFunctionProbes probes)
         {
             ProfilerResolver.InitializeResolver<FunctionProbesManager>();
 
+            _disposalToken = _disposalTokenSource.Token;
+
+            //
+            // CONSIDER:
+            // Use UnmanagedCallersOnlyAttribute for the callbacks
+            // and use the address-of operator to get a function pointer
+            // to give to the profiler.
+            //
+
+            _onRegistrationDelegate = OnRegistration;
+            _onInstallationDelegate = OnInstallation;
+            _onUninstallationDelegate = OnUninstallation;
+            _onFaultDelegate = OnFault;
+
+            RegisterFunctionProbeCallbacks(
+                Marshal.GetFunctionPointerForDelegate(_onRegistrationDelegate),
+                Marshal.GetFunctionPointerForDelegate(_onInstallationDelegate),
+                Marshal.GetFunctionPointerForDelegate(_onUninstallationDelegate),
+                Marshal.GetFunctionPointerForDelegate(_onFaultDelegate));
+
             RequestFunctionProbeRegistration(FunctionProbesStub.GetProbeFunctionId());
+
             FunctionProbesStub.Instance = probes;
         }
 
-        public void StopCapturing()
+        private void OnRegistration(int hresult)
+        {
+            TransitionStateFromHr(_probeRegistrationTaskSource, hresult,
+                expectedState: ProbeStateUninitialized,
+                succeededState: ProbeStateUninstalled,
+                failedState: ProbeStateUnrecoverable);
+        }
+
+        private void OnInstallation(int hresult)
+        {
+            TransitionStateFromHr(_installationTaskSource, hresult,
+                expectedState: ProbeStateInstalling,
+                succeededState: ProbeStateInstalled,
+                failedState: ProbeStateUninstalled);
+        }
+
+        private void OnUninstallation(int hresult)
+        {
+            TransitionStateFromHr(_uninstallationTaskSource, hresult,
+                expectedState: ProbeStateUninstalling,
+                succeededState: ProbeStateUninstalled,
+                failedState: ProbeStateInstalled);
+        }
+
+        private void OnFault(ulong uniquifier)
         {
-            lock (_requestLocker)
+            var methodCache = FunctionProbesStub.InstrumentedMethodCache;
+            if (methodCache == null ||
+                !methodCache.TryGetValue(uniquifier, out InstrumentedMethod? instrumentedMethod))
             {
-                if (!_isCapturing)
-                {
-                    return;
-                }
+                //
+                // The probe fault occurred in a method that is no longer actively instrumented, ignore.
+                // This can happen when we request uninstallation of function probes and there's still a thread
+                // actively in one of the instrumented methods and it happens to fault.
+                //
+                return;
+            }
 
-                FunctionProbesStub.InstrumentedMethodCache = null;
-                RequestFunctionProbeUninstallation();
+            OnProbeFault?.Invoke(this, instrumentedMethod);
+        }
+
+        private void TransitionStateFromHr(TaskCompletionSource? taskCompletionSource, int hresult, long expectedState, long succeededState, long failedState)
+        {
+            Exception? ex = Marshal.GetExceptionForHR(hresult);
+            long newState = (ex == null) ? succeededState : failedState;
 
-                _isCapturing = false;
+            if (expectedState != Interlocked.CompareExchange(ref _probeState, newState, expectedState))
+            {
+                // Unexpected, the profiler is in a different state than us.
+                StateMismatch(expectedState);
+                return;
+            }
+
+            if (ex == null)
+            {
+                _ = taskCompletionSource?.TrySetResult();
+            }
+            else
+            {
+                _ = taskCompletionSource?.TrySetException(ex);
             }
         }
 
-        public void StartCapturing(IList<MethodInfo> methods)
+        private void StateMismatch(long expected)
         {
+            InvalidOperationException ex = new(string.Format(CultureInfo.InvariantCulture, ParameterCapturingStrings.ErrorMessage_ProbeStateMismatchFormatString, expected, _probeState));
+
+            _probeState = ProbeStateUnrecoverable;
+            _ = _probeRegistrationTaskSource.TrySetException(ex);
+            _ = _installationTaskSource?.TrySetException(ex);
+            _ = _uninstallationTaskSource?.TrySetException(ex);
+        }
+
+        public async Task StopCapturingAsync(CancellationToken token)
+        {
+            //
+            // NOTE: Do not await anything until after StopCapturingCore is called
+            // to ensure deterministic cancellation behavior.
+            //
+            DisposableHelper.ThrowIfDisposed<FunctionProbesManager>(ref _disposedState);
+
+            if (ProbeStateInstalled != Interlocked.CompareExchange(ref _probeState, ProbeStateUninstalling, ProbeStateInstalled))
+            {
+                throw new InvalidOperationException(string.Format(CultureInfo.InvariantCulture, ParameterCapturingStrings.ErrorMessage_ProbeStateMismatchFormatString, ProbeStateInstalled, _probeState));
+            }
+
+            _uninstallationTaskSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+            try
+            {
+                StopCapturingCore();
+            }
+            catch
+            {
+                _uninstallationTaskSource = null;
+                _probeState = ProbeStateInstalled;
+                throw;
+            }
+
+            using CancellationTokenSource cts = CancellationTokenSource.CreateLinkedTokenSource(_disposalToken, token);
+            try
+            {
+                using IDisposable _ = cts.Token.Register(() =>
+                {
+                    _uninstallationTaskSource.TrySetCanceled(cts.Token);
+                });
+                await _uninstallationTaskSource.Task.ConfigureAwait(false);
+            }
+            finally
+            {
+                _uninstallationTaskSource = null;
+            }
+        }
+
+        private void StopCapturingCore()
+        {
+            if (_probeState == ProbeStateUninstalled)
+            {
+                return;
+            }
+
+            FunctionProbesStub.InstrumentedMethodCache = null;
+            RequestFunctionProbeUninstallation();
+        }
+
+
+        public async Task StartCapturingAsync(IList<MethodInfo> methods, CancellationToken token)
+        {
+            DisposableHelper.ThrowIfDisposed<FunctionProbesManager>(ref _disposedState);
+
             if (methods.Count == 0)
             {
                 throw new ArgumentException(nameof(methods));
             }
 
-            Dictionary<ulong, InstrumentedMethod> newMethodCache = new(methods.Count);
-            lock (_requestLocker)
+            // _probeRegistrationTaskSource will be cancelled (if needed) on dispose
+            await _probeRegistrationTaskSource.Task.WaitAsync(token).ConfigureAwait(false);
+
+            if (ProbeStateUninstalled != Interlocked.CompareExchange(ref _probeState, ProbeStateInstalling, ProbeStateUninstalled))
             {
-                if (_isCapturing)
-                {
-                    throw new InvalidOperationException();
-                }
+                throw new InvalidOperationException(string.Format(CultureInfo.InvariantCulture, ParameterCapturingStrings.ErrorMessage_ProbeStateMismatchFormatString, ProbeStateUninstalled, _probeState));
+            }
 
+            try
+            {
+                Dictionary<ulong, InstrumentedMethod> newMethodCache = new(methods.Count);
                 List<ulong> functionIds = new(methods.Count);
                 List<uint> argumentCounts = new(methods.Count);
                 List<uint> boxingTokens = new();
@@ -78,7 +251,7 @@ public void StartCapturing(IList<MethodInfo> methods)
                     ulong functionId = method.GetFunctionId();
                     if (functionId == 0)
                     {
-                        return;
+                        throw new InvalidOperationException(string.Format(CultureInfo.InvariantCulture, ParameterCapturingStrings.ErrorMessage_FunctionDoesNotHaveIdFormatString, method.Name));
                     }
 
                     uint[] methodBoxingTokens = BoxingTokens.GetBoxingTokens(method);
@@ -94,13 +267,52 @@ public void StartCapturing(IList<MethodInfo> methods)
                 }
 
                 FunctionProbesStub.InstrumentedMethodCache = new ReadOnlyDictionary<ulong, InstrumentedMethod>(newMethodCache);
+
+                _installationTaskSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
                 RequestFunctionProbeInstallation(
                     functionIds.ToArray(),
                     (uint)functionIds.Count,
                     boxingTokens.ToArray(),
                     argumentCounts.ToArray());
+            }
+            catch
+            {
+                FunctionProbesStub.InstrumentedMethodCache = null;
+                _probeState = ProbeStateUninstalled;
+                _installationTaskSource = null;
+                throw;
+            }
+
+            using CancellationTokenSource cts = CancellationTokenSource.CreateLinkedTokenSource(_disposalToken, token);
+            try
+            {
+                using IDisposable _ = cts.Token.Register(() =>
+                {
+                    //
+                    // 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);
+                });
 
-                _isCapturing = true;
+                await _installationTaskSource.Task.ConfigureAwait(false);
+            }
+            finally
+            {
+                _installationTaskSource = null;
             }
         }
 
@@ -109,14 +321,24 @@ public void Dispose()
             if (!DisposableHelper.CanDispose(ref _disposedState))
                 return;
 
-            FunctionProbesStub.Instance = null;
             try
             {
-                StopCapturing();
+                _disposalTokenSource.Cancel();
             }
             catch
             {
+            }
 
+            _ = _probeRegistrationTaskSource.TrySetCanceled(_disposalTokenSource.Token);
+            _disposalTokenSource.Dispose();
+
+            try
+            {
+                UnregisterFunctionProbeCallbacks();
+                StopCapturingCore();
+            }
+            catch
+            {
             }
         }
     }

src/Microsoft.Diagnostics.Monitoring.HostingStartup/ParameterCapturing/FunctionProbes/IFunctionProbesManager.cs

@@ -4,13 +4,17 @@
 using System;
 using System.Collections.Generic;
 using System.Reflection;
+using System.Threading;
+using System.Threading.Tasks;
 
 namespace Microsoft.Diagnostics.Monitoring.HostingStartup.ParameterCapturing.FunctionProbes
 {
     internal interface IFunctionProbesManager : IDisposable
     {
-        public void StartCapturing(IList<MethodInfo> methods);
+        public Task StartCapturingAsync(IList<MethodInfo> methods, CancellationToken token);
 
-        public void StopCapturing();
+        public Task StopCapturingAsync(CancellationToken token);
+
+        public event EventHandler<InstrumentedMethod> OnProbeFault;
     }
 }

src/Microsoft.Diagnostics.Monitoring.HostingStartup/ParameterCapturing/InstrumentedMethod.cs

@@ -19,6 +19,7 @@ public sealed class InstrumentedMethod
 
         public InstrumentedMethod(MethodInfo method, uint[] boxingTokens)
         {
+            FunctionId = method.GetFunctionId();
             SupportedParameters = BoxingTokens.AreParametersSupported(boxingTokens);
             MethodWithParametersTemplateString = PrettyPrinter.ConstructTemplateStringFromMethod(method, SupportedParameters);
             foreach (bool isParameterSupported in SupportedParameters)
@@ -67,5 +68,7 @@ private static ParameterCaptureMode ComputeCaptureMode(string? typeName)
         /// The number of format items equals NumberOfSupportedParameters.
         /// </summary>
         public string MethodWithParametersTemplateString { get; }
+
+        public ulong FunctionId { get; }
     }
 }