JitProfiler.cs

using System.Runtime.CompilerServices;
using Silhouette;
using System.Runtime.InteropServices;
using System.Xml.Linq;

namespace JitProfiler;

internal unsafe class CorProfiler : CorProfilerCallback4Base
{
    private const string FileName = @"G:\JetBrains\jit_profiler.txt";

    protected override HResult Initialize(int iCorProfilerInfoVersion)
    {
        File.Delete(FileName);
            
        var result = ICorProfilerInfo2.SetEventMask(COR_PRF_MONITOR.COR_PRF_ENABLE_STACK_SNAPSHOT | COR_PRF_MONITOR.COR_PRF_MONITOR_JIT_COMPILATION);
        return result;
    }

    protected override HResult JITCompilationStarted(FunctionId functionId, bool fIsSafeToBlock)
    {
        var typeArgs = Array.Empty<ClassId>();

        var result = ICorProfilerInfo2.GetFunctionInfo2(functionId, default, out var classId, out var moduleId, out var token, 0, out var pcTypeArgs, null);

        if (!result)
        {
            Logger.Log($"GetFunctionInfo2 failed: {HResult.ToString(result)}");
            return result;
        }

        var output = new XDocument();
        var root = new XElement("item");
        output.Add(root);

        var function = new XElement("function");
        root.Add(function);
            
        if (pcTypeArgs > 0)
        {
            typeArgs = new ClassId[pcTypeArgs];

            fixed (ClassId* typeArgsPtr = typeArgs)
            {
                result = ICorProfilerInfo2.GetFunctionInfo2(functionId, default, out classId, out moduleId, out token, pcTypeArgs, out pcTypeArgs, typeArgsPtr);

                if (!result)
                {
                    Logger.Log($"GetFunctionInfo2 failed (2nd call): {HResult.ToString(result)}");
                    return result;
                }
            }
        }

        var functionName = GetFunctionName(moduleId, token);
        function.Add(new XElement("name", functionName));

        var functionType = GetTypeName(classId);
            
        if (functionType.Element("name")?.Value.StartsWith("JetBrains") != true)
        {
            return HResult.S_OK;
        }
            
        function.Add(functionType);
            
        if (typeArgs.Length > 0)
        {
            var genericArguments = new XElement("genericArguments");
            function.Add(genericArguments);
                
            foreach (var typeArg in typeArgs)
            {
                genericArguments.Add(GetTypeName(typeArg));
            }
        }

        root.Add(WalkStack());

        lock (this)
        {
            File.AppendAllText(FileName, output + "\r\n");
        }

        return HResult.S_OK;
    }
        
    private XElement WalkStack()
    {
        var element = new XElement("stacktrace");
            
        StackWalkContext context = default;
        ICorProfilerInfo2.DoStackSnapshot(default, &WalkThread, COR_PRF_SNAPSHOT_INFO.COR_PRF_SNAPSHOT_DEFAULT, &context, null, 0);

        for (int i = 0; i < context.Count; i++)
        {
            string functionName = string.Empty;

            try
            {
                var ip = context.Frames[i];
                var functionId = ICorProfilerInfo2.GetFunctionFromIP(ip).ThrowIfFailed();

                var result = ICorProfilerInfo2.GetFunctionInfo2(functionId, default, out _, out var moduleId, out var token, 0, out _, null);
                if (result)
                {
                    functionName = GetFullFunctionName(moduleId, token);
                }
            }
            catch (Exception ex)
            {
                Logger.Log($"WalkStack failed: {ex}");
            }
                
            if (!string.IsNullOrEmpty(functionName))
            {
                element.Add(new XElement("frame", functionName));
            }
        }

        return element;
    }

    [UnmanagedCallersOnly(CallConvs = [typeof(CallConvStdcall)])]
    private static unsafe HResult WalkThread(FunctionId funcId, nint ip, COR_PRF_FRAME_INFO frameInfo, uint contextSize, byte* context, void* clientData)
    {
        ref StackWalkContext stackWalkContext = ref *(StackWalkContext*)clientData;

        if (stackWalkContext.Count >= StackWalkContext.MaxFrames)
        {
            return HResult.E_ABORT;
        }

        stackWalkContext.Frames[stackWalkContext.Count++] = ip;
        return HResult.S_OK;
    }

