[Parameter Capturing] Support sending messages to managed in-proc features from dotnet-monitor

src/Microsoft.Diagnostics.Monitoring.HostingStartup/Microsoft.Diagnostics.Monitoring.HostingStartup.csproj

@@ -13,17 +13,18 @@
   </ItemGroup>
 
   <ItemGroup>
-    <Compile Include="..\Tools\dotnet-monitor\ToolIdentifiers.cs" />
     <Compile Include="..\Tools\dotnet-monitor\InProcessFeatures\InProcessFeaturesIdentifiers.cs" Link="InProcessFeaturesIdentifiers.cs" />
-    <Compile Include="..\Tools\dotnet-monitor\Profiler\ProfilerIdentifiers.cs" />
-    <Compile Include="..\Tools\dotnet-monitor\DisposableHelper.cs" Link="DisposableHelper.cs" />
   </ItemGroup>
 
   <ItemGroup>
     <InternalsVisibleTo Include="Microsoft.Diagnostics.Monitoring.HostingStartup.UnitTests" />
     <InternalsVisibleTo Include="Microsoft.Diagnostics.Monitoring.UnitTestApp" />
   </ItemGroup>
 
+  <ItemGroup>
+    <ProjectReference Include="..\Microsoft.Diagnostics.Monitoring.StartupHook\Microsoft.Diagnostics.Monitoring.StartupHook.csproj" />
+  </ItemGroup>
+
   <ItemGroup>
     <Compile Update="ParameterCapturing\ParameterCapturingStrings.Designer.cs">
       <DesignTime>True</DesignTime>

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

@@ -1,11 +1,11 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using Microsoft.Diagnostics.Tools.Monitor;
+using Microsoft.Diagnostics.Monitoring.StartupHook;
+using Microsoft.Diagnostics.Monitoring.StartupHook.MonitorMessageDispatcher;
 using Microsoft.Diagnostics.Tools.Monitor.Profiler;
 using System;
 using System.Collections.Generic;
-using System.IO;
 using System.Reflection;
 using System.Runtime.InteropServices;
 
