ReflectiveConnectionMonitoringTransformer.java

package org.stagemonitor.jdbc;

import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;
import net.bytebuddy.utility.JavaModule;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.stagemonitor.core.util.ClassUtils;
import org.stagemonitor.tracing.wrapper.SpanWrapper;
import org.stagemonitor.tracing.wrapper.StatelessSpanEventListener;

import java.lang.reflect.Method;
import java.security.ProtectionDomain;
import java.sql.Connection;

import javax.sql.DataSource;

import __redirected.org.stagemonitor.dispatcher.Dispatcher;

import static net.bytebuddy.matcher.ElementMatchers.not;
import static org.stagemonitor.core.instrument.CachedClassLoaderMatcher.cached;
import static org.stagemonitor.core.instrument.CanLoadClassElementMatcher.canLoadClass;

/**
 * When the {@link DataSource} implementation is not loaded by the application {@link ClassLoader}, like it is common
 * in application servers like JBoss, the calls to stagemonitor can't be inserted directly but only reflectively.
 */
public class ReflectiveConnectionMonitoringTransformer extends ConnectionMonitoringTransformer {

	private static final Logger logger = LoggerFactory.getLogger(ReflectiveConnectionMonitoringTransformer.class);

	private static final String CONNECTION_MONITOR = ConnectionMonitor.class.getName();

	// [0]: ConnectionMonitor [1]: Method
	private static ThreadLocal<Object[]> connectionMonitorThreadLocal;

	public ReflectiveConnectionMonitoringTransformer() throws NoSuchMethodException {
		if (isActive()) {
			Dispatcher.getValues().putIfAbsent(CONNECTION_MONITOR, new ThreadLocal<Object[]>());
			connectionMonitorThreadLocal = Dispatcher.get(CONNECTION_MONITOR);
		}
	}

	/**
	 * If the ThreadLocal is added, the code added in {@link #addReflectiveMonitorMethodCall} gets active
	 * <p/>
	 * Using a ThreadLocal ensures that each application invokes its own instance of the ConnectionMonitor and that
	 * applications which are not monitored by stagemonitor are not influenced
	 */
	public static class ConnectionMonitorAddingSpanEventListener extends StatelessSpanEventListener {

		private final Method monitorGetConnectionMethod;

		public ConnectionMonitorAddingSpanEventListener() throws NoSuchMethodException {
			monitorGetConnectionMethod = ConnectionMonitor.class
					.getMethod("monitorGetConnection", Connection.class, Object.class, long.class);
			makeReflectionInvocationFaster(monitorGetConnectionMethod);
		}

		@Override
		public void onStart(SpanWrapper spanWrapper) {
			if (connectionMonitorThreadLocal != null) {
				connectionMonitorThreadLocal.set(new Object[]{connectionMonitor, monitorGetConnectionMethod});
			}
		}

		@Override
		public void onFinish(SpanWrapper spanWrapper, String operationName, long durationNanos) {
			if (connectionMonitorThreadLocal != null) {
				connectionMonitorThreadLocal.remove();
			}
		}

	}

	private static void makeReflectionInvocationFaster(Method method) {
		try {
			method.setAccessible(true);
		} catch (SecurityException e) {
			// ignore
		}
	}

	/**
	 * Only applies if stagemonitor can't be loaded by this class loader. For example a module class loader which loaded
	 * the DataSource but does not have access to the application classes.
	 */
	@Override
	public ElementMatcher.Junction<ClassLoader> getClassLoaderMatcher() {
		return not(cached(canLoadClass("org.stagemonitor.core.Stagemonitor")));
	}

	@Advice.OnMethodEnter
	private static long addTimestampLocalVariable() {
		return System.nanoTime();
	}

	@Advice.OnMethodExit
	private static void addReflectiveMonitorMethodCall(@Advice.This Object dataSource, @Advice.Return(readOnly = false) Connection connection, @Advice.Enter long startTime) {
		try {
			Object[] connectionMonitor = (Object[]) ((ThreadLocal) Dispatcher.getValues().get("org.stagemonitor.jdbc.ConnectionMonitor")).get();
			if (connectionMonitor != null) {
				final Method connectionMonitorMethod = (Method) connectionMonitor[1];
				final long duration = System.nanoTime() - startTime;
				// In JBoss, this method is executed in the context of the module class loader which loads the DataSource
				// The connectionMonitor is not accessible from this class loader. That's why we have to use reflection.
				connection = (Connection) connectionMonitorMethod.invoke(connectionMonitor[0], connection, dataSource, duration);
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	/**
	 * Makes sure that no DataSources are instrumented twice, even if multiple stagemonitored applications are
	 * deployed on one application server
	 */
	@Override
	public AgentBuilder.RawMatcher getRawMatcher() {
		return new AvoidDuplicateTransformationsRawMatcher();
	}

	private static class AvoidDuplicateTransformationsRawMatcher implements AgentBuilder.RawMatcher {
		@Override
		public boolean matches(TypeDescription typeDescription, ClassLoader classLoader, JavaModule javaModule, Class<?> classBeingRedefined, ProtectionDomain protectionDomain) {
			final String key = getClassAlreadyTransformedKey(typeDescription, classLoader);
			final boolean hasAlreadyBeenTransformed = Dispatcher.getValues().containsKey(key);
			if (DEBUG_INSTRUMENTATION) {
				logger.info("{}: {}", key, hasAlreadyBeenTransformed);
			}
			return !hasAlreadyBeenTransformed;
		}
	}

	private static String getClassAlreadyTransformedKey(TypeDescription typeDescription, ClassLoader classLoader) {
		return typeDescription.getName() + ClassUtils.getIdentityString(classLoader) + ".transformed";
	}

	@Override
	public void beforeTransformation(TypeDescription typeDescription, ClassLoader classLoader) {
		super.beforeTransformation(typeDescription, classLoader);
		final String key = getClassAlreadyTransformedKey(typeDescription, classLoader);
		if (DEBUG_INSTRUMENTATION) {
			logger.info("setting {}", key);
		}
		Dispatcher.getValues().put(key, Boolean.TRUE);
	}
}