com.undefinedlabs.scope.rules.URLScopeAgentRule Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of scope-rule-javanet Show documentation
Scope is a APM for tests to give engineering teams unprecedented visibility into their CI process to quickly identify, troubleshoot and fix failed builds. This artifact contains the classes to instrument the Java NET package.
There is a newer version: 0.15.1-beta.2
Show newest version
package com.undefinedlabs.scope.rules;

import com.undefinedlabs.scope.logger.ScopeLogger;
import com.undefinedlabs.scope.rules.transformer.ScopeAgentAdvicedTransformer;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.asm.Advice;
import org.slf4j.Logger;

import java.io.IOException;
import java.lang.reflect.Method;
import java.net.HttpURLConnection;
import java.net.URL;
import java.net.URLConnection;
import java.util.Collections;
import java.util.UUID;

import static net.bytebuddy.matcher.ElementMatchers.*;

public class URLScopeAgentRule extends AbstractScopeAgentRule {

    private static final Logger LOGGER = ScopeLogger.INSTANCE;

    @Override
    protected String instrumentedClassName() {
        return "java.net.URL";
    }

    @Override
    protected Iterable transformers() {
        return Collections.singleton(new AgentBuilder.Default()
            .ignore(none())
            .with(AgentBuilder.InitializationStrategy.NoOp.INSTANCE)
            .with(AgentBuilder.RedefinitionStrategy.REDEFINITION)
            .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
            .type(is(instrumentedClass()), isBootstrapClassLoader()).transform(new ScopeAgentAdvicedTransformer(HttpURLOpenConnectionAdvice.class, named("openConnection")))
        );
    }

    public void putSpanByUUID(final String uuid, final Object span) {
        URLConnectionScopeUtils.getSpanByUuid().put(uuid, span);
    }

    public void logInfo(final String message) {
        LOGGER.info(message);
    }

    public static class HttpURLOpenConnectionAdvice {

        @Advice.OnMethodExit(onThrowable = IOException.class)
        public static void exit(@Advice.This URL url, @Advice.Return(readOnly = false) URLConnection urlConnection, @Advice.Thrown Throwable throwable) {
            if(!(urlConnection.getClass().getName().contains("HttpURLConnection")) || urlConnection.getRequestProperty(URLConnectionScopeUtils.SCOPE_OPEN_CONNECTION_CUUID_KEY) != null) {
                return;
            }

            try {
                final HttpURLConnection httpURLConnection = (HttpURLConnection) urlConnection;
                final String uuid = UUID.randomUUID().toString();

                //It is necessary to load Scope classes using Reflection with SystemClassLoader because java.net classes are loaded by BootstrapClassloader.
                final Class reflectionContextClass = Class.forName("com.undefinedlabs.scope.jdk.reflection.ReflectionContext", false, ClassLoader.getSystemClassLoader());
                final Object reflCtxtInstance = reflectionContextClass.getDeclaredField("INSTANCE").get(null);
                final Method getScopeMethod = reflectionContextClass.getMethod("getScopeMethod", String.class, String.class, Class[].class);
                final Method getScopeClass = reflectionContextClass.getMethod("getScopeClass", String.class);

                final Class ruleClass = Class.forName("com.undefinedlabs.scope.rules.URLScopeAgentRule", false, ClassLoader.getSystemClassLoader());
                final Method putSpanByUuidMethod = ruleClass.getMethod("putSpanByUUID", String.class, Object.class);
                final Method logInfoMethod = ruleClass.getMethod("logInfo", String.class);
                final Object ruleInstance = ruleClass.newInstance();

                final Class globalTracerClass = (Class) getScopeClass.invoke(reflCtxtInstance,"io.opentracing.util.GlobalTracer");
                final Method getTracerMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, globalTracerClass.getName(), "get", new Class[]{});

                final Object tracer = getTracerMethod.invoke(null);
                final Class tracerClass = tracer.getClass();
                final Method activeSpanMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, tracerClass.getName(), "activeSpan", new Class[]{});
                final Object previousActiveSpan = activeSpanMethod.invoke(tracer);

                //If there is no previousActiveSpan, it is not created Span Client.
                if ( previousActiveSpan == null ){
                    return;
                }

