APM - 零侵入监控JDBC服务

Code

public class JdbcCommonCollects extends AbstractByteTransformCollect implements ICollect {

	public static  JdbcCommonCollects INSTANCE;
	private ApmContext context;

	public JdbcCommonCollects(ApmContext context, Instrumentation instrumentation) {
		super(instrumentation);
		INSTANCE=this;
		this.context=context;
	}


    private final static String[] connection_agent_methods = new String[]{"prepareStatement"};
    private final static String[] prepared_statement_methods = new String[]{"execute", "executeUpdate", "executeQuery"};
    private static final String beginSrc;
    private static final String endSrc;
    private static final String errorSrc;

    static {
        // connect
        beginSrc = "com.artisan.collects.JdbcCommonCollects inst=com.artisan.collects.JdbcCommonCollects.INSTANCE;";
        errorSrc = "inst.error(null,e);";
        endSrc = "result=inst.proxyConnection((java.sql.Connection)result);";
    }

    
    
    public JdbcStatistics begin(String className, String method) {
        JdbcStatistics jdbcStat = new JdbcStatistics();
        jdbcStat.begin=System.currentTimeMillis();
        jdbcStat.setModelType("jdbc");
        return jdbcStat;
    }

    
    
    public void end(JdbcStatistics stat) {
        JdbcStatistics jdbcStat= (JdbcStatistics) stat;
        jdbcStat.end=System.currentTimeMillis();
        jdbcStat.useTime =jdbcStat.end-jdbcStat.begin;
        if (jdbcStat.jdbcUrl != null) {
            jdbcStat.databaseName = getDbName(jdbcStat.jdbcUrl);
        }
        this.context.submitCollectResult(stat);
    }

    public void error(JdbcStatistics stat, Throwable throwable) {
        if (stat != null) {
            stat.error = throwable.getMessage();
            stat.errorType = throwable.getClass().getName();
            if (throwable instanceof InvocationTargetException) {
                stat.errorType = ((InvocationTargetException) throwable).getTargetException().getClass().getName();
                stat.error = ((InvocationTargetException) throwable).getTargetException().getMessage();
            }
        }
    }
    
    public void sendStatistics(JdbcStatistics stat) {
       
    }

    
    public Connection proxyConnection(final Connection connection) {
        Object c = Proxy.newProxyInstance(JdbcCommonCollects.class.getClassLoader()
                , new Class[]{Connection.class}, new ConnectionHandler(connection));
        return (Connection) c;
    }


    public PreparedStatement proxyPreparedStatement(final PreparedStatement statement, JdbcStatistics jdbcStat) {
        Object c = Proxy.newProxyInstance(JdbcCommonCollects.class.getClassLoader()
                , new Class[]{PreparedStatement.class}, new PreparedStatementHandler(statement, jdbcStat));
        return (PreparedStatement) c;
    }

    public byte[] transform(ClassLoader loader, String className) throws Exception {
        // System.out.println("JDBC className:" + className);
        // 如果更換數據庫，需要替換改類  這裡目前僅適配了MySQL
    	 if (!className.equals("com.mysql.cj.jdbc.NonRegisteringDriver")) {
             return null;
         }
    	CtClass ctclass = super.toCtClass(loader, className);
    	AgentByteBuild byteLoade = new AgentByteBuild(className, loader, ctclass);
        CtMethod connectMethod = ctclass.getMethod("connect", "(Ljava/lang/String;Ljava/util/Properties;)Ljava/sql/Connection;");
//      connectMethod.getMethodInfo().getDescriptor();
        AgentByteBuild.MethodSrcBuild build = new AgentByteBuild.MethodSrcBuild();
        build.setBeginSrc(beginSrc);
        build.setErrorSrc(errorSrc);
        build.setEndSrc(endSrc);
        byteLoade.updateMethod(connectMethod, build);
        return byteLoade.toBytecote();
    }


    /**
     * connection 代理处理
     */
    public class ConnectionHandler implements InvocationHandler {
        private final Connection connection;

        private ConnectionHandler(Connection connection) {
            this.connection = connection;
        }

        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
            boolean isTargetMethod = false;
            for (String agentm : connection_agent_methods) {
                if (agentm.equals(method.getName())) {
                    isTargetMethod = true;
                }
            }
            Object result = null;
            JdbcStatistics jdbcStat = null;
            try {
                if (isTargetMethod) { // 获取PreparedStatement 开始统计
                    jdbcStat = (JdbcStatistics) JdbcCommonCollects.this.begin(null, null);
                    jdbcStat.jdbcUrl = connection.getMetaData().getURL();
                    jdbcStat.sql = (String) args[0];
                }
                result = method.invoke(connection, args);
                // 代理 PreparedStatement
                if (isTargetMethod && result instanceof PreparedStatement) {
                    PreparedStatement ps = (PreparedStatement) result;
                    result = proxyPreparedStatement(ps, jdbcStat);
                }
            } catch (Throwable e) {
                JdbcCommonCollects.this.error(jdbcStat, e);
                JdbcCommonCollects.this.end(jdbcStat);
                throw e;
            }
            return result;
        }
    }

    /**
     * PreparedStatement 代理处理
     */
    public class PreparedStatementHandler implements InvocationHandler {
        private final PreparedStatement statement;
        private final JdbcStatistics jdbcStat;

        public PreparedStatementHandler(PreparedStatement statement, JdbcStatistics jdbcStat) {
            this.statement = statement;
            this.jdbcStat = jdbcStat;
        }

        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
            System.out.println("++++++++++++++++++++++++++++++++++++++++++++");
            boolean isTargetMethod = false;
            for (String agentm : prepared_statement_methods) {
                if (agentm.equals(method.getName())) {
                    isTargetMethod = true;
                    break;
                }
            }
            Object result = null;
            try {
                // SQL的參數 輸出  TODO  只能输出一个，后面的会把前面的参数覆盖掉
                if (args != null ){
                    List<JdbcStatistics.ParamValues> paramValues = SplitUtils.split2Parts(args, 2);
                    jdbcStat.setParams(paramValues);
                    //System.out.println("KKKKKKK:" + JsonWriter.objectToJson(jdbcStat));
                }

                result = method.invoke(statement, args);
            } catch (Throwable e) {
                if (isTargetMethod) {
                    JdbcCommonCollects.this.error(jdbcStat, e);
                }
                throw e;
            } finally {
                if (isTargetMethod) {
                    JdbcCommonCollects.this.end(jdbcStat);
                }
            }
            return result;
        }
    }

    
   

    private static String getDbName(String url) {
        int index = url.indexOf("?"); //$NON-NLS-1$
        if (index != -1) {
            String paramString = url.substring(index + 1, url.length());
            url = url.substring(0, index);
        }
        String dbName = url.substring(url.lastIndexOf("/") + 1);
        return dbName;
    }

	
}