Fix sampler and EventSource

src/Ultra.Core/UltraProfilerEventPipe.cs

@@ -5,6 +5,7 @@
 using System.Diagnostics;
 using System.Diagnostics.Tracing;
 using System.IO.Enumeration;
+using System.Net.Sockets;
 using ByteSizeLib;
 using Microsoft.Diagnostics.NETCore.Client;
 using Microsoft.Diagnostics.Tracing;
@@ -117,7 +118,7 @@ private static void SetupUltraSampler(ProcessStartInfo startInfo)
             value = ultraSamplerPath;
         }
 
-        Console.WriteLine($"DYLD_INSERT_LIBRARIES={value}");
+        //Console.WriteLine($"DYLD_INSERT_LIBRARIES={value}");
         startInfo.Environment[key] = value;
     }
 
@@ -156,6 +157,7 @@ private UltraSamplerProfilerState(UltraProfilerOptions options, DiagnosticsClien
         public static async Task<UltraSamplerProfilerState> Connect(string baseName, int pid, CancellationToken token, UltraProfilerOptions options)
         {
 
+            //var ultraTempFolder = Path.Combine(Path.GetTempPath(), ".ultra");
             var ultraTempFolder = Path.Combine(Path.GetTempPath(), ".ultra");
 
             var pattern = $"dotnet-diagnostic-{pid}-*";
@@ -189,15 +191,16 @@ public static async Task<UltraSamplerProfilerState> Connect(string baseName, int
             }
 
             var diagnosticClientMain = new DiagnosticsClient(pid);
-            var diagnosticClientUltra = await DiagnosticsClientConnector.FromDiagnosticPort(ultraDiagnosticPortSocket, token);
+            //DiagnosticsClient? diagnosticClientMain = null;
+            var diagnosticClientUltra = (await DiagnosticsClientConnector.FromDiagnosticPort(ultraDiagnosticPortSocket, token).ConfigureAwait(false))!.Instance;
 
             var timeoutSource = new CancellationTokenSource();
             var linkedCancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(token, timeoutSource.Token);
 
             try
             {
                 timeoutSource.CancelAfter(1000);
-                await diagnosticClientUltra!.Instance.WaitForConnectionAsync(linkedCancellationTokenSource.Token).ConfigureAwait(false);
+                await diagnosticClientUltra!.WaitForConnectionAsync(linkedCancellationTokenSource.Token).ConfigureAwait(false);
                 await diagnosticClientMain.WaitForConnectionAsync(linkedCancellationTokenSource.Token).ConfigureAwait(false);
             }
             catch (OperationCanceledException) when (timeoutSource.IsCancellationRequested)
@@ -209,7 +212,7 @@ public static async Task<UltraSamplerProfilerState> Connect(string baseName, int
                 throw;
             }
 
-            return new UltraSamplerProfilerState(options, diagnosticClientMain, diagnosticClientUltra.Instance, baseName, pid, token);
+            return new UltraSamplerProfilerState(options, diagnosticClientMain, diagnosticClientUltra, baseName, pid, token);
         }
 
         public long TotalFileLength()
@@ -228,7 +231,8 @@ public long TotalFileLength()
         public async Task StartProfiling()
         {
             var ultraEventProvider = new EventPipeProvider(UltraSamplerParser.Name, EventLevel.Verbose);
-            _ultraSession = await _ultraDiagnosticsClient.StartEventPipeSessionAsync([ultraEventProvider], true, 256, _token).ConfigureAwait(false);
+            var config = new EventPipeSessionConfiguration([ultraEventProvider], 256, false, true);
+            _ultraSession = await _ultraDiagnosticsClient.StartEventPipeSessionAsync(config, _token).ConfigureAwait(false);
             _ultraEventStreamCopyTask = _ultraSession.EventStream.CopyToAsync(_ultraNetTraceFileStream, _token);
 
             if (_mainDiagnosticsClient is not null)
@@ -264,14 +268,44 @@ public async Task Stop()
                 await _mainEventStreamCopyTask.ConfigureAwait(false);
             }
 
-            if (_ultraSession is not null)
+            await SafeStopAsync(_ultraSession, _token);
+            _ultraSession = null;
+
+            await SafeStopAsync(_mainSession, _token);
+            _mainSession = null;
+        }
+
+        private static async Task SafeStopAsync(EventPipeSession? session, CancellationToken token)
+        {
+            if (session is null) return;
+
+            try
             {
-                await _ultraSession.StopAsync(_token).ConfigureAwait(false);
+                await session.StopAsync(token).ConfigureAwait(false);
             }
+            catch (EndOfStreamException)
+            {
 
-            if (_mainSession is not null)
+            }
+            catch (TimeoutException)
+            {
+
+            }
+            catch (OperationCanceledException)
             {
-                await _mainSession.StopAsync(_token).ConfigureAwait(false);
+
+            }
+            catch (PlatformNotSupportedException)
+            {
+
+            }
+            catch (ServerNotAvailableException)
+            {
+
+            }
+            catch (SocketException)
+            {
+
             }
         }
 

