feat: auto get base packages (#561)

YongwuHe · web-flow · commit da452e47bd84 · 2024-09-09T16:16:53.000+08:00
diff --git a/arex-agent/pom.xml b/arex-agent/pom.xml
@@ -211,6 +211,11 @@
             <artifactId>arex-httpclient-resttemplate</artifactId>
             <version>${project.version}</version>
         </dependency>
+        <dependency>
+            <groupId>${project.groupId}</groupId>
+            <artifactId>arex-component-scan</artifactId>
+            <version>${project.version}</version>
+        </dependency>
         <!--Agent instrumentation end-->
 
         <dependency>
diff --git a/arex-instrumentation-api/src/main/java/io/arex/inst/runtime/config/Config.java b/arex-instrumentation-api/src/main/java/io/arex/inst/runtime/config/Config.java
@@ -5,6 +5,7 @@
 import io.arex.agent.bootstrap.util.StringUtil;
 import io.arex.inst.runtime.context.RecordLimiter;
 import io.arex.inst.runtime.listener.EventProcessor;
+import io.arex.inst.runtime.model.ArexConstants;
 import io.arex.inst.runtime.model.DynamicClassEntity;
 
 import java.util.ArrayList;
@@ -19,7 +20,7 @@
 import static io.arex.agent.bootstrap.constants.ConfigConstants.STORAGE_SERVICE_MODE;
 
 public class Config {
-
+    private static final char SEPARATOR = ',';
     private static Config INSTANCE = null;
 
     static void update(boolean enableDebug, String serviceName, List<DynamicClassEntity> dynamicClassList,
@@ -45,7 +46,7 @@ public static Config get() {
     private final int recordRate;
     private final String recordVersion;
     private final Set<String> includeServiceOperations;
-    private final String[] coveragePackages;
+    private final Set<String> coveragePackages = new ConcurrentHashSet<>();
 
     Config(boolean enableDebug, String serviceName, List<DynamicClassEntity> dynamicClassList,
         Map<String, String> properties,
@@ -58,11 +59,26 @@ public static Config get() {
         this.dubboStreamReplayThreshold = dubboStreamReplayThreshold;
         this.recordRate = recordRate;
         this.recordVersion = properties.get("arex.agent.version");
-        this.includeServiceOperations = StringUtil.splitToSet(properties.get("includeServiceOperations"), ',');
-        this.coveragePackages = StringUtil.split(properties.get(ConfigConstants.COVERAGE_PACKAGES), ',');
+        this.includeServiceOperations = StringUtil.splitToSet(properties.get("includeServiceOperations"), SEPARATOR);
+        buildCoveragePackages(properties);
         buildDynamicClassInfo();
     }
 
+    /**
+     * only no config coverage package, the coverage package is set to spring scan base packages
+     */
+    private void buildCoveragePackages(Map<String, String> properties) {
+        String configCoveragePackages = properties.get(ConfigConstants.COVERAGE_PACKAGES);
+        if (StringUtil.isNotEmpty(configCoveragePackages)) {
+            this.coveragePackages.addAll(StringUtil.splitToSet(configCoveragePackages, SEPARATOR));
+            return;
+        }
+        String scanBasePackages = System.getProperty(ArexConstants.SPRING_SCAN_PACKAGES);
+        if (StringUtil.isNotEmpty(scanBasePackages)) {
+            this.coveragePackages.addAll(StringUtil.splitToSet(scanBasePackages, SEPARATOR));
+        }
+    }
+
     private Set<String> buildExcludeServiceOperations(Set<String> excludeServiceOperations) {
         if (excludeServiceOperations == null) {
             return Collections.emptySet();
@@ -91,9 +107,10 @@ private void buildDynamicClassInfo() {
         this.dynamicAbstractClassList = list.toArray(StringUtil.EMPTY_STRING_ARRAY);
     }
 
-    public String[] getCoveragePackages() {
+    public Set<String> getCoveragePackages() {
         return coveragePackages;
     }
+
     public String getRecordVersion() {
         return recordVersion;
     }
diff --git a/arex-instrumentation-api/src/main/java/io/arex/inst/runtime/model/ArexConstants.java b/arex-instrumentation-api/src/main/java/io/arex/inst/runtime/model/ArexConstants.java
@@ -71,4 +71,5 @@ private ArexConstants() {}
     public static final String MERGE_MOCKER_TYPE = "java.util.ArrayList-io.arex.agent.bootstrap.model.ArexMocker";
     public static final int DB_SQL_MAX_LEN = 5000;
     public static final String DISABLE_SQL_PARSE = "arex.disable.sql.parse";
+    public static final String SPRING_SCAN_PACKAGES = "arex.spring.scan.packages";
 }
diff --git a/arex-instrumentation-api/src/test/java/io/arex/inst/runtime/config/ConfigTest.java b/arex-instrumentation-api/src/test/java/io/arex/inst/runtime/config/ConfigTest.java
@@ -135,8 +135,16 @@ void testBuildExcludeOperation() {
 
     @Test
     void buildCoveragePackages() {
-        ConfigBuilder.create("mock").addProperty(ConfigConstants.COVERAGE_PACKAGES, "io.arex.inst").build();
-        assertEquals(1, Config.get().getCoveragePackages().length);
-        assertEquals("io.arex.inst", Config.get().getCoveragePackages()[0]);
+        ConfigBuilder configBuilder = ConfigBuilder.create("mock").addProperty(ConfigConstants.COVERAGE_PACKAGES, "io.arex.inst");
+        configBuilder.build();
+        assertEquals(1, Config.get().getCoveragePackages().size());
+        assertEquals("io.arex.inst", Config.get().getCoveragePackages().iterator().next());
+        // no config coverage package
+        System.setProperty(ArexConstants.SPRING_SCAN_PACKAGES, "io.arex.inst.spring");
+        ConfigBuilder.create("mock").build();
+        assertEquals(1, Config.get().getCoveragePackages().size());
+        assertEquals("io.arex.inst.spring", Config.get().getCoveragePackages().iterator().next());
+        System.clearProperty(ArexConstants.SPRING_SCAN_PACKAGES);
+
     }
 }
diff --git a/arex-instrumentation/dynamic/arex-cache/src/main/java/io/arex/inst/cache/util/CacheLoaderUtil.java b/arex-instrumentation/dynamic/arex-cache/src/main/java/io/arex/inst/cache/util/CacheLoaderUtil.java
@@ -1,6 +1,6 @@
 package io.arex.inst.cache.util;
 
-import io.arex.agent.bootstrap.util.ArrayUtils;
+import io.arex.agent.bootstrap.util.CollectionUtil;
 import io.arex.agent.bootstrap.util.MapUtils;
 import io.arex.agent.bootstrap.util.StringUtil;
 import io.arex.inst.runtime.config.Config;
@@ -9,6 +9,7 @@
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 
 public class CacheLoaderUtil {
@@ -113,8 +114,8 @@ public static boolean needRecordOrReplay(Object cacheLoader) {
         if (cacheLoader == null) {
             return false;
         }
-        String[] coveragePackages = Config.get().getCoveragePackages();
-        if (ArrayUtils.isEmpty(coveragePackages)){
+        Set<String> coveragePackages = Config.get().getCoveragePackages();
+        if (CollectionUtil.isEmpty(coveragePackages)){
             return false;
         }
         String loaderClassName = getClassName(cacheLoader);
diff --git a/arex-instrumentation/dynamic/arex-cache/src/test/java/io/arex/inst/cache/util/CacheLoaderUtilTest.java b/arex-instrumentation/dynamic/arex-cache/src/test/java/io/arex/inst/cache/util/CacheLoaderUtilTest.java
@@ -9,7 +9,10 @@
 import org.mockito.Mockito;
 
 import java.lang.reflect.Field;
+import java.util.Collections;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 
 import static org.junit.jupiter.api.Assertions.*;
 
@@ -72,7 +75,7 @@ void needRecordOrReplay() throws NoSuchMethodException {
         try (MockedStatic<Config> configMockedStatic = Mockito.mockStatic(Config.class)) {
             Config config = Mockito.mock(Config.class);
             Mockito.when(Config.get()).thenReturn(config);
-            Mockito.when(config.getCoveragePackages()).thenReturn(new String[]{"io.arex.inst"});
+            Mockito.when(config.getCoveragePackages()).thenReturn(Collections.singleton("io.arex.inst.cache.util"));
             // method in coverage packages
             assertTrue(CacheLoaderUtil.needRecordOrReplay(CacheLoaderUtilTest.class.getDeclaredMethod("getLocatedClass")));
 
@@ -82,13 +85,16 @@ void needRecordOrReplay() throws NoSuchMethodException {
             assertTrue(CacheLoaderUtil.needRecordOrReplay(cacheLoader));
 
             // loader is not in coverage packages
-            Mockito.when(config.getCoveragePackages()).thenReturn(new String[]{});
+            Mockito.when(config.getCoveragePackages()).thenReturn(Collections.EMPTY_SET);
             assertFalse(CacheLoaderUtil.needRecordOrReplay(cacheLoader));
 
-            Mockito.when(config.getCoveragePackages()).thenReturn(new String[]{"io.arexs"});
+            Mockito.when(config.getCoveragePackages()).thenReturn(Collections.singleton("io.arexs"));
             assertFalse(CacheLoaderUtil.needRecordOrReplay(cacheLoader));
 
-            Mockito.when(config.getCoveragePackages()).thenReturn(new String[]{"io.arexs","io.arex.inst.cache.util"});
+            Set<String> coveragePackages = new HashSet<>();
+            coveragePackages.add("io.arex.inst.cache.util");
+            coveragePackages.add("io.arexs");
+            Mockito.when(config.getCoveragePackages()).thenReturn(coveragePackages);
             assertTrue(CacheLoaderUtil.needRecordOrReplay(cacheLoader));
         }
     }
diff --git a/arex-instrumentation/httpclient/arex-httpclient-okhttp-v3/src/main/java/io/arex/inst/httpclient/okhttp/v3/OkHttpCallInstrumentation.java b/arex-instrumentation/httpclient/arex-httpclient-okhttp-v3/src/main/java/io/arex/inst/httpclient/okhttp/v3/OkHttpCallInstrumentation.java
@@ -122,4 +122,4 @@ public static void onExit(@Advice.Argument(value = 0) Callback callback,
             }
         }
     }
-}
+}
diff --git a/arex-instrumentation/pom.xml b/arex-instrumentation/pom.xml
@@ -56,6 +56,7 @@
         <module>authentication/arex-shiro</module>
         <module>authentication/arex-jcasbin</module>
         <module>config/arex-apollo</module>
+        <module>spring/arex-component-scan</module>
     </modules>
     <dependencies>
         <dependency>
diff --git a/arex-instrumentation/servlet/arex-httpservlet/src/test/java/io/arex/inst/httpservlet/SpringUtilTest.java b/arex-instrumentation/servlet/arex-httpservlet/src/test/java/io/arex/inst/httpservlet/SpringUtilTest.java
@@ -66,7 +66,6 @@ static Stream<Arguments> regenerateOperationNameCase() {
 
         RequestMappingInfo matchMapping = Mockito.mock(RequestMappingInfo.class);
         PatternsRequestCondition patternsCondition = Mockito.mock(PatternsRequestCondition.class);
-
         Runnable emptyMocker = () -> {};
         Runnable mockApplicationContextInvoke = () -> {
             ApplicationContext applicationContext = Mockito.mock(ApplicationContext.class);
@@ -98,4 +97,4 @@ static Stream<Arguments> regenerateOperationNameCase() {
                 arguments(mockPatterns, httpServletRequest, servletContext, equalPattern)
         );
     }
-}
+}
diff --git a/arex-instrumentation/servlet/arex-httpservlet/src/test/java/io/arex/inst/httpservlet/adapter/impl/ServletAdapterImplV3Test.java b/arex-instrumentation/servlet/arex-httpservlet/src/test/java/io/arex/inst/httpservlet/adapter/impl/ServletAdapterImplV3Test.java
@@ -5,6 +5,7 @@
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.when;
 
+import io.arex.agent.bootstrap.util.ReflectUtil;
 import io.arex.inst.httpservlet.SpringUtil;
 import io.arex.inst.httpservlet.wrapper.CachedBodyRequestWrapperV3;
 import io.arex.inst.httpservlet.wrapper.CachedBodyResponseWrapperV3;
@@ -36,6 +37,7 @@ class ServletAdapterImplV3Test {
 
     @BeforeAll
     static void setUp() {
+        Mockito.mockStatic(ReflectUtil.class);
         Mockito.mockStatic(SpringUtil.class);
     }
 
diff --git a/arex-instrumentation/servlet/arex-httpservlet/src/test/java/io/arex/inst/httpservlet/adapter/impl/ServletAdapterImplV5Test.java b/arex-instrumentation/servlet/arex-httpservlet/src/test/java/io/arex/inst/httpservlet/adapter/impl/ServletAdapterImplV5Test.java
@@ -5,6 +5,7 @@
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.when;
 
+import io.arex.agent.bootstrap.util.ReflectUtil;
 import io.arex.inst.httpservlet.SpringUtil;
 import io.arex.inst.httpservlet.wrapper.CachedBodyRequestWrapperV5;
 import io.arex.inst.httpservlet.wrapper.CachedBodyResponseWrapperV5;
@@ -34,6 +35,7 @@ class ServletAdapterImplV5Test {
 
     @BeforeAll
     static void setUp() {
+        Mockito.mockStatic(ReflectUtil.class);
         Mockito.mockStatic(SpringUtil.class);
     }
 
diff --git a/arex-instrumentation/spring/arex-component-scan/pom.xml b/arex-instrumentation/spring/arex-component-scan/pom.xml
@@ -0,0 +1,35 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <parent>
+        <groupId>io.arex</groupId>
+        <artifactId>arex-instrumentation-parent</artifactId>
+        <version>${revision}</version>
+        <relativePath>../../pom.xml</relativePath>
+    </parent>
+    <modelVersion>4.0.0</modelVersion>
+
+    <artifactId>arex-component-scan</artifactId>
+    <dependencies>
+        <dependency>
+            <groupId>org.springframework</groupId>
+            <artifactId>spring-context</artifactId>
+            <version>5.2.3.RELEASE</version>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.springframework.boot</groupId>
+            <artifactId>spring-boot-autoconfigure</artifactId>
+            <version>2.5.1</version>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.springframework.boot</groupId>
+            <artifactId>spring-boot</artifactId>
+            <version>2.1.8.RELEASE</version>
+            <scope>test</scope>
+        </dependency>
+    </dependencies>
+
+</project>
diff --git a/arex-instrumentation/spring/arex-component-scan/src/main/java/io/arex/inst/spring/SpringComponentScanInstrumentation.java b/arex-instrumentation/spring/arex-component-scan/src/main/java/io/arex/inst/spring/SpringComponentScanInstrumentation.java
@@ -0,0 +1,34 @@
+package io.arex.inst.spring;
+
+import io.arex.inst.extension.MethodInstrumentation;
+import io.arex.inst.extension.TypeInstrumentation;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.springframework.context.ConfigurableApplicationContext;
+
+import java.util.Collections;
+import java.util.List;
+
+import static net.bytebuddy.matcher.ElementMatchers.*;
+
+public class SpringComponentScanInstrumentation extends TypeInstrumentation {
+    @Override
+    protected ElementMatcher<TypeDescription> typeMatcher() {
+        return named("org.springframework.boot.SpringApplication");
+    }
+
+    @Override
+    public List<MethodInstrumentation> methodAdvices() {
+        return Collections.singletonList(new MethodInstrumentation(
+                isMethod().and(named("run").and(takesArguments(1)).and(not(isStatic()))),
+                SpringRunAdvice.class.getName()));
+    }
+
+    public static class SpringRunAdvice {
+        @Advice.OnMethodExit(suppress = Throwable.class)
+        public static void onExit(@Advice.Return ConfigurableApplicationContext context) {
+            SpringUtil.updateScanBasePackages(context);
+        }
+    }
+}
diff --git a/arex-instrumentation/spring/arex-component-scan/src/main/java/io/arex/inst/spring/SpringModuleInstrumentation.java b/arex-instrumentation/spring/arex-component-scan/src/main/java/io/arex/inst/spring/SpringModuleInstrumentation.java
@@ -0,0 +1,27 @@
+package io.arex.inst.spring;
+
+import com.google.auto.service.AutoService;
+import io.arex.agent.bootstrap.util.CollectionUtil;
+import io.arex.inst.extension.ModuleInstrumentation;
+import io.arex.inst.extension.TypeInstrumentation;
+import io.arex.inst.runtime.config.Config;
+
+import java.util.Collections;
+import java.util.List;
+
+@AutoService(ModuleInstrumentation.class)
+public class SpringModuleInstrumentation extends ModuleInstrumentation {
+
+    public SpringModuleInstrumentation() {
+        super("spring-scan");
+    }
+
+    @Override
+    public List<TypeInstrumentation> instrumentationTypes() {
+        // spring-scan module is special, it will be disabled if coveragePackages is not empty
+        if (CollectionUtil.isNotEmpty(Config.get().getCoveragePackages())) {
+            return Collections.emptyList();
+        }
+        return Collections.singletonList(new SpringComponentScanInstrumentation());
+    }
+}
diff --git a/arex-instrumentation/spring/arex-component-scan/src/main/java/io/arex/inst/spring/SpringUtil.java b/arex-instrumentation/spring/arex-component-scan/src/main/java/io/arex/inst/spring/SpringUtil.java
@@ -0,0 +1,60 @@
+package io.arex.inst.spring;
+
+import io.arex.agent.bootstrap.util.CollectionUtil;
+import io.arex.agent.bootstrap.util.StringUtil;
+import io.arex.inst.runtime.config.Config;
+import io.arex.inst.runtime.log.LogManager;
+import io.arex.inst.runtime.model.ArexConstants;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.context.annotation.ComponentScan;
+
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+public class SpringUtil {
+    public static void updateScanBasePackages(ConfigurableApplicationContext context) {
+        try {
+            Map<String, Object> scanAnnotation = context.getBeansWithAnnotation(ComponentScan.class);
+            Set<String> allBasePackages = new HashSet<>();
+            for (Object value : scanAnnotation.values()) {
+                Class<?> targetClass = value.getClass();
+                String name = targetClass.getName();
+                if (name.contains("$$")) {
+                    targetClass = targetClass.getSuperclass();
+                }
+                SpringBootApplication annotation = targetClass.getAnnotation(SpringBootApplication.class);
+                if (annotation != null) {
+                    allBasePackages.add(targetClass.getPackage().getName());
+                    String[] scanBasePackages = annotation.scanBasePackages();
+                    allBasePackages.addAll(normalizeBasePackages(scanBasePackages));
+                }
+                ComponentScan componentScan = targetClass.getAnnotation(ComponentScan.class);
+                if (componentScan != null) {
+                    String[] basePackages = componentScan.basePackages();
+                    allBasePackages.addAll(normalizeBasePackages(basePackages));
+                }
+            }
+            if (allBasePackages.isEmpty()) {
+                return;
+            }
+            System.setProperty(ArexConstants.SPRING_SCAN_PACKAGES, StringUtil.join(allBasePackages, ","));
+            Set<String> coveragePackages = Config.get().getCoveragePackages();
+            // If there is no config coverage package, the coverage package is set to all base packages
+            if (CollectionUtil.isEmpty(coveragePackages)) {
+                coveragePackages.addAll(allBasePackages);
+            }
+        }catch (Exception e) {
+            LogManager.warn("spring.scan.packages", e);
+        }
+    }
+
+    private static Set<String> normalizeBasePackages(String[] basePackages) {
+        Set<String> normalizedPackages = new HashSet<>();
+        for (String basePackage : basePackages) {
+            normalizedPackages.add(StringUtil.replace(basePackage, "*", ""));
+        }
+        return normalizedPackages;
+    }
+}
diff --git a/arex-instrumentation/spring/arex-component-scan/src/test/java/io/arex/inst/spring/SpringComponentScanInstrumentationTest.java b/arex-instrumentation/spring/arex-component-scan/src/test/java/io/arex/inst/spring/SpringComponentScanInstrumentationTest.java
diff --git a/arex-instrumentation/spring/arex-component-scan/src/test/java/io/arex/inst/spring/SpringModuleInstrumentationTest.java b/arex-instrumentation/spring/arex-component-scan/src/test/java/io/arex/inst/spring/SpringModuleInstrumentationTest.java
diff --git a/arex-instrumentation/spring/arex-component-scan/src/test/java/io/arex/inst/spring/SpringUtilTest.java b/arex-instrumentation/spring/arex-component-scan/src/test/java/io/arex/inst/spring/SpringUtilTest.java
diff --git a/pom.xml b/pom.xml

Original file line number	Diff line number	Diff line change
`@@ -71,4 +71,5 @@ private ArexConstants() {}`
`71`	`71`	`public static final String MERGE_MOCKER_TYPE = "java.util.ArrayList-io.arex.agent.bootstrap.model.ArexMocker";`
`72`	`72`	`public static final int DB_SQL_MAX_LEN = 5000;`
`73`	`73`	`public static final String DISABLE_SQL_PARSE = "arex.disable.sql.parse";`
	`74`	`+ public static final String SPRING_SCAN_PACKAGES = "arex.spring.scan.packages";`
`74`	`75`	`}`
Original file line number	Diff line number	Diff line change
`@@ -122,4 +122,4 @@ public static void onExit(@Advice.Argument(value = 0) Callback callback,`
`122`	`122`	`}`
`123`	`123`	`}`
`124`	`124`	`}`
`125`		`-}`
	`125`	`+}`