DEV Community

Query Filter
Query Filter

Posted on

bridge47


Enter fullscreen mode Exit fullscreen mode

package comet.agent;

import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.matcher.ElementMatchers;
import net.bytebuddy.description.type.TypeDescription;
import java.lang.instrument.Instrumentation;
import java.io.;
import java.nio.file.
;
import java.util.;
import java.util.concurrent.
;
import java.util.concurrent.atomic.LongAdder;
import java.text.SimpleDateFormat;
import java.util.jar.JarFile;

public class ProfilerAgent {
private static final String MAP_KEY = "comet.metrics.global";
private static final SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");

@SuppressWarnings("unchecked")
public static ConcurrentHashMap<String, Stats> getGlobalMap() {
    if (System.getProperties().get(MAP_KEY) == null) {
        System.getProperties().put(MAP_KEY, new ConcurrentHashMap<String, Stats>());
    }
    return (ConcurrentHashMap<String, Stats>) System.getProperties().get(MAP_KEY);
}

public static void agentmain(String agentArgs, Instrumentation inst) {
    setupBootstrap(inst);
    premain(agentArgs, inst);
}

public static void premain(String agentArgs, Instrumentation inst) {
    // Prevent double-initialization
    if (System.getProperty("profiler.active") != null) return;
    System.setProperty("profiler.active", "true");

    String folderPath = "build/profiler-results";
    new File(folderPath).mkdirs();

    // Load targets into a HashSet for O(1) lookup speed
    Set<String> targets = new HashSet<>(loadClasses("profiler_targets.txt"));
    System.out.println("PROFILER: Optimized for " + targets.size() + " targets.");

    startReporter(60, folderPath);

    new AgentBuilder.Default()
        .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
        .with(AgentBuilder.TypeStrategy.Default.REDEFINE)
        .with(new AgentBuilder.CircularityLock.Default()) 
        // CRITICAL: Ignore these to save memory and prevent OOM
        .ignore(ElementMatchers.nameStartsWith("java.")
            .or(ElementMatchers.nameStartsWith("jdk."))
            .or(ElementMatchers.nameStartsWith("sun."))
            .or(ElementMatchers.nameStartsWith("com.sun."))
            .or(ElementMatchers.nameStartsWith("net.bytebuddy."))
            .or(ElementMatchers.nameStartsWith("org.apache."))
            .or(ElementMatchers.nameStartsWith("comcode.agent.")))
        .type(typeDescription -> {
            String name = typeDescription.getName();
            // Exact match or simple wildcard check
            if (targets.contains(name)) return true;
            for (String t : targets) {
                if (t.endsWith("*") && name.startsWith(t.substring(0, t.length() - 1))) return true;
            }
            return false;
        })
        .transform((builder, typeDescription, classLoader, module) ->
            builder.method(ElementMatchers.any()
                    .and(ElementMatchers.not(ElementMatchers.isAbstract()))
                    .and(ElementMatchers.not(ElementMatchers.isNative())))
                    .intercept(Advice.to(ProfilerAdvice.class))
        )
        .installOn(inst);
}

private static void setupBootstrap(Instrumentation inst) {
    try {
        File agentJar = new File(ProfilerAgent.class.getProtectionDomain().getCodeSource().getLocation().toURI());
        if (agentJar.exists()) {
            inst.appendToBootstrapClassLoaderSearch(new JarFile(agentJar));
        }
    } catch (Exception ignored) {}
}

private static List<String> loadClasses(String path) {
    try {
        List<String> lines = Files.readAllLines(Paths.get(path));
        List<String> clean = new ArrayList<>();
        for(String s : lines) if(!s.trim().isEmpty()) clean.add(s.trim());
        return clean;
    } catch (Exception e) { return Collections.emptyList(); }
}

public static class ProfilerAdvice {
    @Advice.OnMethodEnter
    static long enter() { return System.nanoTime(); }

    @Advice.OnMethodExit(onThrowable = Throwable.class)
    static void exit(@Advice.Enter long start, @Advice.Origin("#t.#m") String methodName) {
        if (start == 0L) return;
        long duration = System.nanoTime() - start;
        ConcurrentHashMap<String, Stats> metrics = comet.agent.ProfilerAgent.getGlobalMap();
        Stats s = metrics.get(methodName);
        if (s == null) {
            metrics.putIfAbsent(methodName, new comet.agent.ProfilerAgent.Stats());
            s = metrics.get(methodName);
        }
        s.record(duration);
    }
}

public static class Stats {
    public final LongAdder count = new LongAdder();
    public final LongAdder totalTime = new LongAdder();
    public void record(long nanos) {
        count.increment();
        totalTime.add(nanos);
    }
}

private static void startReporter(int seconds, String outputDir) {
    Executors.newSingleThreadScheduledExecutor().scheduleAtFixedRate(() -> {
        ConcurrentHashMap<String, Stats> metrics = getGlobalMap();
        File csvFile = new File(outputDir, "profiler-report.csv");
        try (PrintWriter pw = new PrintWriter(new BufferedWriter(new FileWriter(csvFile, true)))) {
            String ts = sdf.format(new Date());
            metrics.forEach((method, stats) -> {
                long c = stats.count.sumThenReset();
                long t = stats.totalTime.sumThenReset();
                if (c > 0) {
                    pw.printf("%s,%s,%d,%.4f,%.2f%n", ts, method, c, (t/(double)c)/1000000.0, t/1000000.0);
                }
            });
            pw.flush();
        } catch (Exception ignored) {}
    }, seconds, seconds, TimeUnit.SECONDS);
}
Enter fullscreen mode Exit fullscreen mode

}

Top comments (0)