How to intercept inherited method

Question

I want to build a Java agent with Byte Buddy, which intercepts JDBC connection creation and closing event. I want to support data sources such as HikariCP, DBCP, etc, also plain JDBC with RDBMS drivers.

My code:

Instrumention setup

new AgentBuilder.Default()
                .type(startMatcher.and(isSubTypeOf(java.sql.Connection.class).and(not(isAbstract())))
                .transform(constructorTransformer).transform(methodsTransformer).with(listener).installOn(inst);

Intercepting code

package org.wxt.xtools.agents.jdbcmon;

import java.io.IOException;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.lang.reflect.Constructor;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;

import org.wxt.xtools.agents.utils.StringUtils;
import org.wxt.xtools.agents.utils.TopN;

import ch.qos.logback.classic.Logger;
import io.prometheus.client.Histogram;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.asm.Advice.This;
import net.bytebuddy.implementation.bind.annotation.AllArguments;
import net.bytebuddy.implementation.bind.annotation.Origin;
import net.bytebuddy.implementation.bind.annotation.RuntimeType;
import net.bytebuddy.implementation.bind.annotation.SuperCall;

/**
 * 
 * @author ggfan
 * 
 */
public class JDBCAPIInterceptor {

    // ***ATTENTION*** fields below must be public, as we will access them from the
    // intercepting methods
    public static Map<Integer, ConnectionInfo> conns = new HashMap<Integer, ConnectionInfo>();
    
    public static Map<String, String> stacks = new HashMap<String, String>();

    public static Map<String, Integer> counstsByStack = new HashMap<String, Integer>();
    
    public static TopN<StatementExecutionInfo> slowStatements = new TopN<StatementExecutionInfo>(10, new Comparator<StatementExecutionInfo>() {
        @Override
        public int compare(StatementExecutionInfo o1, StatementExecutionInfo o2) {
            long d1 = o1.getExecutionEndTime() - o1.getExecutionStartTime();
            long d2 = o2.getExecutionEndTime() - o2.getExecutionStartTime();
            return Long.compare(d1, d2);
        }
    });

    public static Logger log = JDBCMonitor.log;

    public static JDBCStatsReport getInfo() {
        JDBCStatsReport report = new JDBCStatsReport();
        List<ConnectionInfo> data = new ArrayList<ConnectionInfo>();
        for (ConnectionInfo ci : conns.values()) {
            ConnectionInfo copy = new ConnectionInfo();
            copy.setHash(ci.getHash());
            copy.setStackHash(ci.getStackHash() + "(" + counstsByStack.get(ci.getStackHash()) + ")");
            copy.setCreationTime(ci.getCreationTime());
            copy.setLastActiveTime(ci.getLastActiveTime());
            copy.setLastMethod(ci.getLastMethod());
            copy.setCurrentStatement(ci.getCurrentStatement());
            data.add(copy);                                                                                                                                                                                                                                                                
        }
        report.setConnectionInfo(data);
        report.setSlowStatementTopN(slowStatements);
        return report;
    }

    @RuntimeType
    public static Object interceptor(@Origin Class<?> clazz, @Origin Method method, @SuperCall Callable<?> callable,
            @net.bytebuddy.implementation.bind.annotation.This Object inst, @AllArguments Object[] args)
            throws Exception {
        // this is equal to original method, will not cause inner calls to other
        // matching methods get intercepted.
        // Object o = callable.call();
        log.debug("intercepting {}#{}", clazz.getName(), method.getName());
        int hashCode = System.identityHashCode(inst);
        if (java.sql.Connection.class.isAssignableFrom(clazz)) {
            ConnectionInfo ci = conns.get(hashCode);
            // TODO fix bug here!!!
            if (ci == null) {
                log.debug("connection@{} is not in records anymore, maybe called #{} after close, ignored", hashCode,
                        method.getName());
            }
            return interceptConnectionMethods(method, hashCode, ci, callable, args);
        } else if (java.sql.Statement.class.isAssignableFrom(clazz)) {
            return interceptStatementMethods(method, hashCode, callable, args);
        } else if (java.sql.PreparedStatement.class.isAssignableFrom(clazz)) {
            return interceptStatementMethods(method, hashCode, callable, args);
        } else if (java.sql.CallableStatement.class.isAssignableFrom(clazz)) {
            return interceptStatementMethods(method, hashCode, callable, args);
        }
        return null;
    }

    private static Object interceptConnectionMethods(Method method, int hashCode, ConnectionInfo ci,
            Callable<?> callable, Object[] args) throws Exception {
        Object o = callable.call();
        log.debug("connection@{} used by {}", hashCode, method.getName());
        ci.setLastActiveTime(System.currentTimeMillis());
        ci.setLastMethod(method.getName());
        int resultHash = System.identityHashCode(o);
        if (method.getName().equals("close")) {
            log.info("connection@{} released", hashCode);
            String stackHash = ci.getStackHash();
            Integer scount = counstsByStack.get(stackHash);
            if (scount != null && scount > 0) {
                int newscount = scount - 1;
                log.info("set connection count to {} by stack hash {}", newscount, stackHash);
                if (newscount == 0) {
                    counstsByStack.remove(stackHash);
                    stacks.remove(stackHash);
                } else {
                    counstsByStack.put(stackHash, newscount);
                }
            } else {
                log.error("connection count by stack hash {} is not supposed to be null or less than zero", stackHash);
            }
            conns.remove(hashCode);
        } else if (method.getName().equals("createStatement")) {
            StatementExecutionInfo stmt = new StatementExecutionInfo();
            stmt.setType(StatementType.NORMAL);
            stmt.setHash(resultHash);
            conns.get(hashCode).setCurrentStatement(stmt);
            log.info("statement@{} created, type {}, attached to connection@{}", resultHash, StatementType.NORMAL, hashCode);
        } else if (method.getName().equals("prepareStatement")) {
            StatementExecutionInfo stmt = new StatementExecutionInfo();
            stmt.setType(StatementType.PREPARED);
            stmt.setSqlText((String) args[0]);
            stmt.setHash(resultHash);
            conns.get(hashCode).setCurrentStatement(stmt);
            log.info("statement@{} created, type {}, attached to connection@{}", resultHash, StatementType.PREPARED, hashCode);
        } else if (method.getName().equals("prepareCall")) {
            StatementExecutionInfo stmt = new StatementExecutionInfo();
            stmt.setType(StatementType.CALLABLE);
            stmt.setSqlText((String) args[0]);
            stmt.setHash(resultHash);
            conns.get(hashCode).setCurrentStatement(stmt);
            log.info("statement@{} created, type {}, attached to connection@{}", resultHash, StatementType.CALLABLE, hashCode);
        }
        return o;
    }

    private static Object interceptStatementMethods(Method method, int hashCode, Callable<?> callable, Object[] args)
            throws Exception {
        log.info("intercepting statement method {}", method.getName());
        Object o = null;
        ConnectionInfo containingCI = conns.values().stream().filter(c -> c.getCurrentStatement().getHash() == hashCode)
                .findFirst().orElse(null);
        if (containingCI == null) {
            log.warn("unexpected situation happened: statement can't found containing connection!");
        }
        if (method.getName().equals("close")) {
            o = callable.call();
            // TODO statement close method not intercepted ??
            if (containingCI != null) {
                containingCI.setCurrentStatement(null);
                log.info("statement@{} closed, detached from connection@{}", hashCode, containingCI.getHash());
            }
        }
        // all statement execution method trigger
        else if (method.getName().startsWith("execute")) {
            String stack = getCurrentThreadStackTrace();
            String stackHash = StringUtils.md5(stack);
            stacks.put(stackHash, stack);
            log.info("statement@{} {} started by {}", hashCode, method.getName(), stackHash);
            long est = System.currentTimeMillis();
            if (containingCI != null) {
                containingCI.getCurrentStatement().setExecutionStartTime(est);
                containingCI.getCurrentStatement().setStackHash(stackHash);
                containingCI.getCurrentStatement().setExecutionCallStack(stack);
                if (args.length > 0) {
                    containingCI.getCurrentStatement().setSqlText((String) args[0]);
                }
            }
            Histogram.Timer timer = JDBCMetrics.SQL_EXECUTION_TIME.startTimer();
            try {
                o = callable.call();
            } finally {
                timer.observeDuration();
            }
            long eet = System.currentTimeMillis();
            log.info("statement@{} {} finished, duration is {}", hashCode, method.getName(), (eet - est));
            if (containingCI != null) {
                containingCI.getCurrentStatement().setExecutionEndTime(eet);
                slowStatements.add(containingCI.getCurrentStatement());
            }
        }
        return o;
    }

    public static String getCurrentThreadStackTrace() throws IOException {
        StringWriter sw = new StringWriter();
        Throwable t = new Throwable("");
        t.printStackTrace(new PrintWriter(sw));
        String stackTrace = sw.toString();
        sw.close();
        return stackTrace;
    }

    @Advice.OnMethodExit
    public static void intercept(@net.bytebuddy.asm.Advice.Origin Constructor<?> m, @This Object inst)
            throws Exception {
        log.debug("----------------- constructor intercept -------------------------");
        if (java.sql.Connection.class.isAssignableFrom(m.getDeclaringClass())) {
            // some CP library override hashCode method
            int hashCode = System.identityHashCode(inst);
            ConnectionInfo ci = new ConnectionInfo();
            ci.setHash(hashCode);
            ci.setCreationTime(System.currentTimeMillis());
            String stackTrace = getCurrentThreadStackTrace();
            String shash = StringUtils.md5(stackTrace);
            stacks.put(shash, stackTrace);
            // ci.setStack(stackTrace);
            ci.setStackHash(shash);
            log.info("connection@{} acquired by stack@{}", hashCode, shash);
            log.debug(stackTrace);
            conns.put(hashCode, ci);
            Integer scount = counstsByStack.get(ci.getStackHash());
            if (scount == null) {
                counstsByStack.put(ci.getStackHash(), 1);
            } else {
                counstsByStack.put(ci.getStackHash(), scount + 1);
            }
        }
    }

}

However, this will not work for some cases. Take HikariCP for example:

HikariCP's implementation has an abstract class ProxyConnection , which has the close method implemented, then a class HikariProxyConnection extends ProxyConnection , which overides some methods, except close . If I setup instrumentation on HikariProxyConnection , the close method will not be intercepted. If I change my code to:

new AgentBuilder.Default()
                .type(startMatcher.and(isSubTypeOf(java.sql.Connection.class).and(isAbstract()))
                .transform(constructorTransformer).transform(methodsTransformer).with(listener).installOn(inst);

it will work for HikariCP, but not for other connection pool implementations.

For my use case, is there a unified way? No matter what connection pool is used, and even with plain JDBC.

Answer 1

Your matcher:

isSubTypeOf(java.sql.Connection.class).and(not(isAbstract()) 

explicitly excludes abstract classes. You would need to instrument all methods to achieve what you wanted if you are using Advice .

Also, it seems like you are blending concepts:

import net.bytebuddy.asm.Advice;
import net.bytebuddy.asm.Advice.This;
import net.bytebuddy.implementation.bind.annotation.AllArguments;
import net.bytebuddy.implementation.bind.annotation.Origin;
import net.bytebuddy.implementation.bind.annotation.RuntimeType;
import net.bytebuddy.implementation.bind.annotation.SuperCall;

The latter annotations belong to MethodDelegation , not to Advice . Refer to the javadoc to see how advice should be applied.