                final Method buildSpanMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, tracerClass.getName(), "buildSpan", new Class[]{String.class});


                final Object spanBuilder = buildSpanMethod.invoke(tracer, "temporal_op_name");
                final Class spanBuilderClass = spanBuilder.getClass();

                final Method withTagStringStringMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, spanBuilderClass.getName(), "withTag", new Class[]{String.class, String.class});
                final Method withTagStringNumberMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, spanBuilderClass.getName(), "withTag", new Class[]{String.class, Number.class});
                final Method startMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, spanBuilderClass.getName(), "start", new Class[]{});

                final Class tagKeysClass = (Class) getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.utils.tag.TagKeys");
                final Class tagKeysNetworkClass = (Class) getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.utils.tag.TagKeys$Network");
                final Class tagValuesNetworkClass = (Class) getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.utils.tag.TagValues$Network");
                final Class baggageKeysClass = (Class) getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.utils.baggage.BaggageKeys");
                final Class baggageValuesClass = (Class) getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.utils.baggage.BaggageValues");

                withTagStringStringMethod.invoke(spanBuilder, tagKeysClass.getDeclaredField("SPAN_KIND").get(null), tagValuesNetworkClass.getDeclaredField("SPAN_KIND_CLIENT").get(null));
                withTagStringStringMethod.invoke(spanBuilder, tagKeysNetworkClass.getDeclaredField("HTTP_URL").get(null), httpURLConnection.getURL().toString());
                withTagStringStringMethod.invoke(spanBuilder, tagKeysNetworkClass.getDeclaredField("PEER_HOSTNAME").get(null), httpURLConnection.getURL().getHost());
                withTagStringNumberMethod.invoke(spanBuilder, tagKeysNetworkClass.getDeclaredField("PEER_PORT").get(null), Integer.valueOf(httpURLConnection.getURL().getPort()));
                final Object span = startMethod.invoke(spanBuilder);
                final Class spanClass = span.getClass();
                final Method setBaggageItemMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, spanClass.getName(), "setBaggageItem", new Class[]{String.class, String.class});
                setBaggageItemMethod.invoke(span, baggageKeysClass.getDeclaredField("TRACE_KIND").get(null), baggageValuesClass.getDeclaredField("TRACE_KIND_TEST").get(null));

                final Method activateSpanMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, tracerClass.getName(), "activateSpan", new Class[]{spanClass.getInterfaces()[0]});
                activateSpanMethod.invoke(tracer, span);

                final Class spanContextClass = (Class) getScopeClass.invoke(reflCtxtInstance, "io.opentracing.SpanContext");
                final Class formatClass = (Class) getScopeClass.invoke(reflCtxtInstance, "io.opentracing.propagation.Format");
                final Class formatBuiltinClass = (Class) getScopeClass.invoke(reflCtxtInstance, "io.opentracing.propagation.Format$Builtin");
                final Object httpHeadersConstant = formatBuiltinClass.getDeclaredField("HTTP_HEADERS").get(null);
                final Class urlConnInjectAdaptClass = (Class) getScopeClass.invoke(reflCtxtInstance, "com.undefinedlabs.scope.rules.propagation.carriers.URLConnectionInjectAdapter");
                final Method getContextMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, spanClass.getName(), "context", new Class[]{});
                final Object spanContext = getContextMethod.invoke(span);
                final Object urlConnInjectAdaptInstance = urlConnInjectAdaptClass.getConstructor(URLConnection.class).newInstance(httpURLConnection);
                final Method injectMethod = (Method) getScopeMethod.invoke(reflCtxtInstance, tracerClass.getName(), "inject", new Class[]{spanContextClass, formatClass, Object.class});
                injectMethod.invoke(tracer, spanContext, httpHeadersConstant, urlConnInjectAdaptInstance);

                putSpanByUuidMethod.invoke(ruleInstance, uuid, span);
                logInfoMethod.invoke(ruleInstance, "HttpURLOpenConnectionAdvice#exit - Active Span: " + span);

                httpURLConnection.setRequestProperty(URLConnectionScopeUtils.SCOPE_OPEN_CONNECTION_CUUID_KEY, uuid);
            } catch (Exception e){
                throw new RuntimeException(e);
            }
        }

    }
}