src/Ultra.Sampler/MacOS/MacOSLibSystem.cs

@@ -343,8 +343,38 @@ public struct uuid_command {
         public uint	cmd;		/* LC_UUID */
         public uint	cmdsize;	/* sizeof(struct uuid_command) */
         public Guid	uuid;	  /* the 128-bit uuid */
-    };
+    }
+
+    public unsafe struct segment_command_64 { /* for 64-bit architectures */
+        public uint	cmd;		/* LC_SEGMENT_64 */
+        public uint	cmdsize;	/* includes sizeof section_64 structs */
+        public fixed byte segname[16];	/* segment name */
+        public ulong	vmaddr;		/* memory address of this segment */
+        public ulong	vmsize;		/* memory size of this segment */
+        public ulong	fileoff;	/* file offset of this segment */
+        public ulong	filesize;	/* amount to map from the file */
+        public int		maxprot;	/* maximum VM protection */
+        public int		initprot;	/* initial VM protection */
+        public uint	nsects;		/* number of sections in segment */
+        public uint	flags;		/* flags */
+    }
+
+    public unsafe struct section_64 { /* for 64-bit architectures */
+        public fixed byte		sectname[16];	/* name of this section */
+        public fixed byte		segname[16];	/* segment this section goes in */
+        public ulong	addr;		/* memory address of this section */
+        public ulong	size;		/* size in bytes of this section */
+        public uint	offset;		/* file offset of this section */
+        public uint	align;		/* section alignment (power of 2) */
+        public uint	reloff;		/* file offset of relocation entries */
+        public uint	nreloc;		/* number of relocation entries */
+        public uint	flags;		/* flags (section type and attributes)*/
+        public uint	reserved1;	/* reserved (for offset or index) */
+        public uint	reserved2;	/* reserved (for count or sizeof) */
+        public uint	reserved3;	/* reserved */
+    }
 
     public const uint LC_UUID = 0x1b;
+    public const uint LC_SEGMENT_64 = 0x19;
 
 }

src/Ultra.Sampler/MacOS/MacOSUltraSampler.cs

@@ -2,7 +2,9 @@
 // Licensed under the BSD-Clause 2 license.
 // See license.txt file in the project root for full license information.
 
+using System.Collections;
 using System.Diagnostics;
+using System.Diagnostics.Tracing;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 using XenoAtom.Collections;
@@ -24,6 +26,7 @@ internal unsafe class MacOSUltraSampler : UltraSampler
     private bool _initializingModules;
     private readonly object _moduleEventLock = new();
     private int _nextModuleEventIndexToLog;
+    private readonly UltraSamplerSource _samplerEventSource;
 
     private readonly MacOSLibSystem.dyld_register_callback _callbackDyldAdded;
     private readonly MacOSLibSystem.dyld_register_callback _callbackDyldRemoved;
@@ -40,6 +43,9 @@ public MacOSUltraSampler()
         MacOSLibSystem._dyld_register_func_for_add_image(Marshal.GetFunctionPointerForDelegate(_callbackDyldAdded));
         _initializingModules = false;
         MacOSLibSystem._dyld_register_func_for_remove_image(Marshal.GetFunctionPointerForDelegate(_callbackDyldRemoved));
+
+        // Make sure to use the instance to trigger the constructor of the EventSource so that it is registered in the runtime!
+        _samplerEventSource = UltraSamplerSource.Log;
     }
 
     protected override void StartImpl()
@@ -140,6 +146,7 @@ private void AddModuleEvent(NativeModuleEventKind kind, nint loadAddress)
         var path = MemoryMarshal.CreateReadOnlySpanFromNullTerminated((byte*)info.dli_fname);
         evt.Path = path.ToArray();
         evt.TimestampUtc = DateTime.UtcNow;
+        evt.Size = GetMaximumCodeAddress(loadAddress);
 
         lock (_moduleEventLock)
         {
@@ -163,7 +170,7 @@ private void NotifyPendingNativeModuleEvents()
             for(; _nextModuleEventIndexToLog < events.Length; _nextModuleEventIndexToLog++)
             {
                 var evt = events[_nextModuleEventIndexToLog];
-                UltraSamplerSource.Log.OnNativeModuleEvent(evt.Kind, evt.LoadAddress, evt.Path, evt.TimestampUtc.Ticks);
+                UltraSamplerSource.Log.OnNativeModuleEvent(evt.Kind, evt.LoadAddress, evt.Size, evt.Path, evt.TimestampUtc.Ticks);
             }
         }
     }
