Merge pull request #9 from ZhangLin16/zhanglin

#8
Nepxion · Nov 23, 2023 · 0c120d2 · 0c120d2
2 parents 2b503c1 + e040cdc
commit 0c120d2
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 2 deletions.
diff --git a/...agent-starter/src/main/java/com/nepxion/discovery/agent/plugin/thread/ThreadConstant.java b/...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());
-}
+}
diff --git a/...y-agent-starter/src/main/java/com/nepxion/discovery/agent/plugin/thread/ThreadPlugin.java b/...y-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<?>
             }
         }
     }
-}
+
+    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;
+            }
+        }
+    }
+
+}