natsdotnet/tools/DtpSnapshotExtractor/Program.cs

using System.Text.Json;
using JetBrains.Common.CallTreeStorage.Dfs;
using JetBrains.Common.Util.Metadata;
using JetBrains.DotTrace.Dal.Performance.RemoteTree.CoreLogic;
using JetBrains.DotTrace.Dal.Performance.SnapshotComponents.BothSides;
using JetBrains.DotTrace.Dal.Performance.SnapshotComponents.Remote;
using JetBrains.DotTrace.Dal.Performance.SnapshotStorage;
using JetBrains.DotTrace.Dal.Timeline.SnapshotComponents.Remote;
using JetBrains.DotTrace.DalInterface.Performance.CallTree;
using JetBrains.DotTrace.DalInterface.Performance.Fuids;
using JetBrains.DotTrace.DataStructures.CallTree;
using JetBrains.DotTrace.DataStructures.Payloads;
using JetBrains.Lifetimes;
using JetBrains.Metadata.Access;
using JetBrains.Profiler.Snapshot.Converters;
using JetBrains.Profiler.Snapshot.Interface.Section;
using JetBrains.Profiler.Snapshot.Interface.Section.Metadata;
using JetBrains.Profiler.Snapshot.Interface.Section.Metadata.Helpers;
using JetBrains.Profiler.Snapshot.Section.Metadata;
using JetBrains.Util;
using JetBrains.Util.Logging;

return await ProgramMain(args);

