Converts .collapsed output produced by async-profiler to VisualVM format

NpsConverter.java

/*
 * Copyright 2020 Andrei Pangin
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import java.io.BufferedReader;
import java.io.ByteArrayOutputStream;
import java.io.DataOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.IOException;
import java.io.OutputStream;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.TreeMap;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.zip.DeflaterOutputStream;

/**
 * Converts .collapsed output produced by async-profiler to nps format,
 * supported by NetBeans and VisualVM.
 */
public class NpsConverter {

    static class Frame extends TreeMap<String, Frame> {
        long total;
        long self;
    }

    private final Frame root = new Frame();
    private final Map<String, Integer> methods = new LinkedHashMap<>();

    public NpsConverter(String fileName) throws IOException {
        try (BufferedReader br = new BufferedReader(new FileReader(fileName))) {
            for (String line; (line = br.readLine()) != null; ) {
                int space = line.lastIndexOf(' ');
                String[] trace = line.substring(0, space).split(";");
                long ticks = Long.parseLong(line.substring(space + 1));

                Frame frame = root;
                for (String name : trace) {
                    if (!name.isEmpty()) {
                        frame.total += ticks;
                        frame = frame.computeIfAbsent(name, key -> new Frame());
                    }
                }
                frame.total += ticks;
                frame.self += ticks;
            }
        }
    }

    public void dump(OutputStream outFile) throws IOException {
        ByteArrayOutputStream profile = new ByteArrayOutputStream(32000);
        DataOutputStream out = new DataOutputStream(profile);

        out.writeInt(1);
        out.writeLong(System.currentTimeMillis());
        out.writeLong(0);
        out.writeByte(0);

        collectAllMethods();
        writeMethods(out);
        writeThreads(out);

        ByteArrayOutputStream wrapper = new ByteArrayOutputStream(out.size());
        wrapper.write(new byte[]{'n', 'B', 'p', 'R', 'o', 'F', 'i', 'L', 'e', 'R', 1, 2, 0, 0, 0, 1});
        wrapper.write(new byte[8]);
        try (DeflaterOutputStream compressedWrapper = new DeflaterOutputStream(wrapper)) {
            profile.writeTo(compressedWrapper);
        }

        byte[] result = wrapper.toByteArray();
        put(result, 16, result.length - 24, 4);
        put(result, 20, profile.size(), 4);

        outFile.write(result);
        outFile.write(new byte[6]);
    }

    private void collectAllMethods() {
        methods.put("Thread", 0);
        for (Frame frame : root.values()) {
            collectMethods(frame);
        }
    }

    private void collectMethods(Frame frame) {
        frame.forEach((name, child) -> {
            methods.putIfAbsent(name, methods.size());
            if (!child.isEmpty()) {
                collectMethods(child);
            }
        });
    }

    private void writeMethods(DataOutputStream out) throws IOException {
        out.writeInt(methods.size());

        for (String method : methods.keySet()) {
            String methodName = stripMethodSuffix(method);
            int sig = methodName.indexOf('(');
            int dot = sig >= 0 ? methodName.lastIndexOf('.', sig) : methodName.lastIndexOf('.');

            if (dot >= 0) {
                out.writeUTF(methodName.substring(0, dot));
                out.writeUTF(sig >= 0 ? methodName.substring(dot + 1, sig) : methodName.substring(dot + 1));
                out.writeUTF(sig >= 0 ? methodName.substring(sig) : "");
            } else {
                out.writeUTF(methodName);
                out.writeUTF("");
                out.writeUTF("");
            }
        }
    }

    private void writeThreads(DataOutputStream out) throws IOException {
        out.writeInt(root.size());

        int threadId = 0;
        for (Map.Entry<String, Frame> entry : root.entrySet()) {
            out.writeInt(++threadId);
            out.writeUTF(stripThreadName(entry.getKey()));
            out.writeByte(0);

            byte[] profile = encodeProfile(entry.getValue());
            out.writeInt(profile.length);
            out.write(profile);
            out.writeInt(18);
            out.writeLong(0);
            out.writeLong(0);
            out.writeDouble(0);
            out.writeDouble(0);
            out.writeLong(0);
            out.writeLong(0);
            out.writeLong(0);
            out.writeLong(0);
            out.writeLong(0);
            out.writeByte(0);
        }
    }

    private byte[] encodeProfile(Frame threadRoot) {
        int children = countChildren(threadRoot);
        int size = (children + 1) * 18 + children * 3;

        byte[] profile = new byte[size + (size <= 0xffffff ? 0 : children)];
        encodeFrame(profile, 0, "Thread", threadRoot);
        return profile;
    }

    private int encodeFrame(byte[] profile, int pos, String name, Frame frame) {
        int methodId = methods.get(name);
        put(profile, pos, methodId, 2);
        put(profile, pos + 2, frame.total, 4);
        put(profile, pos + 6, frame.total * 1000, 5);
        put(profile, pos + 11, frame.self * 1000, 5);
        put(profile, pos + 16, frame.size(), 2);

        int slotSize = profile.length <= 0xffffff ? 3 : 4;
        int childOffset = pos + 18;
        pos = childOffset + frame.size() * slotSize;

        for (Map.Entry<String, Frame> entry : frame.entrySet()) {
            put(profile, childOffset, pos, slotSize);
            childOffset += slotSize;
            pos = encodeFrame(profile, pos, entry.getKey(), entry.getValue());
        }

        return pos;
    }

    private static int countChildren(Frame frame) {
        int result = frame.size();
        for (Frame child : frame.values()) {
            if (!child.isEmpty()) {
                result += countChildren(child);
            }
        }
        return result;
    }

    private static void put(byte[] array, int pos, long n, int bytes) {
        switch (bytes) {
            case 5:
                array[pos++] = (byte) (n >>> 32);
            case 4:
                array[pos++] = (byte) (n >>> 24);
            case 3:
                array[pos++] = (byte) (n >>> 16);
            case 2:
                array[pos++] = (byte) (n >>> 8);
            case 1:
                array[pos] = (byte) n;
        }
    }

    private static final Pattern THREAD_NAME_PATTERN = Pattern.compile("\\[(.*) tid=[0-9]+]");

    private static String stripThreadName(String name) {
        Matcher m = THREAD_NAME_PATTERN.matcher(name);
        return m.matches() ? m.group(1) : name;
    }

    private static String stripMethodSuffix(String name) {
        int len = name.length();
        if (len >= 4 && name.charAt(len - 1) == ']' && name.regionMatches(len - 4, "_[", 0, 2)) {
            return name.substring(0, len - 4);
        }
        return name;
    }

    private static String extractFileName(String fileName) {
        int nameStart = fileName.lastIndexOf(File.separatorChar) + 1;
        int dot = fileName.lastIndexOf('.');
        if (dot >= nameStart) {
            return fileName.substring(nameStart, dot);
        } else {
            return fileName.substring(nameStart);
        }
    }

    public static void main(String[] args) throws Exception {
        if (args.length < 2) {
            System.out.println("Usage: java " + NpsConverter.class.getName() + " input.collapsed output.nps");
            System.exit(1);
        }

        File dst = new File(args[1]);
        if (dst.isDirectory()) {
            dst = new File(dst, extractFileName(args[0]) + ".nps");
        }

        NpsConverter converter = new NpsConverter(args[0]);
        try (FileOutputStream out = new FileOutputStream(dst)) {
            converter.dump(out);
        }
    }
}