JavaAgent实现数据路由

前言：本篇主要采用JavaAgent实现字节码层面的代码修改；

目标实现：DRDS，MongoDB，KVS，MQ以及ES数据路由（数据路由：请求线程添加标记，判断标记是否存在，以此为依据切换数据源）。

一：JavaAgent实现方法

请参考：https://blog.csdn.net/21aspnet/article/details/81671777

总结的比较清楚，基本使用没有问题

本人初学，简单实现如下：

以DRDS数据路由为例，当前实现Druid数据源的数据路由。需要编写的一共两个部分：

1，agent premain 方法实现入口
public class DruidAgent {
 private static Instrumentation inst = null;
 public static void premain(String agentArgs, Instrumentation _inst) throws ClassNotFoundException, UnmodifiableClassException
 {
 System.out.println("===进入Druid-Agent方法===");
 inst = _inst;
 inst.addTransformer(new DruidTransformer());
 }
}
2，实现ClassFileTransformer的类转换器
public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
 ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
 className = className.replace("/", ".");
 /**判断加载的class的包路径是不是需要监控的类*/
 if (methodMap.containsKey(className)) {
 System.out.println("当前处理类名:" + className);
 CtClass ctclass = null;
 System.out.println("获取pool,开始,classLoader:"+loader.getClass()+","+loader);
 ClassPool pool = PoolUtil.getClassPool(loader);
 try {
 /**使用javassist取得字节码类*/
 ctclass = pool.get(className);
 // 解冻CtClass对象
 ctclass.defrost();
 ctclass.addField(CtField.make("private static com.alibaba.druid.pool.DruidDataSource shadowDS = null ;", ctclass));
 ctclass.addField(CtField.make("private static final java.lang.String CONFIG_PATH_ARG = \"propertyFile\";", ctclass));
 //植入shadowInit方法，其中shadow库配置文件以VM参数-DpropertyFile指定，文件中需要有配置参数。略
/**自行补充*/
 for (String methodName : methodMap.get(className)) {
 System.out.println("当前处理方法名:" + methodName);
 String newMethod = 
 "if(StressTestFlagHolder.isStressTest() && null == druid.DruidFlagHolder.getFlag()){\n" +
 "\t druid.DruidFlagHolder.setFlag(\"DRUID\");\n" +
 "\ttry{\n" +
 "\t\t System.out.println(\"DruidDataSource的CL:\"+com.alibaba.druid.pool.DruidDataSource.class.getClassLoader().getClass()); \n"+
 "\t\t shadowInit();\n" +
 "\t\t return shadowDS.getConnection();\n" +
 "\t\t}finally{\n" +
 "\t\t druid.DruidFlagHolder.clearFlag();\n" +
 "\t } \n" +
 " }";
 /**获取方法实例*/
 CtMethod ctmethod = ctclass.getDeclaredMethod(methodName);
 /**构建方法体*/
 StringBuilder bodyStr = new StringBuilder();
 bodyStr.append(newMethod);
 System.out.println("新方法内容：" + bodyStr.toString());
 ctmethod.insertBefore(bodyStr.toString());// 替换新方法
 }
 return ctclass.toBytecode();
 } catch (Exception e) {
 e.printStackTrace();
 }
 }
 return null;
}

.

二，数据路由切面，需要具体阅读源码，分析代码植入点，以下为本人浅显实现

• DRDS

切入******.druid方法，判断是否存在压测标识，提换数据源

• ES（本方法切入构造方法之后，启动报错，后封装工具层切入实现，原生切入待研究）

index加前缀，拦截CreateIndexRequestBuilder构造器
IndicesExistsRequest构造器
ReplicationRequest的index方法，在index加上es的前缀，前缀在agent的配置文件配置

• Mongo

org.springframework.data.mongodb.core.MongoTemplate.getDb切换数据源

• MQ

植入代码，在com.aliyun.openservices.ons.api.impl.rocketmq.ProducerImpl的send方法中植入：
msg.putUserProperties("StressTestFlag",StressTestFlagHolder.getFlag());
植入代码，在MqMessageListener的consume方法中植 入： StressTestFlagHolder.setFlag(msg.getUserProperties("StressTestFlag"));
StressTestFlagHolder.clearString();

• Redis

切入redis.clients.jedis.BinaryJedis.set 和redis.clients.jedis.BinaryJedis.get方法，key添加前缀

标记类参考：

public class FlagHolder {
 private static final ThreadLocal<String> redisFlagHolder = new ThreadLocal<String>();
 /**添加标记*/
 public static void setFlag(String flag) {
 redisFlagHolder.set(flag);
 }
 /**获取标记*/
 public static String getFlag() {
 String flag = (String)redisFlagHolder.get();
 return flag;
 }
 /**清除标记*/
 public static void clearFlag() {
 redisFlagHolder.remove();
 }
 /**标记判断*/
 public static boolean isStressTest() {
 return "DRUID".equals(getFlag());
 }
}

三，打包启动

本次实现中不再赘述打包过程，有兴趣参考https://blog.csdn.net/xuemengrui12/article/details/74984731

打包结束启动之前：

你需要有：待切入工程jar包，agent代码jar包，javassist依赖包，如果你采用外部配置文件，还需要一份外部配置文件，这样你可以在把agent的一些配置信息放在这里。

指令参考：java -javaagent:E:/Agent/agent.jar -jar -Xmx1g -Xms1g -Xmn512m -XX:SurvivorRatio=8 -DpropertyFile=E:/Agent/外部配置文件.properties 你的服务打包.jar