static async Task<int> ProgramMain(string[] args)
{
    try
    {
        if (args.Length == 0)
        {
            Console.Error.WriteLine("Usage: DtpSnapshotExtractor <snapshot.dtp>");
            return 2;
        }

        var snapshotPath = Path.GetFullPath(args[0]);
        if (!File.Exists(snapshotPath))
        {
            Console.Error.WriteLine($"Snapshot not found: {snapshotPath}");
            return 2;
        }

        var dotTraceAppDir = AppContext.GetData("APP_CONTEXT_DEPS_FILES") is not null
            ? Environment.GetEnvironmentVariable("DOTTRACE_APP_DIR")
            : Environment.GetEnvironmentVariable("DOTTRACE_APP_DIR");
        RegisterAssemblyResolver(dotTraceAppDir);

        using var lifetimeDef = Lifetime.Define();
        var lifetime = lifetimeDef.Lifetime;
        var logger = Logger.GetLogger("DtpSnapshotExtractor");
        var masks = new SnapshotMasksComponent();
        var snapshotFile = FileSystemPath.Parse(snapshotPath);
        var container = new SnapshotStorageContainer(lifetime, logger, masks, snapshotFile);
        var callTreeSections = new CallTreeSections(container, masks);
        var headerSections = new HeaderSections(container);
        var headerIndexSections = new HeaderIndexSections(container, masks);
        var headerCombinedSections = new HeaderCombinedSections(headerIndexSections, headerSections, masks);
        var singles = new HeaderIndexOptSections(container, masks);
        var accessor = new CallTreeAndIndexesSnapshotAccessor(callTreeSections, headerCombinedSections, singles);
        var dfsReaders = new DotTraceDfsReaders(accessor);
        var totalPayloadReader = new NodePayloadBatchReaderTotal(accessor);
        var ownPayloadReader = new NodePayloadBatchReaderOwn(accessor);

        var metadataSection = (MetadataSection)new MetadataSectionContainer(container).Get(lifetime);
        var fuidConverter = new FuidToMetadataSectionContainer(container).Get();
        metadataSection.SetFuidConverter(fuidConverter);
        var assemblyProvider = metadataSection.GetMetadataAssemblyProvider(lifetime, loadReferencedAssemblies: false, guardMultiThreading: true);
        var resolver = new SignatureResolver(masks, fuidConverter, assemblyProvider);

        var headers = callTreeSections.AllHeaders().ToArray();
        var totalNodeCount = headers.Sum(header =>
            checked((int)((header.HeaderFull.SectionSize - header.HeaderFull.SectionHeaderSize) / header.HeaderFull.RecordSize())));

        var nodes = new DfsNode<CallTreeSectionOffset, FunctionUID>[totalNodeCount];
        var totals = new DotTracePayload[totalNodeCount];
        var owns = new DotTracePayload[totalNodeCount];

        var readNodes = dfsReaders.GetNodesReaders(lifetime).ReadForwardOffsetsAscending(accessor.MinOffset, totalNodeCount, nodes, 0);
        var readTotals = totalPayloadReader.ReadForwardOffsetsAscending(accessor.MinOffset, totalNodeCount, totals, 0);
        var readOwns = ownPayloadReader.ReadForwardOffsetsAscending(accessor.MinOffset, totalNodeCount, owns, 0);

        if (readNodes != totalNodeCount || readTotals != totalNodeCount || readOwns != totalNodeCount)
            throw new InvalidOperationException($"Snapshot read mismatch. nodes={readNodes}, totals={readTotals}, owns={readOwns}, expected={totalNodeCount}.");

        var nodeMap = new Dictionary<CallTreeSectionOffset, MutableNode>(totalNodeCount);
        foreach (var index in Enumerable.Range(0, totalNodeCount))
        {
            var node = nodes[index];
            var signature = resolver.Resolve(node.Key);
            nodeMap[node.Offset] = new MutableNode
            {
                Id = node.Offset.ToString(),
                Name = signature.Name,
                Kind = signature.Kind,
                InclusiveTime = GetTotalTime(totals[index]),
                ExclusiveTime = GetTotalTime(owns[index]),
                CallCount = GetCallCount(totals[index]),
                Children = []
            };
        }

        foreach (var index in Enumerable.Range(0, totalNodeCount))
        {
            var node = nodes[index];
            if (!node.ParentOffset.IsValid)
                continue;

            if (nodeMap.TryGetValue(node.ParentOffset, out var parent))
                parent.Children.Add(nodeMap[node.Offset]);
        }

        var threadNodes = new List<SerializableNode>(headers.Length);
        foreach (var header in headers)
        {
            if (!nodeMap.TryGetValue(header.Root, out var rootNode))
                continue;

            var threadName = string.IsNullOrWhiteSpace(header.ThreadName)
                ? (string.IsNullOrWhiteSpace(header.GroupName) ? $"Thread {header.HeaderFull.GroupId}" : header.GroupName)
                : header.ThreadName;

            threadNodes.Add(new SerializableNode
            {
                Id = $"thread:{header.HeaderFull.GroupId}:{rootNode.Id}",
                Name = threadName,
                Kind = "thread",
                ThreadName = threadName,
                InclusiveTime = rootNode.InclusiveTime,
                ExclusiveTime = 0,
                CallCount = rootNode.CallCount,
                Children = [CloneTree(header.Root, nodeMap, [])]
            });
        }

        var syntheticRoot = new SerializableNode
        {
            Id = "root",
            Name = "<root>",
            Kind = "root",
            ThreadName = null,
            InclusiveTime = threadNodes.Sum(node => node.InclusiveTime),
            ExclusiveTime = 0,
            CallCount = threadNodes.Sum(node => node.CallCount),
            Children = threadNodes
        };

        var hotspots = BuildHotspots(nodeMap.Values);
        var payload = new OutputDocument
        {
            Snapshot = new SnapshotInfo
            {
                Path = snapshotPath,
                PayloadType = "time",
                ThreadCount = threadNodes.Count,
                NodeCount = nodeMap.Count
            },
            ThreadRoots = threadNodes.Select(node => new ThreadRootInfo
            {
                Id = node.Id,
                Name = node.Name,
                InclusiveTime = node.InclusiveTime
            }).ToList(),
            Hotspots = hotspots,
            CallTree = syntheticRoot
        };

        await JsonSerializer.SerializeAsync(Console.OpenStandardOutput(), payload, CreateJsonOptions());
        await Console.Out.WriteLineAsync();
        return 0;
    }
    catch (Exception exception)
    {
        Console.Error.WriteLine(exception);
        return 1;
    }
}

static JsonSerializerOptions CreateJsonOptions() => new()
{
    MaxDepth = 4096,
    PropertyNamingPolicy = JsonNamingPolicy.CamelCase,
    WriteIndented = true
};