    private struct StackWalkContext
    {
        public const int MaxFrames = 1024;
        public FramesArray Frames;
        public int Count;

        [InlineArray(MaxFrames)]
        public struct FramesArray
        {
            private nint _frames;
        }
    }

    private XElement GetTypeName(ClassId classId)
    {
        var element = new XElement("type");

        try
        {   
            Span<ClassId> typeArgs = [];

            var classIdInfo = ICorProfilerInfo2.GetClassIDInfo2(classId, typeArgs, out var numTypeArgs).ThrowIfFailed();

            if (numTypeArgs > 0)
            {
                typeArgs = stackalloc ClassId[(int)numTypeArgs];

                classIdInfo = ICorProfilerInfo2.GetClassIDInfo2(classId, typeArgs, out numTypeArgs).ThrowIfFailed();
            }

            var moduleInfo = ICorProfilerInfo2.GetModuleInfo(classIdInfo.ModuleId).ThrowIfFailed();
            element.Add(new XElement("module", moduleInfo.ModuleName));

            using var metaDataImport = ICorProfilerInfo2.GetModuleMetaData(classIdInfo.ModuleId, CorOpenFlags.ofRead, KnownGuids.IMetaDataImport).ThrowIfFailed().Wrap();
            var typeDefProps = metaDataImport.Value.GetTypeDefProps(classIdInfo.TypeDef).ThrowIfFailed();                
                
            element.Add(new XElement("name", typeDefProps.TypeName));

            if ((typeDefProps.TypeDefFlags & 0x7) /* tdVisibilityMask */ > 0x2 /* tdNestedPublic */)
            {
                var enclosingTypeDef = metaDataImport.Value.GetNestedClassProps(classIdInfo.TypeDef).ThrowIfFailed();
                var enclosingTypeDefProps = metaDataImport.Value.GetTypeDefProps(enclosingTypeDef).ThrowIfFailed();
                element.Add(new XElement("parent", enclosingTypeDefProps.TypeName));
            }

            if (numTypeArgs > 0)
            {
                var genericArguments = new XElement("genericArguments");
                element.Add(genericArguments);

                foreach (var typeArg in typeArgs)
                {
                    genericArguments.Add(GetTypeName(typeArg));
                }
            }
        }
        catch (Exception ex)
        {
            Logger.Log($"GetTypeName failed: {ex}");
        }

        return element;
    }

    private string GetFunctionName(ModuleId moduleId, MdToken token)
    {
        try
        {
            using var metaDataImport = ICorProfilerInfo2.GetModuleMetaData(moduleId, CorOpenFlags.ofRead, KnownGuids.IMetaDataImport).ThrowIfFailed().Wrap();
            var methodProperties = metaDataImport.Value.GetMethodProps(new(token)).ThrowIfFailed();

            return methodProperties.Name;
        }
        catch (Exception ex)
        {
            Logger.Log($"GetFunctionName failed: {ex}");
            return string.Empty;
        }
    }

    private string GetFullFunctionName(ModuleId moduleId, MdToken token)
    {
        try
        {
            using var metaDataImport = ICorProfilerInfo2.GetModuleMetaData(moduleId, CorOpenFlags.ofRead, KnownGuids.IMetaDataImport).ThrowIfFailed().Wrap();
            var methodProperties = metaDataImport.Value.GetMethodProps(new(token)).ThrowIfFailed();
            var typeDefProps = metaDataImport.Value.GetTypeDefProps(methodProperties.Class).ThrowIfFailed();

            return $"{typeDefProps.TypeName}.{methodProperties.Name}";
        }
        catch (Exception ex)
        {
            Logger.Log($"GetFunctionName failed: {ex}");
            return string.Empty;
        }
    }
}