[Parameter Capturing] Make probe management APIs asynchronous

cspell.json

@@ -98,6 +98,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
 {
@@ -25,50 +28,172 @@ private static extern void RequestFunctionProbeInstallation(
             uint count,
             [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(
+            FunctionProbeRegistrationCallback onRegistration,
+            FunctionProbeInstallationCallback onInstallation,
+            FunctionProbeUninstallationCallback onUninstallation,
+            FunctionProbeFaultCallback onFault);
+
+        [DllImport(ProfilerIdentifiers.MutatingProfiler.LibraryRootFileName, CallingConvention = CallingConvention.StdCall, PreserveSig = false)]
+        private static extern void UnregisterFunctionProbeCallbacks();
+
+        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 long _disposedState;
-        private bool _isCapturing;
+
+        public event EventHandler<InstrumentedMethod>? OnProbeFault;
 
         public FunctionProbesManager(IFunctionProbes probes)
         {
             ProfilerResolver.InitializeResolver<FunctionProbesManager>();
 
+            RegisterFunctionProbeCallbacks(OnRegistration, OnInstallation, OnUninstallation, OnFault);
             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;
+
+            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);
+            }
+        }
+
+        private void StateMismatch(long expected)
+        {
+            InvalidOperationException ex = new(string.Format(CultureInfo.InvariantCulture, ParameterCapturingStrings.ErrorMessage_ProbeStateMismatchFormatString, expected, _probeState));
+
+            _probeState = ProbeStateUnrecoverable;
+            _ = _installationTaskSource?.TrySetException(ex);
+            _ = _uninstallationTaskSource?.TrySetException(ex);
+            _ = _probeRegistrationTaskSource?.TrySetException(ex);
+        }
+
+        public async Task StopCapturingAsync(CancellationToken token)
+        {
+            if (ProbeStateInstalled != Interlocked.CompareExchange(ref _probeState, ProbeStateUninstalling, ProbeStateInstalled))
+            {
+                throw new InvalidOperationException();
+            }
 
-                _isCapturing = false;
+            _uninstallationTaskSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+            try
+            {
+                StopCapturingCore();
             }
+            catch
+            {
+                _uninstallationTaskSource = null;
+                _probeState = ProbeStateInstalled;
+                throw;
+            }
+
+            await _uninstallationTaskSource.Task.WaitAsync(token).ConfigureAwait(false);
         }
 
-        public void StartCapturing(IList<MethodInfo> methods)
+        private void StopCapturingCore()
+        {
+            if (_probeState == ProbeStateUninstalled)
+            {
+                return;
+            }
+
+            FunctionProbesStub.InstrumentedMethodCache = null;
+            RequestFunctionProbeUninstallation();
+        }
+
+
+        public async Task StartCapturingAsync(IList<MethodInfo> methods, CancellationToken token)
         {
             if (methods.Count == 0)
             {
                 throw new ArgumentException(nameof(methods));
             }
 
-            Dictionary<ulong, InstrumentedMethod> newMethodCache = new(methods.Count);
-            lock (_requestLocker)
+            await _probeRegistrationTaskSource.Task.WaitAsync(token).ConfigureAwait(false);
+
+            if (ProbeStateUninstalled != Interlocked.CompareExchange(ref _probeState, ProbeStateInstalling, ProbeStateUninstalled))
             {
-                if (_isCapturing)
-                {
-                    throw new InvalidOperationException();
-                }
+                throw new InvalidOperationException();
+            }
 
+            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 +203,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,14 +219,23 @@ 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());
-
-                _isCapturing = true;
             }
+            catch
+            {
+                FunctionProbesStub.InstrumentedMethodCache = null;
+                _probeState = ProbeStateUninstalled;
+                _installationTaskSource = null;
+                throw;
+            }
+
+            await _installationTaskSource.Task.WaitAsync(token).ConfigureAwait(false);
         }
 
         public void Dispose()
@@ -110,13 +244,17 @@ public void Dispose()
                 return;
 
             FunctionProbesStub.Instance = null;
+
+            _ = _installationTaskSource?.TrySetCanceled();
+            _ = _uninstallationTaskSource?.TrySetCanceled();
+
             try
             {
-                StopCapturing();
+                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; }
     }
 }