@@ -29,40 +29,14 @@ private static extern void RequestFunctionProbeInstallation(
         private readonly object _requestLocker = new();
         private long _disposedState;
 
-        private readonly string? _profilerModulePath;
-
         public FunctionProbesManager(IFunctionProbes probes)
         {
-            _profilerModulePath = Environment.GetEnvironmentVariable(ProfilerIdentifiers.EnvironmentVariables.ModulePath);
-            if (!File.Exists(_profilerModulePath))
-            {
-                throw new FileNotFoundException(_profilerModulePath);
-            }
-
-            NativeLibrary.SetDllImportResolver(typeof(ParameterCapturingService).Assembly, ResolveDllImport);
+            ProfilerResolver.InitializeResolver<FunctionProbesManager>();
 
             RequestFunctionProbeRegistration(FunctionProbesStub.GetProbeFunctionId());
             FunctionProbesStub.Instance = probes;
         }
 
-        private IntPtr ResolveDllImport(string libraryName, Assembly assembly, DllImportSearchPath? searchPath)
-        {
-            // DllImport for Windows automatically loads in-memory modules (such as the profiler). This is not the case for Linux/MacOS.
-            // If we fail resolving the DllImport, we have to load the profiler ourselves.
-            if (_profilerModulePath == null ||
-                libraryName != ProfilerIdentifiers.LibraryRootFileName)
-            {
-                return IntPtr.Zero;
-            }
-
-            if (NativeLibrary.TryLoad(_profilerModulePath, out IntPtr handle))
-            {
-                return handle;
-            }
-
-            return IntPtr.Zero;
-        }
-
         public void StopCapturing()
         {
             lock (_requestLocker)

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

@@ -2,7 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using Microsoft.Diagnostics.Monitoring.HostingStartup.ParameterCapturing.FunctionProbes;
-using Microsoft.Diagnostics.Tools.Monitor;
+using Microsoft.Diagnostics.Monitoring.StartupHook;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;

src/Microsoft.Diagnostics.Monitoring.StartupHook/Microsoft.Diagnostics.Monitoring.StartupHook.csproj

@@ -16,10 +16,13 @@
     <Compile Include="..\Tools\dotnet-monitor\Exceptions\ExceptionEvents.cs" Link="Exceptions\Eventing\ExceptionEvents.cs" />
     <Compile Include="..\Tools\dotnet-monitor\StartupHook\StartupHookIdentifiers.cs" Link="StartupHookIdentifiers.cs" />
     <Compile Include="..\Tools\dotnet-monitor\ToolIdentifiers.cs" Link="ToolIdentifiers.cs" />
+    <Compile Include="..\Tools\dotnet-monitor\Profiler\ProfilerIdentifiers.cs" />
+    <Compile Include="..\Microsoft.Diagnostics.Monitoring.WebApi\ProfilerMessage.cs" Link="MonitorMessageDispatcher\ProfilerMessage.cs" />
   </ItemGroup>
 
   <ItemGroup>
     <InternalsVisibleTo Include="Microsoft.Diagnostics.Monitoring.StartupHook.UnitTests" />
+    <InternalsVisibleTo Include="Microsoft.Diagnostics.Monitoring.HostingStartup" />
   </ItemGroup>
 
 </Project>

src/Microsoft.Diagnostics.Monitoring.StartupHook/MonitorMessageDispatcher/IMonitorMessageSource.cs

@@ -0,0 +1,26 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+
+namespace Microsoft.Diagnostics.Monitoring.StartupHook.MonitorMessageDispatcher
+{
+    internal sealed class MonitorMessageArgs : EventArgs
+    {
+        public MonitorMessageArgs(IpcCommand command, IntPtr nativeBuffer, long bufferSize)
+        {
+            Command = command;
+            NativeBuffer = nativeBuffer;
+            BufferSize = bufferSize;
+        }
+
+        public IpcCommand Command { get; private set; }
+        public IntPtr NativeBuffer { get; private set; }
+        public long BufferSize { get; private set; }
+    }
+
+    internal interface IMonitorMessageSource : IDisposable
+    {
+        public event EventHandler<MonitorMessageArgs> MonitorMessage;
+    }
+}

src/Microsoft.Diagnostics.Monitoring.StartupHook/MonitorMessageDispatcher/MonitorMessageDispatcher.cs

@@ -0,0 +1,96 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+
+using System.IO;
+using System;
+using System.Text.Json;
+using System.Collections.Generic;
+
+namespace Microsoft.Diagnostics.Monitoring.StartupHook.MonitorMessageDispatcher
+{
+    internal sealed class MonitorMessageDispatcher : IDisposable
+    {
+        internal struct MessageDispatchEntry
+        {
+            public Type DeserializeType { get; set; }
+            public Action<object> Callback { get; set; }
+        }
+
+        private readonly object _dispatchTableLocker = new();
+        private readonly Dictionary<IpcCommand, MessageDispatchEntry> _dispatchTable = new();
+
+        private readonly IMonitorMessageSource _messageSource;
+
+        private long _disposedState;
+
+        public MonitorMessageDispatcher(IMonitorMessageSource messageSource)
+        {
+            _messageSource = messageSource;
+            _messageSource.MonitorMessage += OnMessage;
+        }
+
+        public void RegisterCallback<T>(IpcCommand command, Action<T> callback)
+        {
+            MessageDispatchEntry dispatchEntry = new()
+            {
+                DeserializeType = typeof(T),
+                Callback = (obj) =>
+                {
+                    callback((T)obj);
+                }
+            };
+
+            lock (_dispatchTableLocker)
+            {
+                if (!_dispatchTable.TryAdd(command, dispatchEntry))
+                {
+                    throw new InvalidOperationException("Callback already registered for the requested command.");
+                }
+            }
+        }
+
+        public void UnregisterCallback(IpcCommand command)
+        {
+            lock (_dispatchTableLocker)
+            {
+                _dispatchTable.Remove(command, out _);
+            }
+        }
+
+        private void OnMessage(object? sender, MonitorMessageArgs args)
+        {
+            lock (_dispatchTableLocker)
+            {
+                if (!_dispatchTable.TryGetValue(args.Command, out MessageDispatchEntry dispatchEntry))
+                {
+                    throw new NotSupportedException("Unsupported message type.");
+                }
+
+                object? payload = null;
+                unsafe
+                {
+                    using UnmanagedMemoryStream memoryStream = new((byte*)args.NativeBuffer.ToPointer(), args.BufferSize);
+                    // Exceptions thrown during deserialization will be handled by the message source
+                    payload = JsonSerializer.Deserialize(memoryStream, dispatchEntry.DeserializeType);
+                }
+
+                if (payload == null)
+                {
+                    throw new ArgumentException("Could not deserialize.");
+                }
+
+                dispatchEntry.Callback(payload);
+            }
+        }
+
+        public void Dispose()
+        {
+            if (!DisposableHelper.CanDispose(ref _disposedState))
+                return;
+
+            _messageSource.MonitorMessage -= OnMessage;
+            _messageSource.Dispose();
+        }
+    }
+}

src/Microsoft.Diagnostics.Monitoring.StartupHook/MonitorMessageDispatcher/ProfilerMessageSource.cs

@@ -0,0 +1,75 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+
+using System.Runtime.InteropServices;
+using System;
+using Microsoft.Diagnostics.Tools.Monitor.Profiler;
+
+namespace Microsoft.Diagnostics.Monitoring.StartupHook.MonitorMessageDispatcher
+{
+    internal sealed class ProfilerMessageSource : IMonitorMessageSource
+    {
+        public event EventHandler<MonitorMessageArgs>? MonitorMessage;
+
+        public delegate int ProfilerMessageCallback(IpcCommand command, IntPtr nativeBuffer, long bufferSize);
+
+        [DllImport(ProfilerIdentifiers.LibraryRootFileName, CallingConvention = CallingConvention.StdCall, PreserveSig = false)]
+        private static extern void RegisterMonitorMessageCallback(ProfilerMessageCallback callback);
+
+        [DllImport(ProfilerIdentifiers.LibraryRootFileName, CallingConvention = CallingConvention.StdCall, PreserveSig = false)]
+        private static extern void UnregisterMonitorMessageCallback();
+
+        private long _disposedState;
+
+        public ProfilerMessageSource()
+        {
+            ProfilerResolver.InitializeResolver<ProfilerMessageSource>();
+            RegisterMonitorMessageCallback(OnProfilerMessage);
+        }
+
+        private void RaiseMonitorMessage(MonitorMessageArgs e)
+        {
+            MonitorMessage?.Invoke(this, e);
+        }
+
+        private int OnProfilerMessage(IpcCommand command, IntPtr nativeBuffer, long bufferSize)
+        {
+            try
+            {
+                if (bufferSize == 0)
+                {
+                    throw new ArgumentException(nameof(bufferSize));
+                }
+
+                if (nativeBuffer == IntPtr.Zero)
+                {
+                    throw new ArgumentException(nameof(nativeBuffer));
+                }
+
+                RaiseMonitorMessage(new MonitorMessageArgs(command, nativeBuffer, bufferSize));
+            }
+            catch (Exception ex)
+            {
+                return Marshal.GetHRForException(ex);
+            }
+
+            return 0;
+        }
+
+        public void Dispose()
+        {
+            if (!DisposableHelper.CanDispose(ref _disposedState))
+                return;
+
+            try
+            {
+                UnregisterMonitorMessageCallback();
+            }
+            catch
+            {
+
+            }
+        }
+    }
+}

src/Microsoft.Diagnostics.Monitoring.StartupHook/ProfilerResolver.cs

@@ -0,0 +1,64 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+
+using System.Runtime.InteropServices;
+using System.IO;
+using System.Reflection;
+using System;
+using Microsoft.Diagnostics.Tools.Monitor.Profiler;
+using System.Collections.Generic;
+
+namespace Microsoft.Diagnostics.Monitoring.StartupHook
+{
+    internal static class ProfilerResolver
+    {
+        private static readonly Lazy<string> s_profilerModulePath = new Lazy<string>(() => Environment.GetEnvironmentVariable(ProfilerIdentifiers.EnvironmentVariables.ModulePath) ?? string.Empty);
+        private static readonly Lazy<bool> s_profilerModuleExists = new Lazy<bool>(() => File.Exists(s_profilerModulePath.Value));
+
+        private static readonly HashSet<Assembly> s_registeredAssemblies = new HashSet<Assembly>();
+        private static readonly object s_registeredAssembliesLocker = new();
+
+        public static void InitializeResolver(Type type)
+        {
+            if (!s_profilerModuleExists.Value)
+            {
+                throw new FileNotFoundException(s_profilerModulePath.Value);
+            }
+
+            Assembly assembly = type.Assembly;
+            lock (s_registeredAssembliesLocker)
+            {
+                if (!s_registeredAssemblies.Add(assembly))
+                {
+                    return;
+                }
+
+                NativeLibrary.SetDllImportResolver(assembly, ResolveProfilerDllImport);
+            }
+        }
+
+        public static void InitializeResolver<T>()
+        {
+            InitializeResolver(typeof(T));
+        }
+
+        private static IntPtr ResolveProfilerDllImport(string libraryName, Assembly assembly, DllImportSearchPath? searchPath)
+        {
+            // DllImport for Windows automatically loads in-memory modules (such as the profiler). This is not the case for Linux/MacOS.
+            // If we fail resolving the DllImport, we have to load the profiler ourselves.
+            if (s_profilerModulePath.Value == null ||
+                libraryName != ProfilerIdentifiers.LibraryRootFileName)
+            {
+                return IntPtr.Zero;
+            }
+
+            if (NativeLibrary.TryLoad(s_profilerModulePath.Value, out IntPtr handle))
+            {
+                return handle;
+            }
+
+            return IntPtr.Zero;
+        }
+    }
+}

src/Microsoft.Diagnostics.Monitoring.StartupHook/SharedInternals.cs

@@ -0,0 +1,12 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Microsoft.Diagnostics.Monitoring.StartupHook.MonitorMessageDispatcher;
+
+namespace Microsoft.Diagnostics.Tools.Monitor.StartupHook
+{
+    internal static class SharedInternals
+    {
+        public static MonitorMessageDispatcher? MessageDispatcher { get; set; }
+    }
+}