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

Go to download
package com.undefinedlabs.scope.rules;

import com.undefinedlabs.scope.ScopeGlobalTracer;
import com.undefinedlabs.scope.events.EventFieldsFactory;
import com.undefinedlabs.scope.events.exception.ThrowableEvent;
import com.undefinedlabs.scope.rules.sql.model.ConnectionInfo;
import com.undefinedlabs.scope.rules.sql.provider.ConnectionInfoProviderRegistry;
import com.undefinedlabs.scope.rules.transformer.ScopeAgentAdvicedTransformer;
import com.undefinedlabs.scope.utils.event.EventValues;
import com.undefinedlabs.scope.utils.sourcecode.ExceptionSourceCodeFactory;
import com.undefinedlabs.scope.utils.sourcecode.ExceptionSourceCodeFrame;
import com.undefinedlabs.scope.utils.tag.TagKeys;
import com.undefinedlabs.scope.utils.tag.TagValues;
import io.opentracing.Span;
import io.opentracing.Tracer;
import io.opentracing.tag.Tags;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.asm.Advice;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.Collections;

import static net.bytebuddy.matcher.ElementMatchers.isSubTypeOf;
import static net.bytebuddy.matcher.ElementMatchers.named;

public class ConnectionScopeAgentRule extends AbstractScopeAgentRule {

    @Override
    protected Iterable transformers() {
        final Class instrumentedClass = lookUp("java.sql.Connection");

        return (instrumentedClass != null) ?
                Collections.singleton(newAgentBuilder()
                        .type(isSubTypeOf(instrumentedClass)).transform(new ScopeAgentAdvicedTransformer(ConnectionAdvice.class, named("commit")))
                        .type(isSubTypeOf(instrumentedClass)).transform(new ScopeAgentAdvicedTransformer(ConnectionAdvice.class, named("rollback")))
                        .type(isSubTypeOf(instrumentedClass)).transform(new ScopeAgentAdvicedTransformer(ConnectionAdvice.class, named("close"))))
                : emptyCollection();
    }

    public static class ConnectionAdvice {

        @Advice.OnMethodEnter
        public static Span enter(@Advice.This Object thiz, @Advice.Origin("#m") Object methodName) {
            final Connection connection = (Connection) thiz;

            final ConnectionInfo connectionInfo = ConnectionInfoProviderRegistry.INSTANCE.getProvider(connection).extractInfo(connection);
            if (ConnectionInfo.EMPTY.equals(connectionInfo)) {
                return null;
            }

            final Tracer tracer = ScopeGlobalTracer.get();
            final Span previousSpan = tracer.activeSpan();
            if (previousSpan == null) {
                return null;
            }

            try {

                final String method = (String) methodName;
                final Span span = tracer.buildSpan(connectionInfo.getProductName() + ":" + method)
                        .withTag(TagKeys.COMPONENT, TagValues.Component.SQL)
                        .withTag(TagKeys.SPAN_KIND, Tags.SPAN_KIND_CLIENT)
                        .withTag(TagKeys.DB.DB_CONNECTION, connectionInfo.getUrl())
                        .withTag(TagKeys.DB.DB_STATEMENT, method)
                        .withTag(TagKeys.DB.DB_PARAMS, (String) null)
                        .withTag(TagKeys.DB.DB_TYPE, TagValues.DB.TYPE_SQL)
                        .withTag(TagKeys.DB.DB_USER, connectionInfo.getUserName())
                        .withTag(TagKeys.DB.DB_INSTANCE, connection.getCatalog())
                        .withTag(TagKeys.DB.DB_PRODUCT_NAME, connectionInfo.getProductName())
                        .withTag(TagKeys.DB.DB_PRODUCT_VERSION, connectionInfo.getProductVersion())
                        .withTag(TagKeys.DB.DB_DRIVER_NAME, connectionInfo.getDriverName())
                        .withTag(TagKeys.DB.DB_DRIVER_VERSION, connectionInfo.getDriverVersion())
                        .withTag(TagKeys.Network.PEER_SERVICE, connectionInfo.getPeerService())
                        .start();

                return span;

            } catch (Exception e) {
                throw new RuntimeException(e);
            }

        }

        @Advice.OnMethodExit(onThrowable = SQLException.class)
        public static void exit(@Advice.This Object thiz, @Advice.Enter Object spanObj, @Advice.Thrown Throwable throwable) {
            if (spanObj == null) {
                return;
            }

            final Span span = (Span) spanObj;

            if (throwable != null) {
                span.setTag(TagKeys.ERROR, true);

                final ExceptionSourceCodeFrame exceptionSourceCodeFrame = ExceptionSourceCodeFactory.INSTANCE.createFrame(throwable);
                final ThrowableEvent.Builder throwableEventBuilder = ThrowableEvent.newBuilder();
                throwableEventBuilder
                        .withEventType(EventValues.General.ERROR)
                        .withThrowable(exceptionSourceCodeFrame.getUserThrowable())
                        .withSource(exceptionSourceCodeFrame.getSourceCodeFrame().getLinkPathWithMethodLine());

                span.log(EventFieldsFactory.INSTANCE.createFields(throwableEventBuilder.build()));
            }

            span.finish();
        }
    }
}