static void RegisterAssemblyResolver(string? dotTraceAppDir)
{
    var probeDir = string.IsNullOrWhiteSpace(dotTraceAppDir)
        ? Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.UserProfile), "Applications", "dotTrace.app", "Contents", "DotFiles")
        : dotTraceAppDir;

    AppDomain.CurrentDomain.AssemblyResolve += (_, eventArgs) =>
    {
        var assemblyName = new System.Reflection.AssemblyName(eventArgs.Name).Name;
        if (string.IsNullOrWhiteSpace(assemblyName))
            return null;

        var candidatePath = Path.Combine(probeDir, $"{assemblyName}.dll");
        return File.Exists(candidatePath) ? System.Reflection.Assembly.LoadFrom(candidatePath) : null;
    };
}

static long GetTotalTime(DotTracePayload payload) => payload.PlusTime - payload.MinusTime;

static long GetCallCount(DotTracePayload payload) => payload.PlusCallCount - payload.MinusCallCount;

static HotspotLists BuildHotspots(IEnumerable<MutableNode> nodes)
{
    var candidates = nodes
        .Where(node => node.Kind == "method")
        .Select(node => new HotspotEntry
        {
            Id = node.Id,
            Name = node.Name,
            Kind = node.Kind,
            InclusiveTime = node.InclusiveTime,
            ExclusiveTime = node.ExclusiveTime,
            CallCount = node.CallCount
        })
        .ToArray();

    return new HotspotLists
    {
        Inclusive = candidates
            .OrderByDescending(node => node.InclusiveTime)
            .ThenBy(node => node.Name, StringComparer.Ordinal)
            .Take(50)
            .ToList(),
        Exclusive = candidates
            .OrderByDescending(node => node.ExclusiveTime)
            .ThenBy(node => node.Name, StringComparer.Ordinal)
            .Take(50)
            .ToList()
    };
}

static SerializableNode CloneTree(
    CallTreeSectionOffset offset,
    IReadOnlyDictionary<CallTreeSectionOffset, MutableNode> nodeMap,
    HashSet<CallTreeSectionOffset> ancestry)
{
    var source = nodeMap[offset];
    var nextAncestry = new HashSet<CallTreeSectionOffset>(ancestry) { offset };
    var children = new List<SerializableNode>(source.Children.Count);

    foreach (var child in source.Children)
    {
        var childOffset = ParseOffset(child.Id);
        if (nextAncestry.Contains(childOffset))
            continue;

        children.Add(CloneTree(childOffset, nodeMap, nextAncestry));
    }

    return new SerializableNode
    {
        Id = source.Id,
        Name = source.Name,
        Kind = source.Kind,
        ThreadName = null,
        InclusiveTime = source.InclusiveTime,
        ExclusiveTime = source.ExclusiveTime,
        CallCount = source.CallCount,
        Children = children
    };
}

static CallTreeSectionOffset ParseOffset(string value)
{
    var parts = value.Split('/');
    return new CallTreeSectionOffset(Convert.ToInt64(parts[0], 16), int.Parse(parts[1]));
}

