Start of rework of instrumentation to support shading.

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:<p>
+     *    &lt;original package>:&lt;shaded package>
+     * <p>
+     * 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 <code>org.mariadb.jdbc.MySQLDataSource</code>,
+     * that has been shaded to <code>com.other.maria.jdbc.MySQLDataSource</code>.  In this case,
+     * the following mapping could be used:<p>
+     *    org.mariadb:com.other.maria
+     * <br>
+     * @param mapping a mapping of the form: &lt;original package>:&lt;shaded package>
+     */
+    public void setShadedCodexMapping(String mapping)
+    {
+        this.shadedCodexMapping = mapping;
+    }
+
     public void validate()
     {
         Logger logger = LoggerFactory.getLogger(getClass());

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();

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

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<String, Set<String>> 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<String> hash = (HashSet<String>) codex.get(dsClassName);
+            String keyPrefix = getKeyPrefix(dsClassName);
-            if (hash == null)
+            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<String> classes = (HashSet<String>) codex.get(keyPrefix + ".baseConnection");
             loadClasses(classes, HikariClassTransformer.CONNECTION);
 
             classes = (HashSet<String>) codex.get(keyPrefix + ".subConnection");
-            loadClasses(            classes, HikariClassTransformer.CONNECTION_SUBCLASS);
+            loadClasses(classes, HikariClassTransformer.CONNECTION_SUBCLASS);
 
             classes = (HashSet<String>) codex.get(keyPrefix + ".baseStatement");
             loadClasses(classes, HikariClassTransformer.STATEMENT);
@@ -109,13 +111,37 @@ public class HikariClassScanner
         }
     }
 
+    private String getKeyPrefix(String dsClassName)
+    {
+        if (shadedCodexMapping == null)
+        {
+            HashSet<String> hash = (HashSet<String>) 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<String> hash = (HashSet<String>) codex.get(key);
+                return hash.iterator().next();
+            }
+        }
+
+        return null;
+    }
+
     /**
      * @throws IOException 
      */
-    @SuppressWarnings("unchecked")
     private boolean loadCodex() throws IOException
     {
-        codex = new Properties();
+        codex = new HashMap<String, Set<String>>();
 
         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<String> shadedClasses = new HashSet<>();
+            for (String clazz : classes)
+            {
+                shadedClasses.add(clazz.replace(origPackage, shadePackage));
+            }
+
+            classes = shadedClasses;
+        }
+
         transformer.setScanClass(classes, classType);
         for (String clazz : classes)
         {

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<String> 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;
         }
     }