bridge50

package comet.agent;

import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.matcher.ElementMatchers;

import java.io.*;
import java.lang.instrument.Instrumentation;
import java.nio.file.*;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.LongAdder;

public class ProfilerAgent {

    // Metrics storage
    public static final ConcurrentHashMap<String, Stats> metrics = new ConcurrentHashMap<>();

    // Track methods already reported to avoid duplicates
    private static final Set<String> reportedMethods = ConcurrentHashMap.newKeySet();

    private static final SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");

    // ----------------- Entry points -----------------

    public static void premain(String agentArgs, Instrumentation inst) {
        System.out.println("PROFILER: Starting in premain mode.");
        initAgent(agentArgs, inst);
    }

    public static void agentmain(String agentArgs, Instrumentation inst) {
        System.out.println("PROFILER: Starting in dynamic attach mode.");
        initAgent(agentArgs, inst);
    }

    // ----------------- Initialization -----------------

    private static void initAgent(String agentArgs, Instrumentation inst) {
        String configPath = "profiler_targets.txt";
        int interval = 60; // seconds
        String outputDir = ".";

        if (agentArgs != null && !agentArgs.isEmpty()) {
            String[] parts = agentArgs.split(";");
            if (parts.length > 0) configPath = parts[0];
            if (parts.length > 1) interval = Integer.parseInt(parts[1]);
            if (parts.length > 2) outputDir = parts[2];
        }

        List<String> targetClasses = loadClasses(configPath);
        startReporter(interval, outputDir);

        // ----------------- ByteBuddy Transformer -----------------
        AgentBuilder agentBuilder = new AgentBuilder.Default()
                .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
                .type(type -> targetClasses.stream().anyMatch(name -> type.getName().equals(name.trim())))
                .transform((builder, typeDescription, classLoader, module) ->
                        builder.method(ElementMatchers.any()
                                        .and(ElementMatchers.not(ElementMatchers.isAbstract()))
                                        .and(ElementMatchers.not(ElementMatchers.isNative())))
                                .intercept(Advice.to(ProfilerAdvice.class))
                );

        agentBuilder.installOn(inst);

        // ----------------- Robust retransformation -----------------
        if (inst.isRetransformClassesSupported()) {
            System.out.println("PROFILER: Retransform supported. Instrumenting already-loaded classes...");
            for (Class<?> loadedClass : inst.getAllLoadedClasses()) {
                if (targetClasses.contains(loadedClass.getName()) && inst.isModifiableClass(loadedClass)) {
                    try {
                        inst.retransformClasses(loadedClass);
                        System.out.println("PROFILER: Retransformed " + loadedClass.getName());
                    } catch (Throwable t) {
                        System.err.println("PROFILER: Failed to retransform " + loadedClass.getName() + ": " + t);
                    }
                }
            }
        } else {
            System.err.println("PROFILER WARNING: JVM does NOT support retransformation. " +
                    "Only classes loaded after agent attachment will be instrumented.");
        }
    }

    // ----------------- Load target classes -----------------

    private static List<String> loadClasses(String path) {
        try {
            return Files.readAllLines(Paths.get(path));
        } catch (Exception e) {
            System.err.println("PROFILER: Could not read config file " + path + ": " + e);
            return Collections.emptyList();
        }
    }

    // ----------------- Advice -----------------

    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) {
            // Record method for CSV reporting
            reportedMethods.add(methodName);

            // Record duration
            Stats s = metrics.computeIfAbsent(methodName, k -> new Stats());
            long duration = System.nanoTime() - start;
            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);
        }
    }

    // ----------------- Reporter -----------------

    private static void startReporter(int intervalSeconds, final String outputDir) {
        File csvFile = new File(outputDir, "profiler_combined_report.csv");
        ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
        executor.scheduleAtFixedRate(() -> {
            String timestamp = sdf.format(new Date());
            boolean writeHeader = !csvFile.exists() || csvFile.length() == 0;

            try (PrintWriter pw = new PrintWriter(new BufferedWriter(new FileWriter(csvFile, true)))) {
                if (writeHeader) {
                    pw.println("Timestamp,Method,Calls,Avg_MS,Total_MS");
                }

                int methodsReported = 0;
                for (String method : reportedMethods) {
                    Stats stats = metrics.get(method);
                    if (stats != null) {
                        long calls = stats.count.sumThenReset();
                        long totalNanos = stats.totalTime.sumThenReset();
                        if (calls > 0) {
                            double avgMs = totalNanos / (double) calls / 1_000_000.0;
                            double totalMs = totalNanos / 1_000_000.0;
                            pw.printf("%s,%s,%d,%.4f,%.2f%n", timestamp, method, calls, avgMs, totalMs);
                            methodsReported++;
                        }
                    }
                }

                System.out.println("PROFILER: Report flushed at " + timestamp +
                        " | Methods reported: " + methodsReported);
                System.out.flush();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }, intervalSeconds, intervalSeconds, TimeUnit.SECONDS);
    }
}