@@ -190,6 +197,41 @@ private static bool TryGetUuidFromMacHeader(nint headerPtr, out Guid guid)
         return false;
     }
 
+    private static ulong GetMaximumCodeAddress(nint headerPtr)
+    {
+        ulong startAddress = 0;
+
+        ulong size = 0;
+        var header = (MacOSLibSystem.mach_header_64*)headerPtr;
+        if (header->magic != MacOSLibSystem.MH_MAGIC_64) throw new InvalidOperationException("Invalid magic header");
+
+        var nbCommands = header->ncmds;
+        var commands = (MacOSLibSystem.load_command*)((byte*)header + sizeof(MacOSLibSystem.mach_header_64));
+        for(uint i = 0; i < nbCommands; i++)
+        {
+            ref var command = ref commands[i];
+            if (command.cmd == MacOSLibSystem.LC_SEGMENT_64)
+            {
+                ref var segment = ref Unsafe.As<MacOSLibSystem.load_command,MacOSLibSystem.segment_command_64>(ref command);
+                if (segment.vmaddr != 0)
+                {
+                    if (startAddress == 0)
+                    {
+                        startAddress = segment.vmaddr;
+                    }
+
+                    var newSize = (ulong)((long)segment.vmaddr + (long)segment.vmsize - (long)startAddress);
+                    if (newSize > size)
+                    {
+                        size = newSize;
+                    }
+                }
+            }
+        }
+
+        return size;
+    }
+
     public void Sample(NativeCallstackDelegate nativeCallstack)
     {
         MacOS.MacOSLibSystem.task_for_pid(MacOS.MacOSLibSystem.mach_task_self(), Process.GetCurrentProcess().Id, out var rootTask)
@@ -251,7 +293,7 @@ private static unsafe void Sample(MacOS.MacOSLibSystem.mach_port_t rootTask, ulo
 
                     //Console.WriteLine($"sp: 0x{armThreadState.__sp:X8}, fp: 0x{armThreadState.__fp:X8}, lr: 0x{armThreadState.__lr:X8}");
                     int frameCount = WalkNativeCallStack(armThreadState.__sp, armThreadState.__fp, armThreadState.__lr, pFrames);
-                    nativeCallstack(threadInfo.thread_id, pFrames, frameCount);
+                    nativeCallstack(threadInfo.thread_id, (nint)pFrames, frameCount);
                 }
                 finally
                 {

src/Ultra.Sampler/MacOS/NativeCallstackDelegate.cs

@@ -4,4 +4,4 @@
 
 namespace Ultra.Sampler.MacOS;
 
-public unsafe delegate void NativeCallstackDelegate(ulong threadId, ulong* pFrames, int frameCount);
+public unsafe delegate void NativeCallstackDelegate(ulong threadId, nint pFrames, int frameCount);

src/Ultra.Sampler/MacOS/NativeModuleEvent.cs

@@ -10,11 +10,12 @@ internal struct NativeModuleEvent
 {
     public NativeModuleEventKind Kind;
     public ulong LoadAddress;
+    public ulong Size;
     public byte[]? Path;
     public DateTime TimestampUtc;
 
     public override string ToString()
     {
-        return $"{nameof(LoadAddress)}: 0x{LoadAddress:X8}, {nameof(Path)}: {Encoding.UTF8.GetString(Path ?? [])}, {nameof(TimestampUtc)}: {TimestampUtc:O}";
+        return $"{nameof(LoadAddress)}: 0x{LoadAddress:X8}, {nameof(Size)}: {Size}, {nameof(Path)}: {Encoding.UTF8.GetString(Path ?? [])}, {nameof(TimestampUtc)}: {TimestampUtc:O}";
     }
 }

src/Ultra.Sampler/UltraSamplerParser.cs

@@ -27,10 +27,4 @@ public static class UltraSamplerParser
     public const int NativeCallStackEvent = 1;
 
     public const int NativeModuleEvent = 2;
-
-
-    public static void TestId(Guid guid)
-    {
-        guid = Id;
-    }
 }

src/Ultra.Sampler/UltraSamplerSource.cs

@@ -2,6 +2,7 @@
 // Licensed under the BSD-Clause 2 license.
 // See license.txt file in the project root for full license information.
 
+using System.Diagnostics.CodeAnalysis;
 using System.Diagnostics.Tracing;
 using System.Runtime.CompilerServices;
 using Ultra.Sampler.MacOS;
@@ -17,38 +18,40 @@ private UltraSamplerSource()
     {
     }
 
-    [Event(UltraSamplerParser.NativeCallStackEvent, Level = EventLevel.Verbose, Message = "NativeCallstackEvent Thread {0} with {2} frames")]
-    [SkipLocalsInit]
-    public unsafe void OnNativeCallstack(ulong threadId, ulong* pFrames, int count)
+    [Event(UltraSamplerParser.NativeCallStackEvent, Level = EventLevel.Informational)]
+    [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access otherwise can break functionality when trimming application code", Justification = "<Pending>")]
+    public unsafe void OnNativeCallstack(ulong threadId, nint pFrames, int count)
     {
-
-        Unsafe.SkipInit(out EventData2 evt);
+        EventData2 evt = default;
         evt.Data1.DataPointer = (nint)(void*)&threadId;
         evt.Data1.Size = sizeof(ulong);
         evt.Data2.DataPointer = (nint)pFrames;
         evt.Data2.Size = count * sizeof(ulong);
         WriteEventCore(UltraSamplerParser.NativeCallStackEvent, 2, &evt.Data1);
     }
 
-    [Event(UltraSamplerParser.NativeModuleEvent, Level = EventLevel.Verbose, Message = "NativeModuleEvent {0} LoadAddress: {1}")]
-    [SkipLocalsInit]
-    public unsafe void OnNativeModuleEvent(NativeModuleEventKind nativeModuleEventKind, ulong loadAddress, byte[]? modulePathUtf8, long timestampUtc)
+    [Event(UltraSamplerParser.NativeModuleEvent, Level = EventLevel.Informational)]
+    [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access otherwise can break functionality when trimming application code", Justification = "<Pending>")]
+    public unsafe void OnNativeModuleEvent(NativeModuleEventKind nativeModuleEventKind, ulong loadAddress, ulong size, byte[]? modulePathUtf8, long timestampUtc)
     {
-        Unsafe.SkipInit(out EventData4 evt);
+        EventData5 evt = default;
         evt.Data1.DataPointer = (nint)(void*)&nativeModuleEventKind;
         evt.Data1.Size = sizeof(int);
         evt.Data2.DataPointer = (nint)(void*)&loadAddress;
         evt.Data2.Size = sizeof(ulong);
+        evt.Data3.DataPointer = (nint)(void*)&size;
+        evt.Data3.Size = sizeof(ulong);
         fixed (byte* evtPathPtr = modulePathUtf8)
         {
-            evt.Data3.DataPointer = (nint)evtPathPtr;
-            evt.Data3.Size = modulePathUtf8?.Length ?? 0;
-            evt.Data4.DataPointer = (nint)(void*)&timestampUtc;
-            evt.Data4.Size = sizeof(long);
+            evt.Data4.DataPointer = (nint)evtPathPtr;
+            evt.Data4.Size = modulePathUtf8?.Length ?? 0;
+            evt.Data5.DataPointer = (nint)(void*)&timestampUtc;
+            evt.Data5.Size = sizeof(long);
             WriteEventCore(UltraSamplerParser.NativeModuleEvent, 4, &evt.Data1);
         }
     }
 
+    [NonEvent]
     protected override void OnEventCommand(EventCommandEventArgs command)
     {
         if (command.Command == EventCommand.Enable)
@@ -58,6 +61,9 @@ protected override void OnEventCommand(EventCommandEventArgs command)
         else if (command.Command == EventCommand.Disable)
         {
             UltraSampler.Instance.Disable();
+
+            // Wait a bit to let the sampler thread finishing
+            Thread.Sleep(100);
         }
     }
 
@@ -68,7 +74,7 @@ private struct EventData2
         public EventData Data2;
     }
 
-    private struct EventData4
+    private struct EventData5
     {
         public EventData Data1;
 
@@ -77,5 +83,7 @@ private struct EventData4
         public EventData Data3;
 
         public EventData Data4;
+
+        public EventData Data5;
     }
 }

src/Ultra.Sampler/ultra_sampler_indirect.cpp

@@ -45,7 +45,7 @@ extern "C" {
         free(new_tmpdir);
 
         pid_t pid = getpid();
-        printf("Current Process pid: %d tmpdir: %s\n", pid, getenv("TMPDIR"));
+        //printf("Current Process pid: %d tmpdir: %s\n", pid, getenv("TMPDIR"));
 
         // Call the arbitrary function
         func();