From b5101d5e18d3fa68de4e8982cb94a681c0430d8b Mon Sep 17 00:00:00 2001 From: Brett Wooldridge Date: Wed, 13 Nov 2013 11:11:52 +0900 Subject: [PATCH] Start of rework of instrumentation to support shading. --- .../java/com/zaxxer/hikari/HikariConfig.java | 30 ++++++++- .../java/com/zaxxer/hikari/HikariPool.java | 8 ++- .../javassist/AgentRegistrationElf.java | 6 +- .../hikari/javassist/HikariClassScanner.java | 63 +++++++++++++++---- .../javassist/HikariClassTransformer.java | 9 ++- 5 files changed, 94 insertions(+), 22 deletions(-) diff --git a/src/main/java/com/zaxxer/hikari/HikariConfig.java b/src/main/java/com/zaxxer/hikari/HikariConfig.java index 3dadf264..c52d18f6 100644 --- a/src/main/java/com/zaxxer/hikari/HikariConfig.java +++ b/src/main/java/com/zaxxer/hikari/HikariConfig.java @@ -46,12 +46,11 @@ public final class HikariConfig implements HikariConfigMBean private String poolName; private String connectionTestQuery; private String dataSourceClassName; + private String shadedCodexMapping; private boolean isJdbc4connectionTest; private boolean isAutoCommit; - - private Properties driverProperties; - private boolean isUseInstrumentation; + private Properties driverProperties; /** * Default constructor @@ -314,6 +313,31 @@ public final class HikariConfig implements HikariConfigMBean this.poolName = poolName; } + public String getShadedCodexMapping() + { + return shadedCodexMapping; + } + + /** + * Set a package mapping for "shaded" drivers so that we can find the DataSource + * in the instrumentation codex even though the package/class name is different. + * The mapping should be of the form:

+ * <original package>:<shaded package> + *

+ * Where the original package name is a widely scoped as possible while still being + * unique to the driver. Typically, this is the first two segments of a package + * name. For example, take the DataSource org.mariadb.jdbc.MySQLDataSource, + * that has been shaded to com.other.maria.jdbc.MySQLDataSource. In this case, + * the following mapping could be used:

+ * org.mariadb:com.other.maria + *
+ * @param mapping a mapping of the form: <original package>:<shaded package> + */ + public void setShadedCodexMapping(String mapping) + { + this.shadedCodexMapping = mapping; + } + public void validate() { Logger logger = LoggerFactory.getLogger(getClass()); diff --git a/src/main/java/com/zaxxer/hikari/HikariPool.java b/src/main/java/com/zaxxer/hikari/HikariPool.java index 555e4beb..cab3ac37 100644 --- a/src/main/java/com/zaxxer/hikari/HikariPool.java +++ b/src/main/java/com/zaxxer/hikari/HikariPool.java @@ -78,21 +78,23 @@ public final class HikariPool implements HikariPoolMBean this.jdbc4ConnectionTest = configuration.isJdbc4ConnectionTest(); this.leakDetectionThreshold = configuration.getLeakDetectionThreshold(); + String dsClassName = configuration.getDataSourceClassName(); try { - delegationProxies = !configuration.isUseInstrumentation() || !AgentRegistrationElf.loadTransformerAgent(configuration.getDataSourceClassName()); + String shadedCodexMapping = configuration.getShadedCodexMapping(); + delegationProxies = !configuration.isUseInstrumentation() || !AgentRegistrationElf.loadTransformerAgent(dsClassName, shadedCodexMapping); if (delegationProxies) { LOGGER.info("Using Javassist delegate-based proxies."); } - Class clazz = ClassLoaderUtils.loadClass(configuration.getDataSourceClassName()); + Class clazz = ClassLoaderUtils.loadClass(dsClassName); this.dataSource = (DataSource) clazz.newInstance(); PropertyBeanSetter.setTargetFromProperties(dataSource, configuration.getDataSourceProperties()); } catch (Exception e) { - throw new RuntimeException("Could not create datasource class: " + configuration.getDataSourceClassName(), e); + throw new RuntimeException("Could not create datasource class: " + dsClassName, e); } registerMBean(); diff --git a/src/main/java/com/zaxxer/hikari/javassist/AgentRegistrationElf.java b/src/main/java/com/zaxxer/hikari/javassist/AgentRegistrationElf.java index 5fb2688b..e61977cb 100644 --- a/src/main/java/com/zaxxer/hikari/javassist/AgentRegistrationElf.java +++ b/src/main/java/com/zaxxer/hikari/javassist/AgentRegistrationElf.java @@ -37,7 +37,7 @@ public class AgentRegistrationElf { private static final Logger LOGGER = LoggerFactory.getLogger(HikariClassScanner.class); - public static boolean loadTransformerAgent(String dsClassName) + public static boolean loadTransformerAgent(String dsClassName, String shadedCodexMapping) { String agentJarPath = getSelfJarPath(); if (agentJarPath == null) @@ -48,7 +48,6 @@ public class AgentRegistrationElf try { - Properties systemProperties = System.getProperties(); HikariClassTransformer transformer = new HikariClassTransformer(); @@ -57,7 +56,7 @@ public class AgentRegistrationElf registerInstrumentation(agentJarPath); LOGGER.info("Successfully loaded instrumentation agent. Scanning classes..."); - HikariClassScanner scanner = new HikariClassScanner(transformer); + HikariClassScanner scanner = new HikariClassScanner(transformer, shadedCodexMapping); return scanner.scanClasses(dsClassName); } catch (Exception e) @@ -107,6 +106,7 @@ public class AgentRegistrationElf * dynamically. * * @param jarPath the path to our own jar file + * @param shadedCodexMapping * @param dsClassName * @throws AttachNotSupportedException thrown if the JVM does not support attachment * @throws IOException thrown if the instrumentation JAR cannot be read diff --git a/src/main/java/com/zaxxer/hikari/javassist/HikariClassScanner.java b/src/main/java/com/zaxxer/hikari/javassist/HikariClassScanner.java index f1f8c216..24296c1c 100644 --- a/src/main/java/com/zaxxer/hikari/javassist/HikariClassScanner.java +++ b/src/main/java/com/zaxxer/hikari/javassist/HikariClassScanner.java @@ -20,8 +20,10 @@ import java.io.BufferedReader; import java.io.IOException; import java.io.InputStream; import java.io.InputStreamReader; +import java.util.HashMap; import java.util.HashSet; -import java.util.Properties; +import java.util.Map; +import java.util.Set; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -38,14 +40,16 @@ public class HikariClassScanner private HikariClassTransformer transformer; - private Properties codex; + private Map> codex; - public HikariClassScanner(HikariClassTransformer transformer) + private String shadedCodexMapping; + + public HikariClassScanner(HikariClassTransformer transformer, String shadedCodexMapping) { this.transformer = transformer; + this.shadedCodexMapping = shadedCodexMapping; } - @SuppressWarnings("unchecked") public boolean scanClasses(String dsClassName) { try @@ -59,21 +63,19 @@ public class HikariClassScanner return false; } - HashSet hash = (HashSet) codex.get(dsClassName); - if (hash == null) + String keyPrefix = getKeyPrefix(dsClassName); + if (keyPrefix == null) { LOGGER.warn("DataSource {} not found in the instrumentation codex. Please report at http://github.com/brettwooldridge/HikariCP.", dsClassName); LOGGER.info("Using delegation instead of instrumentation"); return false; } - String keyPrefix = hash.iterator().next(); - HashSet classes = (HashSet) codex.get(keyPrefix + ".baseConnection"); loadClasses(classes, HikariClassTransformer.CONNECTION); classes = (HashSet) codex.get(keyPrefix + ".subConnection"); - loadClasses( classes, HikariClassTransformer.CONNECTION_SUBCLASS); + loadClasses(classes, HikariClassTransformer.CONNECTION_SUBCLASS); classes = (HashSet) codex.get(keyPrefix + ".baseStatement"); loadClasses(classes, HikariClassTransformer.STATEMENT); @@ -109,13 +111,37 @@ public class HikariClassScanner } } + private String getKeyPrefix(String dsClassName) + { + if (shadedCodexMapping == null) + { + HashSet hash = (HashSet) codex.get(dsClassName); + + return (hash == null ? null : hash.iterator().next()); + } + + String[] split = shadedCodexMapping.split(":"); + String origPackage = split[0]; + String shadePackage = split[1]; + + for (String key : codex.keySet()) + { + if (key.replace(origPackage, shadePackage).equals(dsClassName)) + { + HashSet hash = (HashSet) codex.get(key); + return hash.iterator().next(); + } + } + + return null; + } + /** * @throws IOException */ - @SuppressWarnings("unchecked") private boolean loadCodex() throws IOException { - codex = new Properties(); + codex = new HashMap>(); InputStream inputStream = this.getClass().getResourceAsStream("/META-INF/codex.properties"); if (inputStream == null) @@ -167,6 +193,21 @@ public class HikariClassScanner return; } + if (shadedCodexMapping != null) + { + String[] split = shadedCodexMapping.split(":"); + String origPackage = split[0]; + String shadePackage = split[1]; + + HashSet shadedClasses = new HashSet<>(); + for (String clazz : classes) + { + shadedClasses.add(clazz.replace(origPackage, shadePackage)); + } + + classes = shadedClasses; + } + transformer.setScanClass(classes, classType); for (String clazz : classes) { diff --git a/src/main/java/com/zaxxer/hikari/javassist/HikariClassTransformer.java b/src/main/java/com/zaxxer/hikari/javassist/HikariClassTransformer.java index 186c22a0..80da928b 100644 --- a/src/main/java/com/zaxxer/hikari/javassist/HikariClassTransformer.java +++ b/src/main/java/com/zaxxer/hikari/javassist/HikariClassTransformer.java @@ -48,7 +48,7 @@ import org.slf4j.LoggerFactory; */ public class HikariClassTransformer implements ClassFileTransformer { - private static final Logger LOGGER = LoggerFactory.getLogger(HikariClassTransformer.class); + private final Logger LOGGER; public static final int UNDEFINED = 0; public static final int CONNECTION = 1; @@ -71,11 +71,13 @@ public class HikariClassTransformer implements ClassFileTransformer /** * Private constructor. + * @param shadedCodexMapping * * @param sniffPackage the package name used to filter only classes we are interested in */ public HikariClassTransformer() { + LOGGER = LoggerFactory.getLogger(HikariClassTransformer.class); } public void setScanClass(HashSet scanClasses, int classType) @@ -104,8 +106,11 @@ public class HikariClassTransformer implements ClassFileTransformer classPool.appendClassPath(new LoaderClassPath(loader)); } + ClassLoader contextClassLoader = Thread.currentThread().getContextClassLoader(); try { + Thread.currentThread().setContextClassLoader(loader); + ClassFile classFile = new ClassFile(new DataInputStream(new ByteArrayInputStream(classfileBuffer))); LOGGER.info("Instrumenting class {}", className); @@ -145,8 +150,8 @@ public class HikariClassTransformer implements ClassFileTransformer } finally { + Thread.currentThread().setContextClassLoader(contextClassLoader); LOGGER.debug("--------------------------------------------------------------------------"); - //classType = UNDEFINED; } }