diff --git a/discovery-agent-starter/src/main/java/com/nepxion/discovery/agent/plugin/thread/ThreadConstant.java b/discovery-agent-starter/src/main/java/com/nepxion/discovery/agent/plugin/thread/ThreadConstant.java index 93c3cd3..f767dc1 100644 --- a/discovery-agent-starter/src/main/java/com/nepxion/discovery/agent/plugin/thread/ThreadConstant.java +++ b/discovery-agent-starter/src/main/java/com/nepxion/discovery/agent/plugin/thread/ThreadConstant.java @@ -18,9 +18,10 @@ public class ThreadConstant { public static final String THREAD_REQUEST_DECORATOR_ENABLED = "thread.request.decorator.enabled"; public static final String RUNNABLE_CLASS_NAME = "java.lang.Runnable"; public static final String CALLABLE_CLASS_NAME = "java.util.concurrent.Callable"; + public static final String SUPPLIER_CLASS_NAME = "java.util.function.Supplier"; public static final String THREAD_SCAN_PACKAGES_DELIMITERS = ";"; public static final String CONSTRUCTOR_INTERCEPTOR = String.format("%s.before(this);\n", ThreadConstructorInterceptor.class.getName()); public static final String RUN_BEFORE_INTERCEPTOR = String.format("%s.before(this);\n", ThreadCallInterceptor.class.getName()); public static final String RUN_AFTER_INTERCEPTOR = String.format("%s.after(this);\n", ThreadCallInterceptor.class.getName()); -} \ No newline at end of file +} diff --git a/discovery-agent-starter/src/main/java/com/nepxion/discovery/agent/plugin/thread/ThreadPlugin.java b/discovery-agent-starter/src/main/java/com/nepxion/discovery/agent/plugin/thread/ThreadPlugin.java index 4c33014..150f94e 100644 --- a/discovery-agent-starter/src/main/java/com/nepxion/discovery/agent/plugin/thread/ThreadPlugin.java +++ b/discovery-agent-starter/src/main/java/com/nepxion/discovery/agent/plugin/thread/ThreadPlugin.java @@ -61,11 +61,14 @@ public void install(TransformTemplate transformTemplate) { RunnableTransformCallback runnableTransformCallback = new RunnableTransformCallback(); CallableTransformCallback callableTransformCallback = new CallableTransformCallback(); + SupplierTransformCallback supplierTransformCallback = new SupplierTransformCallback(); for (String basePackage : scanPackages) { MatcherOperator runnableInterfaceMatcherOperator = MatcherFactory.newPackageBasedMatcher(basePackage, ThreadConstant.RUNNABLE_CLASS_NAME); MatcherOperator callableInterfaceMatcherOperator = MatcherFactory.newPackageBasedMatcher(basePackage, ThreadConstant.CALLABLE_CLASS_NAME); + MatcherOperator supplierInterfaceMatcherOperator = MatcherFactory.newPackageBasedMatcher(basePackage, ThreadConstant.SUPPLIER_CLASS_NAME); transformTemplate.transform(runnableInterfaceMatcherOperator, runnableTransformCallback); transformTemplate.transform(callableInterfaceMatcherOperator, callableTransformCallback); + transformTemplate.transform(supplierInterfaceMatcherOperator, supplierTransformCallback); } LOG.info(String.format("%s install successfully", this.getClass().getSimpleName())); @@ -169,4 +172,34 @@ public byte[] doInTransform(ClassLoader classLoader, String className, Class } } } -} \ No newline at end of file + + public static class SupplierTransformCallback implements TransformCallback { + @Override + public byte[] doInTransform(ClassLoader classLoader, String className, Class classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) { + try { + ClassInfo classInfo = new ClassInfo(className, classfileBuffer, classLoader); + CtClass ctClass = classInfo.getCtClass(); + + addField(ctClass, AsyncContextAccessor.class); + + CtConstructor[] declaredConstructors = ctClass.getDeclaredConstructors(); + for (CtConstructor ctConstructor : declaredConstructors) { + ctConstructor.insertAfter(ThreadConstant.CONSTRUCTOR_INTERCEPTOR); + } + + CtMethod ctMethod = ctClass.getDeclaredMethod("get"); + if (null != ctMethod) { + ctMethod.insertBefore(ThreadConstant.RUN_BEFORE_INTERCEPTOR); + ctMethod.insertAfter(ThreadConstant.RUN_AFTER_INTERCEPTOR); + } + + return ctClass.toBytecode(); + } catch (Exception e) { + LOG.warn(String.format("Transform %s error, message:", className), e); + + return null; + } + } + } + +}