TimingService.java

package com.gmail.woodyc40.common.timing;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import javassist.*;
import net.tridentsdk.util.TridentLogger;

import java.io.IOException;
import java.lang.instrument.ClassDefinition;
import java.lang.instrument.UnmodifiableClassException;
import java.util.List;
import java.util.Map;
import java.util.UUID;

/**
 * Provides the entry point necessary to provide sampling data
 *
 * @author Pierre C
 */
public class TimingService {
    private static final Map<String, List<TimingMethod>> TIMING_METHODS = Maps.newHashMap();

    private final Class<?> def;
    private final String clazz;

    private TimingService(Class<?> clazz) {
        this.def = clazz;
        this.clazz = clazz.getName();
    }

    public static TimingService in(Class<?> c) {
        return new TimingService(c);
    }

    public static TimingService in(String clazz) {
        try {
            return new TimingService(Class.forName(clazz));
        } catch (ClassNotFoundException e) {
            e.printStackTrace();
        }

        return null;
    }

    public void sampleMethod(String method) {
        String id = new StringBuilder(clazz).append("::").append(method).toString();
        String varId = UUID.randomUUID().toString().replaceAll("-", "");

        TridentLogger.log("Instrumenting " + id);

        try {
            final ClassPool pool = ClassPool.getDefault();
            pool.appendClassPath(new LoaderClassPath(getClass().getClassLoader()));

            final CtClass compiledClass = pool.get(clazz);
            final CtMethod me = compiledClass.getDeclaredMethod(method);

            me.addLocalVariable("start_" + varId, CtClass.longType);
            me.insertBefore("start_" + varId + " = System.nanoTime();");
            me.insertAfter("{" +
                    "final long end_" + varId + " = System.nanoTime();" +
                    "com.gmail.woodyc40.common.timing.TimerService" +
                    ".putTime(\"" + id + "\", " + "(end_" + varId + " - start_" + varId + "));" +
                    "}");

            Instrument.get().redefineClasses(new ClassDefinition(def, compiledClass.toBytecode()));
        } catch (CannotCompileException | NotFoundException | IOException | ClassNotFoundException | UnmodifiableClassException e) {
            TridentLogger.error("Failed to write class to sample method " + method + "(...)");
            TridentLogger.error(e);
        }
    }

    public static void useMethod(String id, TimingMethod method) {
        List<TimingMethod> methods = TIMING_METHODS.get(id);
        if (methods == null) {
            TIMING_METHODS.put(id, Lists.newArrayList(method));
        } else {
            methods.add(method);
        }
    }

    public static void putTime(String id, long time) {
        List<TimingMethod> methods = TIMING_METHODS.get(id);
        if (methods == null) {
            TIMING_METHODS.put(id, methods = Lists.newArrayList(new DefaultTimingMethod()));
        }

        for (TimingMethod method : methods) {
            method.submit(time);
        }
    }
}