public class ApacheHttpClientClassHandler extends SubTypeInstrumentationHandler {
    public ApacheHttpClientClassHandler() {
        super("org.apache.http.client.HttpClient");
    }

    protected boolean transform(CtClass cc, ClassPool pool) throws NotFoundException, CannotCompileException, IOException {
        boolean transformedOneMethod = false;

        for (CtMethod method : cc.getDeclaredMethods()) {
            if ("execute".equals(method.getName()) && !Modifier.isAbstract(method.getModifiers())) {
                String hostParameter = "null";
                String requestParameter = "null";
                String requestUriParameter = "null";
                String handlerParameter = "null";
                CtClass[] parameterTypes = method.getParameterTypes();
                for (int i = 0; i < parameterTypes.length; i++) {
                    if (parameterTypes[i].subtypeOf(pool.get("org.apache.http.HttpHost"))) {
                        hostParameter = "$" + (i + 1);
                    } else if (parameterTypes[i].subtypeOf(pool.get("org.apache.http.client.methods.HttpUriRequest"))) {
                        requestParameter = "$" + (i + 1);
                        requestUriParameter = "$" + (i + 1);
                    } else if (parameterTypes[i].subtypeOf(pool.get("org.apache.http.HttpRequest"))) {
                        requestParameter = "$" + (i + 1);
                    }
                }

                // Copy the method to a uniquely named location that won't conflict with anything.
                CtMethod copiedMethod = CtNewMethod.copy(method, cc.makeUniqueName("execute"), cc, null);
                copiedMethod.setModifiers(javassist.Modifier.PRIVATE);
                cc.addMethod(copiedMethod);

                final String source = Resources.toString(Resources.getResource(getClass(), "HttpClient_execute.javassist"), Charset.forName("utf-8"));
                pool.importPackage("java.net");
                pool.importPackage("org.apache.http");
                pool.importPackage("org.apache.http.client.methods");
                pool.importPackage("com.yammer.telemetry.tracing");
                pool.importPackage("com.yammer.telemetry.agent.handlers");
                String body = source.replace("%HTTP_HOST_PARAM%", hostParameter)
                                    .replace("%HTTP_REQUEST_PARAM%", requestParameter)
                                    .replace("%HTTP_REQUEST_URI_PARAM%", requestUriParameter)
                                    .replace("%RESPONSE_HANDLER_PARAM%", handlerParameter);
                method.setBody(body, "this", copiedMethod.getName());

                transformedOneMethod = true;
            }
        }

        return transformedOneMethod;
    }
} public class EnvironmentExecutorClassHandler extends SubTypeInstrumentationHandler {
    public EnvironmentExecutorClassHandler() {
        super("com.yammer.dropwizard.config.Environment");
    }

    @Override
    protected boolean transform(CtClass cc, ClassPool pool) throws NotFoundException, CannotCompileException, IOException {
        if ("com.yammer.dropwizard.config.Environment".equals(cc.getName())) {
            switchImplementation(cc, pool, "java.util.concurrent.ThreadPoolExecutor", "com.yammer.telemetry.agent.handlers.InstrumentedThreadPoolExecutor");
            switchImplementation(cc, pool, "java.util.concurrent.ScheduledThreadPoolExecutor", "com.yammer.telemetry.agent.handlers.InstrumentedScheduledThreadPoolExecutor");
            return true;
        }
        return false;
    }
} public class HttpServletClassHandler extends SubTypeInstrumentationHandler {
    public HttpServletClassHandler() {
        super("javax.servlet.http.HttpServlet");
    }

    protected boolean transform(CtClass cc, ClassPool pool) throws NotFoundException, CannotCompileException, IOException {
        // Find the service(HttpServletRequest, HttpServletResponse) method so we can wrap it.
        CtClass request = pool.get("javax.servlet.http.HttpServletRequest");
        CtClass response = pool.get("javax.servlet.http.HttpServletResponse");
        CtMethod serviceMethod;
        try {
            serviceMethod = cc.getDeclaredMethod("service", new CtClass[]{request, response});
        } catch (NotFoundException e) {
            // Not every servlet will implement this method and that's okay.
            return false;
        }

        // Copy the method to a uniquely named location that won't conflict with anything.
        CtMethod copiedServiceMethod = CtNewMethod.copy(serviceMethod, cc.makeUniqueName("service"), cc, null);
        copiedServiceMethod.setModifiers(Modifier.PRIVATE);
        cc.addMethod(copiedServiceMethod);

        // Swap in a new method body for service() that invokes the copied version of service().
        final String source = Resources.toString(Resources.getResource(getClass(), "HttpServlet_service.javassist"), Charset.forName("utf-8"));
        pool.importPackage("java.util");
        pool.importPackage("java.math");
        pool.importPackage("com.yammer.telemetry.agent");
        pool.importPackage("com.yammer.telemetry.agent.handlers");
        pool.importPackage("com.yammer.telemetry.tracing");
        serviceMethod.setBody(source, "this", copiedServiceMethod.getName());

        return true;
    }
} public class MetricsRegistryHandler extends SubTypeInstrumentationHandler {
    public MetricsRegistryHandler() {
        super("com.yammer.metrics.core.MetricsRegistry");
    }

    @Override
    protected boolean transform(CtClass cc, ClassPool pool) throws NotFoundException, CannotCompileException, IOException {
        if (!superTypeName.equals(cc.getName())) {
            return false;
        }

        switchImplementation(cc, pool, "com.yammer.metrics.core.Timer", "com.yammer.metrics.core.InstrumentedTimer");
        switchImplementation(cc, pool, "com.yammer.metrics.core.Meter", "com.yammer.metrics.core.InstrumentedMeter");
        switchImplementation(cc, pool, "com.yammer.metrics.core.Counter", "com.yammer.metrics.core.InstrumentedCounter");

        CtMethod getOrAdd = cc.getDeclaredMethod("getOrAdd");
        getOrAdd.insertBefore(
                "{" +
                        "  if ($2 instanceof com.yammer.metrics.core.MetricNameAware) {" +
                        "    ((com.yammer.metrics.core.MetricNameAware)$2).setMetricName((com.yammer.metrics.core.MetricName)$1);" +
                        "  }" +
                        "}");

        return true;
    }

    public static void switchImplementation(CtClass cc, ClassPool pool, String from, String to) throws NotFoundException, CannotCompileException {
        CtClass oldClass = pool.get(from);
        CtClass newClass = pool.get(to);

        CodeConverter converter = new CodeConverter();
        converter.replaceNew(oldClass, newClass);

        cc.instrument(converter);
    }
} public class JdbcDriverClassHandler extends SubTypeInstrumentationHandler {
    public JdbcDriverClassHandler() {
        super("java.sql.Driver");
    }

    @Override
    protected boolean transform(CtClass cc, ClassPool pool) throws NotFoundException, CannotCompileException, IOException {
        CtMethod connect;
        try {
            connect = cc.getDeclaredMethod("connect", new CtClass[]{pool.get("java.lang.String"), pool.get("java.util.Properties")});
            if (Modifier.isAbstract(connect.getModifiers())) {
                // Don't modify abstract connect() methods.
                return false;
            }
        } catch (NotFoundException e) {
            // This is okay, it just means this particular driver doesn't have a connect(String, Properties).
            return false;
        }

        // Copy the method to a uniquely named location that won't conflict with anything.
        CtMethod copiedConnect = CtNewMethod.copy(connect, cc.makeUniqueName("connect"), cc, null);
        copiedConnect.setModifiers(Modifier.PRIVATE);
        cc.addMethod(copiedConnect);

        pool.importPackage("com.yammer.telemetry.agent.jdbc");
        connect.setBody("{ return new InstrumentedConnection($1, $proceed($$)); }", "this", copiedConnect.getName());

        return true;
    }
} public class TelemetryTransformer implements ClassFileTransformer {
    private final Set<ClassInstrumentationHandler> handlers = new HashSet<>();

    public void addHandler(ClassInstrumentationHandler handler) {
        handlers.add(handler);
    }

    @Override
    public byte[] transform(ClassLoader loader,
                            String className,
                            Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain,
                            byte[] classfileBuffer) throws IllegalClassFormatException {
        return transform(loader, className, classfileBuffer, ClassPool.getDefault());
    }

    /**
     * Allows specifying the ClassPool, this allows tests to essentially 'reload' classes.
     *
     * @param loader
     * @param className
     * @param classfileBuffer
     * @param cp
     * @return
     * @throws IllegalClassFormatException
     */
    public byte[] transform(ClassLoader loader,
                            String className,
                            byte[] classfileBuffer,
                            ClassPool cp) throws IllegalClassFormatException {
        try {
            final String realClassName = className.replace('/', '.');
            cp.insertClassPath(new LoaderClassPath(loader));
            cp.insertClassPath(new ByteArrayClassPath(realClassName, classfileBuffer));
            CtClass cc = cp.get(realClassName);

            boolean classUpdated = false;
            for (ClassInstrumentationHandler handler : handlers) {
                if (classUpdated = handler.transformed(cc, cp)) {
                    break;
                }
            }

            if (classUpdated) {
                return cc.toBytecode();
            } else {
                return null;
            }
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        } catch (Throwable t) {
            t.printStackTrace();
            throw t;
        }
    }
} public class Annotations {
    private static ServiceAnnotations serviceAnnotations = new ServiceAnnotations();

    public static void setServiceAnnotations(ServiceAnnotations serviceAnnotations) {
        Annotations.serviceAnnotations = serviceAnnotations;
    }

    public static ServiceAnnotations getServiceAnnotations() {
        return serviceAnnotations;
    }
} public class SpanHelper {
    private static final ThreadLocal<SpanContext> spanContext = new ThreadLocal<>();
    static final Logger LOG = Logger.getLogger(SpanHelper.class.getName());
    static Sampling sampler = Sampling.ON;
    private static final IDGenerator idGenerator = new IDGenerator();

    public static Sampling getSampler() {
        return sampler;
    }

    public static void setSampler(Sampling sampler) {
        SpanHelper.sampler = sampler;
    }

    /**
     * Starts a new trace.
     *
     * @param name Name to be given to the root span in the trace.
     * @return The root span of the newly created trace.
     */
    public static Span startTrace(String name) {
        return start(name, Optional.<BigInteger>absent(), Optional.<BigInteger>absent(), Optional.<BigInteger>absent(), sampler.trace() ? TraceLevel.ON : TraceLevel.OFF);
    }

    /**
     * Starts a new span within a trace. Uses the current thread context to determine the
     * trace ID and parent span ID.
     *
     * @param name Name to be given to the span.
     * @return The newly started span.
     */
    public static Span startSpan(String name) {
        return start(name, Optional.<BigInteger>absent(), Optional.<BigInteger>absent(), Optional.<BigInteger>absent(), TraceLevel.INHERIT);
    }

    /**
     * Attach to an existing span. This is useful when a span has been created elsewhere
     * and you'd like to log annotations against that span locally. For example across thread boundaries in the local
     * vm.
     *
     * @param traceId ID of the trace of the span being attached.
     * @param spanId  ID of the span being attached.
     * @param name    Name for the span - useful for debug
     * @return The attached span.
     */
    public static Span attachSpan(BigInteger traceId, BigInteger spanId, String name) {
        return start(name, Optional.of(traceId), Optional.of(spanId), Optional.<BigInteger>absent(), TraceLevel.ON);
    }

    /**
     * Starts a new span under the specified trace and parent spanId.
     *
     * This is used for attaching to an external span for example with an incoming http request.
     *
     * @param traceId - the current traceId
     * @param parentSpanId - the parent span, if available
     * @param name - the name for this span
     * @return the newly created span
     */
    public static Span startSpan(BigInteger traceId, BigInteger parentSpanId, String name) {
        return start(name, Optional.of(traceId), Optional.<BigInteger>absent(), Optional.of(parentSpanId), TraceLevel.ON);
    }

    public static Optional<Span> currentSpan() {
        SpanContext context = spanContext.get();
        if (context == null) {
            return Optional.absent();
        } else {
            return context.currentSpan();
        }
    }

    private static Span start(String name, Optional<BigInteger> traceId, Optional<BigInteger> spanId, Optional<BigInteger> parentSpanId, TraceLevel traceLevel) {
        SpanContext context = spanContext.get();
        if (context == null) {
            context = new SpanContext();
            spanContext.set(context);
        }

        if (traceLevel == TraceLevel.INHERIT) {
            traceLevel = context.currentTraceLevel();
        }

        if (!traceId.isPresent()) {
            traceId = context.currentTraceId();
            if (!traceId.isPresent()) {
                traceId = Optional.of(idGenerator.generateTraceId());
            }
        }

        if (!parentSpanId.isPresent()) {
            parentSpanId = context.currentSpanId();
        }

        if (!spanId.isPresent()) {
            spanId = Optional.of(idGenerator.generateSpanId());
        }

        final Span span = (traceLevel == TraceLevel.OFF) ?
                new DisabledSpan() :
                new EnabledSpan(traceId.get(), spanId.get(), parentSpanId, name, traceLevel);

        context.startSpan(span);
        return span;
    }

    static long nowInNanoseconds() {
        return System.currentTimeMillis() * 1000000;
    }

    static Optional<SpanContext> currentContext() {
        return Optional.fromNullable(spanContext.get());
    }

    /**
     * This is available for testing to allow checking before and after states are as expected on the span context.
     *
     * @return an immutable list of the current state of spans in the thread local context.
     */
    static ImmutableList<Span> captureSpans() {
        SpanContext context = spanContext.get();
        if (context != null) {
            return context.captureSpans();
        }

        return ImmutableList.of();
    }
} public class SpanSinkRegistry {
    private static AtomicReference<ImmutableList<SpanSink>> spanSinks = new AtomicReference<>(ImmutableList.<SpanSink>of());

    public static void register(SpanSink sink) {
        final ImmutableList<SpanSink> oldSinks = spanSinks.get();
        final ImmutableList<SpanSink> newSinks = new ImmutableList.Builder<SpanSink>().addAll(oldSinks).add(sink).build();

        if (!spanSinks.compareAndSet(oldSinks, newSinks)) {
            throw new RuntimeException("Failed to add new SpanSink, concurrent add");
        }
    }

    public static Iterable<SpanSink> getSpanSinks() {
        return spanSinks.get();
    }

    public static void clear() {
        spanSinks.set(ImmutableList.<SpanSink>of());
    }
} public class LoggingSpanSink extends AsynchronousSpanSink {
    public LoggingSpanSink(String file) throws IOException {
        super(LogJobFactory.withFile(file));
    }
}
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
import com.yammer.telemetry.agent.handlers.ApacheHttpClientClassHandler;
import com.yammer.telemetry.agent.handlers.EnvironmentExecutorClassHandler;
import com.yammer.telemetry.agent.handlers.HttpServletClassHandler;
import com.yammer.telemetry.agent.handlers.MetricsRegistryHandler;
import com.yammer.telemetry.agent.jdbc.JdbcDriverClassHandler;
import com.yammer.telemetry.instrumentation.TelemetryTransformer;
import com.yammer.telemetry.tracing.Annotations;
import com.yammer.telemetry.tracing.SpanHelper;
import com.yammer.telemetry.tracing.SpanSinkRegistry;
import com.yammer.telemetry.tracing.logging.LoggingSpanSink;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.lang.instrument.Instrumentation;
package com.yammer.telemetry.agent;



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

    public static void premain(String agentArgs, Instrumentation inst) {
        if (!agentArgs.isEmpty()) {
            try {
                TelemetryConfiguration config = loadConfiguration(agentArgs);

                SpanHelper.setSampler(config.getSampler());