DEV Community

Query Filter
Query Filter

Posted on

bridge2

import java.nio.file.Files

// Configuration for the profiler
def profilerConfig = [
    targetClassesFile: file("profiler_targets.txt"), // One class name per line
    intervalSeconds: 60,
    agentJar: file("$buildDir/libs/ProfilerAgent.jar")
]

// 1. Task to generate the Java Agent source code
task generateProfilerSource {
    def agentSrcDir = file("$buildDir/profiler-src/comet/agent")
    doFirst {
        agentSrcDir.mkdirs()
        def javaFile = new File(agentSrcDir, "ProfilerAgent.java")
        javaFile.text = """
package comet.agent;

import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.implementation.bind.annotation.*;
import net.bytebuddy.matcher.ElementMatchers;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.Method;
import java.io.*;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.LongAdder;

public class ProfilerAgent {
    public static final ConcurrentHashMap<String, Stats> metrics = new ConcurrentHashMap<>();

    public static void premain(String agentArgs, Instrumentation inst) {
        String[] parts = agentArgs.split(";");
        String configFilePath = parts[0];
        int interval = Integer.parseInt(parts[1]);

        List<String> classesToTrack = loadClasses(configFilePath);
        startReporter(interval);

        AgentBuilder agentBuilder = new AgentBuilder.Default()
            .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION);

        for (String className : classesToTrack) {
            agentBuilder = agentBuilder
                .type(ElementMatchers.named(className))
                .transform((builder, td, cl, m) -> 
                    builder.method(ElementMatchers.any())
                           .intercept(MethodDelegation.to(Interceptor.class))
                );
        }
        agentBuilder.installOn(inst);
    }

    private static List<String> loadClasses(String path) {
        try { return Files.readAllLines(new File(path).toPath()); }
        catch (Exception e) { return Collections.emptyList(); }
    }

    public static class Interceptor {
        @RuntimeType
        public static Object intercept(@Origin Method method, @SuperCall Callable<?> callable) throws Exception {
            long start = System.nanoTime();
            try { return callable.call(); }
            finally {
                long duration = System.nanoTime() - start;
                String key = method.getDeclaringClass().getSimpleName() + "." + method.getName();
                metrics.computeIfAbsent(key, k -> new Stats()).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) {
        Executors.newSingleThreadScheduledExecutor().scheduleAtFixedRate(() -> {
            File csv = new File("profiler_out_" + System.currentTimeMillis() + ".csv");
            try (PrintWriter pw = new PrintWriter(new FileWriter(csv))) {
                pw.println("Method,Calls,Avg_MS,Total_MS");
                metrics.forEach((k, v) -> {
                    long c = v.count.sumThenReset();
                    long t = v.totalTime.sumThenReset();
                    if (c > 0) {
                        pw.printf("%s,%d,%.4f,%.2f%n", k, c, (t/c)/1e6, t/1e6);
                    }
                });
            } catch (Exception e) { e.printStackTrace(); }
        }, seconds, seconds, TimeUnit.SECONDS);
    }
}
"""
    }
}

// 2. Compile and package the Agent
task buildProfilerAgent(type: Jar, dependsOn: generateProfilerSource) {
    archiveFileName = "ProfilerAgent.jar"
    from compileJava.destinationDirectory // Placeholder, usually you'd have a separate sourceSet
    // For simplicity in a single script, we use the buildDir path
    from (fileTree("$buildDir/profiler-classes"))

    manifest {
        attributes(
            'Premain-Class': 'comet.agent.ProfilerAgent',
            'Can-Retransform-Classes': 'true'
        )
    }
}

// 3. The Actual Run Task
task spawnBridgeWithProfiler(type: JavaExec, dependsOn: buildProfilerAgent) {
    group = "verification"
    main = 'com.citigroup.gffcometbridge.server.foBridgeServer'
    classpath = sourceSets.main.runtimeClasspath

    // Pass the target file and interval via the agent string
    jvmArgs "-javaagent:${profilerConfig.agentJar}=${profilerConfig.targetClassesFile.absolutePath};${profilerConfig.intervalSeconds}"

    // Standard system properties from your previous COMET logic
    jvmArgs "-Dinstance=$instanceName", "-Denv=$env"
}
Enter fullscreen mode Exit fullscreen mode

Top comments (0)