sealed class SignatureResolver(
    SnapshotMasksComponent masks,
    FuidToMetadataIdConverter fuidConverter,
    IMetadataSectionAssemblyProvider assemblyProvider)
{
    private readonly Dictionary<FunctionUID, SignatureResult> _cache = [];

    public SignatureResult Resolve(FunctionUID fuid)
    {
        if (_cache.TryGetValue(fuid, out var cached))
            return cached;

        var result = ResolveCore(fuid);
        _cache[fuid] = result;
        return result;
    }

    private SignatureResult ResolveCore(FunctionUID fuid)
    {
        var fid = fuid.ToDotTraceFid(masks);
        if (fid == KnownFunctionId.RootId)
            return new SignatureResult("<thread-root>", "special");
        if (fid == KnownFunctionId.NativeCallId)
            return new SignatureResult("<native>", "special");
        if (fid == KnownFunctionId.ThreadStoppedByUnknownReasonCallId)
            return new SignatureResult("<thread-stopped>", "special");
        if (fid == KnownFunctionId.GcCallId)
            return new SignatureResult("<gc>", "special");
        if (fid == KnownFunctionId.AppDomainShutdownCallId)
            return new SignatureResult("<appdomain-shutdown>", "special");
        if (fid == KnownFunctionId.CodePitchingCallId)
            return new SignatureResult("<code-pitching>", "special");
        if (fid == KnownFunctionId.UnmanagedStackFrameCallId)
            return new SignatureResult("<unmanaged-stack-frame>", "special");
        if (fid == KnownFunctionId.UnsafeStackFrameId)
            return new SignatureResult("<unsafe-stack-frame>", "special");
        if (fid == KnownFunctionId.ContinuationsId)
            return new SignatureResult("<continuations>", "special");
        if (fid == KnownFunctionId.AwaitsId)
            return new SignatureResult("<awaits>", "special");
        if (fid == KnownFunctionId.TaskScheduledId)
            return new SignatureResult("<task-scheduled>", "special");
        if (fid == KnownFunctionId.TaskExecutionId)
            return new SignatureResult("<task-execution>", "special");
        if (fid == KnownFunctionId.TaskRecursionId)
            return new SignatureResult("<task-recursion>", "special");
        if (!KnownFunctionId.IsProcessable(fid))
            return new SignatureResult($"[special:0x{(uint)fid:X8}]", "special");

        try
        {
            var metadataId = fuidConverter.Convert(fid);
            if (metadataId.IsSynthetic)
                return new SignatureResult($"[synthetic:{metadataId}]", "special");

            var function = assemblyProvider.GetFunctionItemsLight(metadataId);
            var className = function.ClassToken != MetadataToken.Nil
                ? assemblyProvider.GetClassItems(MetadataId.Create(metadataId.MetadataIndex, function.ClassToken)).FullName
                : null;
            var qualifiedMethod = string.IsNullOrWhiteSpace(className)
                ? function.FunctionName
                : $"{className}.{function.FunctionName}";
            var methodName = string.IsNullOrEmpty(function.GenericParameters)
                ? qualifiedMethod
                : $"{qualifiedMethod}<{function.GenericParameters}>";
            return new SignatureResult(methodName, "method");
        }
        catch
        {
            return new SignatureResult($"[unresolved:0x{(uint)fid:X8}]", "special");
        }
    }
}

readonly record struct SignatureResult(string Name, string Kind);

sealed class OutputDocument
{
    public required SnapshotInfo Snapshot { get; init; }

    public required List<ThreadRootInfo> ThreadRoots { get; init; }

    public required HotspotLists Hotspots { get; init; }

    public required SerializableNode CallTree { get; init; }
}

sealed class SnapshotInfo
{
    public required string Path { get; init; }

    public required string PayloadType { get; init; }

    public required int ThreadCount { get; init; }

    public required int NodeCount { get; init; }
}

sealed class ThreadRootInfo
{
    public required string Id { get; init; }

    public required string Name { get; init; }

    public required long InclusiveTime { get; init; }
}

sealed class HotspotLists
{
    public required List<HotspotEntry> Inclusive { get; init; }

    public required List<HotspotEntry> Exclusive { get; init; }
}

sealed class HotspotEntry
{
    public required string Id { get; init; }

    public required string Name { get; init; }

    public required string Kind { get; init; }

    public required long InclusiveTime { get; init; }

    public required long ExclusiveTime { get; init; }

    public required long CallCount { get; init; }
}

class MutableNode
{
    public required string Id { get; init; }

    public required string Name { get; init; }

    public required string Kind { get; init; }

    public required long InclusiveTime { get; init; }

    public required long ExclusiveTime { get; init; }

    public required long CallCount { get; init; }

    public required List<MutableNode> Children { get; init; }
}

sealed class SerializableNode
{
    public required string Id { get; init; }

    public required string Name { get; init; }

    public required string Kind { get; init; }

    public string? ThreadName { get; init; }

    public required long InclusiveTime { get; init; }

    public required long ExclusiveTime { get; init; }

    public required long CallCount { get; init; }

    public required List<SerializableNode> Children { get; init; }
}