pinpoint/src/main/java/com/profiler/TomcatProfiler.java

package com.profiler;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.util.logging.Level;
import java.util.logging.Logger;

import com.profiler.config.ProfilerConfig;
import com.profiler.interceptor.bci.ByteCodeInstrumentor;
import com.profiler.interceptor.bci.JavaAssistByteCodeInstrumentor;
import com.profiler.modifier.DefaultModifierRegistry;
import com.profiler.modifier.Modifier;
import com.profiler.modifier.ModifierRegistry;

public class TomcatProfiler implements ClassFileTransformer {

	private static final Logger logger = Logger.getLogger(TomcatProfiler.class.getName());

	private String agentArgString = "";

	private Instrumentation instrumentation;
	private final ByteCodeInstrumentor byteCodeInstrumentor;

	private final ModifierRegistry modifierRepository;

	private final ProfilerConfig profilerConfig;
	private final Agent agent;

	public static void premain(String agentArgs, Instrumentation instrumentation) {
		try {
			ProfilerConfig profilerConfig = new ProfilerConfig();
			profilerConfig.readConfigFile();
			if (!profilerConfig.isProfileEnable()) {
				logger.warning("Profiler Agent not started. PROFILE_ENABLE=" + profilerConfig.isProfileEnable());
				return;
			}
			Agent agent = new Agent(profilerConfig);
			new TomcatProfiler(agentArgs, instrumentation, agent, profilerConfig);
		} catch (Exception e) {
			logger.log(Level.SEVERE, "Profiler Agent start fail. Cause:" + e.getMessage(), e);
		}
	}

	public TomcatProfiler(String agentArgs, Instrumentation instrumentation, Agent agent, ProfilerConfig profilerConfig) {
		this.agentArgString = agentArgs;
		this.profilerConfig = profilerConfig;
		this.agent = agent;

		this.instrumentation = instrumentation;
		this.instrumentation.addTransformer(this);

		String[] paths = getTomcatlibPath();
		this.byteCodeInstrumentor = new JavaAssistByteCodeInstrumentor(paths);

		this.modifierRepository = createModifierRegistry();
	}

	private String[] getTomcatlibPath() {
		String catalinaHome = System.getProperty("catalina.home");

		if (catalinaHome == null) {
			logger.info("CATALINA_HOME is null");
			return null;
		}

		if (logger.isLoggable(Level.INFO)) {
			logger.info("CATALINA_HOME=" + catalinaHome);
		}

		// TODO This is draft. How can we support both Tomcat and BLOC without this configuration?
		String type = System.getProperty("hippo.servertype", "tomcat");

		if (type.equals("bloc")) {
			return new String[] { catalinaHome + "/server/lib/catalina.jar", catalinaHome + "/common/lib/servlet-api.jar" };
		} else {
			return new String[] { catalinaHome + "/lib/servlet-api.jar", catalinaHome + "/lib/catalina.jar" };
		}
	}

	private ModifierRegistry createModifierRegistry() {
		DefaultModifierRegistry modifierRepository = new DefaultModifierRegistry(byteCodeInstrumentor, agent, profilerConfig);
		
		modifierRepository.addDummyModifier();
		
		modifierRepository.addTomcatModifier();

		// jdbc
		modifierRepository.addJdbcModifier();

		// rpc
		modifierRepository.addConnectorModifier();

		// bloc
		modifierRepository.addBLOCModifier();

		return modifierRepository;
	}

	@Override
	public byte[] transform(ClassLoader classLoader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classFileBuffer) throws IllegalClassFormatException {
		// fast java class skip
		if (className.startsWith("java")) {
			if (className.startsWith("/", 4) || className.startsWith("x/", 4)) {
				return classFileBuffer;
			}
		}
		Modifier findModifier = this.modifierRepository.findModifier(className);
		if (findModifier == null) {
			findModifier = this.modifierRepository.findModifier("*");
			// return null;
		}
		if (logger.isLoggable(Level.FINE)) {
			logger.fine("[transform] cl" + classLoader + " className:" + className + " Modifier:" + findModifier.getClass().getName());
		}
		String javassistClassName = className.replace('/', '.');

		try {
			return findModifier.modify(classLoader, javassistClassName, protectionDomain, classFileBuffer);
		} catch (Exception e) {
			logger.log(Level.SEVERE, "Modifier:" + findModifier.getTargetClass() + " modify fail. Cause:" + e.getMessage(), e);
			return null;
		}
	}
}