feat: filter cache mocks under non-application packages (#425)

YongwuHe · web-flow · commit d3f17ac22a8e · 2024-03-15T17:48:16.000+08:00
feat: filter cache mock under non-application packages
diff --git a/arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/util/MapUtils.java b/arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/util/MapUtils.java
@@ -61,4 +61,22 @@ public static <K> boolean getBoolean(final Map<? super K, ?> map, final K key, b
         }
         return defaultValue;
     }
+
+    public static <K> String getString(final Map<? super K, ?> map, final K key, final String defaultValue) {
+        String answer = getString(map, key);
+        if (answer == null) {
+            answer = defaultValue;
+        }
+        return answer;
+    }
+
+    public static <K> String getString(final Map<? super K, ?> map, final K key) {
+        if (map != null) {
+            final Object answer = map.get(key);
+            if (answer != null) {
+                return answer.toString();
+            }
+        }
+        return null;
+    }
 }
diff --git a/arex-agent-bootstrap/src/test/java/io/arex/agent/bootstrap/util/MapUtilsTest.java b/arex-agent-bootstrap/src/test/java/io/arex/agent/bootstrap/util/MapUtilsTest.java
@@ -71,4 +71,19 @@ static Stream<Arguments> getBooleanCase() {
                 arguments(map, "key", asserts1)
         );
     }
+
+    @Test
+    void getString() {
+        // null map
+        assertNull(MapUtils.getString(null, "key"));
+        Map<String, String> map = new HashMap<>();
+        // empty map
+        assertNull(MapUtils.getString(map, "key"));
+        // key exist return value
+        map.put("key", "value");
+        assertEquals("value", MapUtils.getString(map, "key", "default"));
+        // key not exist return default value
+        assertEquals("default", MapUtils.getString(map, "key1", "default"));
+
+    }
 }
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
@@ -44,6 +44,7 @@ public static Config get() {
     private final int recordRate;
     private final String recordVersion;
     private final Set<String> includeServiceOperations;
+    private final String[] coveragePackages;
 
     Config(boolean enableDebug, String serviceName, List<DynamicClassEntity> dynamicClassList,
         Map<String, String> properties,
@@ -57,6 +58,7 @@ public static Config get() {
         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), ',');
         buildDynamicClassInfo();
     }
 
@@ -88,6 +90,9 @@ private void buildDynamicClassInfo() {
         this.dynamicAbstractClassList = list.toArray(StringUtil.EMPTY_STRING_ARRAY);
     }
 
+    public String[] getCoveragePackages() {
+        return coveragePackages;
+    }
     public String getRecordVersion() {
         return recordVersion;
     }
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
@@ -105,4 +105,11 @@ void testBuildExcludeOperation() {
         assertEquals(2, Config.get().excludeServiceOperations().size());
         assertTrue(Config.get().excludeServiceOperations() instanceof ConcurrentHashSet);
     }
+
+    @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]);
+    }
 }
diff --git a/arex-instrumentation-foundation/src/main/java/io/arex/foundation/config/ConfigManager.java b/arex-instrumentation-foundation/src/main/java/io/arex/foundation/config/ConfigManager.java
@@ -387,15 +387,7 @@ private void appendCoveragePackages(String collectCoveragePackages) {
         if (StringUtil.isEmpty(collectCoveragePackages)) {
             return;
         }
-        String defaultPackages = System.getProperty(COVERAGE_PACKAGES);
-        String finalPackages;
-        if (StringUtil.isEmpty(defaultPackages)) {
-            finalPackages = collectCoveragePackages;
-        } else {
-            finalPackages = StringUtil.join(",", defaultPackages, collectCoveragePackages);
-        }
-
-        System.setProperty(COVERAGE_PACKAGES, finalPackages);
+        System.setProperty(COVERAGE_PACKAGES, collectCoveragePackages);
     }
 
     private void publish(Config config) {
diff --git a/arex-instrumentation-foundation/src/test/java/io/arex/foundation/config/ConfigManagerTest.java b/arex-instrumentation-foundation/src/test/java/io/arex/foundation/config/ConfigManagerTest.java
@@ -314,6 +314,6 @@ void appendCoveragePackages() throws Exception {
         // defaultPackage is com.a.b
         System.setProperty(ConfigConstants.COVERAGE_PACKAGES, "com.a.c");
         appendCoveragePackages.invoke(configManager, "com.a.b");
-        assertEquals("com.a.c,com.a.b", System.getProperty(ConfigConstants.COVERAGE_PACKAGES));
+        assertEquals("com.a.b", System.getProperty(ConfigConstants.COVERAGE_PACKAGES));
     }
 }
diff --git a/arex-instrumentation/dynamic/arex-cache/src/main/java/io/arex/inst/cache/caffeine/CaffeineAsyncInstrumentation.java b/arex-instrumentation/dynamic/arex-cache/src/main/java/io/arex/inst/cache/caffeine/CaffeineAsyncInstrumentation.java
@@ -44,7 +44,7 @@ public static boolean onEnter(@Advice.Argument(0) Object key,
             if (ContextManager.needRecord()) {
                 RepeatedCollectManager.enter();
             }
-            if (ContextManager.needReplay()) {
+            if (ContextManager.needReplay() && CacheLoaderUtil.needRecordOrReplay(cacheLoader)) {
                 String className = CacheLoaderUtil.getLocatedClass(cacheLoader);
                 DynamicClassExtractor extractor = new DynamicClassExtractor(className, methodName, new Object[]{key}, methodReturnType);
                 mockResult = extractor.replay();
@@ -68,7 +68,7 @@ public static void onExit(@Advice.Origin("#m") String methodName,
                 }
                 return;
             }
-            if (ContextManager.needRecord() && RepeatedCollectManager.exitAndValidate()) {
+            if (ContextManager.needRecord() && RepeatedCollectManager.exitAndValidate() && CacheLoaderUtil.needRecordOrReplay(cacheLoader)) {
                 String className = CacheLoaderUtil.getLocatedClass(cacheLoader);
                 DynamicClassExtractor extractor = new DynamicClassExtractor(className, methodName, new Object[]{key}, methodReturnType);
                 extractor.recordResponse(throwable != null ? throwable : result);
diff --git a/arex-instrumentation/dynamic/arex-cache/src/main/java/io/arex/inst/cache/caffeine/CaffeineSyncInstrumentation.java b/arex-instrumentation/dynamic/arex-cache/src/main/java/io/arex/inst/cache/caffeine/CaffeineSyncInstrumentation.java
@@ -43,7 +43,7 @@ public static boolean onEnter(@Advice.Argument(0) Object key,
             if (ContextManager.needRecord()) {
                 RepeatedCollectManager.enter();
             }
-            if (ContextManager.needReplay()) {
+            if (ContextManager.needReplay() && CacheLoaderUtil.needRecordOrReplay(cacheLoader)) {
                 String className = CacheLoaderUtil.getLocatedClass(cacheLoader);
                 DynamicClassExtractor extractor = new DynamicClassExtractor(className, methodName, new Object[]{key}, methodReturnType);
                 mockResult = extractor.replay();
@@ -67,7 +67,7 @@ public static void onExit(@Advice.Origin("#m") String methodName,
                 }
                 return;
             }
-            if (ContextManager.needRecord() && RepeatedCollectManager.exitAndValidate()) {
+            if (ContextManager.needRecord() && RepeatedCollectManager.exitAndValidate() && CacheLoaderUtil.needRecordOrReplay(cacheLoader)) {
                 String className = CacheLoaderUtil.getLocatedClass(cacheLoader);
                 DynamicClassExtractor extractor = new DynamicClassExtractor(className, methodName, new Object[]{key}, methodReturnType);
                 extractor.recordResponse(throwable != null ? throwable : result);
diff --git a/arex-instrumentation/dynamic/arex-cache/src/main/java/io/arex/inst/cache/guava/GuavaCacheInstrumentation.java b/arex-instrumentation/dynamic/arex-cache/src/main/java/io/arex/inst/cache/guava/GuavaCacheInstrumentation.java
@@ -13,7 +13,6 @@
 import net.bytebuddy.implementation.bytecode.assign.Assigner;
 import net.bytebuddy.matcher.ElementMatcher;
 
-import java.lang.reflect.Method;
 import java.util.Arrays;
 import java.util.List;
 
@@ -46,7 +45,7 @@ public static boolean onEnter(@Advice.Origin("#m") String methodName,
                 RepeatedCollectManager.enter();
             }
 
-            if (ContextManager.needReplay()) {
+            if (ContextManager.needReplay() && CacheLoaderUtil.needRecordOrReplay(loader)) {
                 String className = CacheLoaderUtil.getLocatedClass(loader);
                 DynamicClassExtractor extractor = new DynamicClassExtractor(className, methodName, new Object[]{key}, methodReturnType);
                 mockResult = extractor.replay();
@@ -72,7 +71,7 @@ public static void onExit(@Advice.Origin("#m") String methodName,
                 }
             }
 
-            if (ContextManager.needRecord() && RepeatedCollectManager.exitAndValidate()) {
+            if (ContextManager.needRecord() && RepeatedCollectManager.exitAndValidate() && CacheLoaderUtil.needRecordOrReplay(loader)) {
                 String className = CacheLoaderUtil.getLocatedClass(loader);
                 DynamicClassExtractor extractor = new DynamicClassExtractor(className, methodName, new Object[]{key}, methodReturnType);
                 extractor.recordResponse(throwable != null ? throwable : result);
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,9 @@
 package io.arex.inst.cache.util;
 
+import io.arex.agent.bootstrap.util.ArrayUtils;
+import io.arex.agent.bootstrap.util.MapUtils;
 import io.arex.agent.bootstrap.util.StringUtil;
+import io.arex.inst.runtime.config.Config;
 import io.arex.inst.runtime.log.LogManager;
 
 import java.lang.reflect.Field;
@@ -74,4 +77,21 @@ private static String getReferenceClass(Object cacheLoader, Object referenceObje
         }
         return referenceObject.getClass().getName();
     }
+
+    public static boolean needRecordOrReplay(Object cacheLoader) {
+        if (cacheLoader == null) {
+            return false;
+        }
+        String[] coveragePackages = Config.get().getCoveragePackages();
+        if (ArrayUtils.isEmpty(coveragePackages)){
+            return false;
+        }
+        String loaderClassName = MapUtils.getString(NO_REFERENCE_MAP, System.identityHashCode(cacheLoader), cacheLoader.getClass().getName());
+        for (String packageName : coveragePackages) {
+            if (StringUtil.startWith(loaderClassName, packageName)) {
+                return true;
+            }
+        }
+        return false;
+    }
 }
diff --git a/arex-instrumentation/dynamic/arex-cache/src/test/java/io/arex/inst/cache/guava/GuavaCacheInstrumentationTest.java b/arex-instrumentation/dynamic/arex-cache/src/test/java/io/arex/inst/cache/guava/GuavaCacheInstrumentationTest.java
@@ -56,6 +56,7 @@ public String load(String key) throws Exception {
                 return "test";
             }
         };
+        Mockito.when(CacheLoaderUtil.needRecordOrReplay(loader)).thenReturn(true);
         assertFalse(GuavaCacheInstrumentation.GetAdvice.onEnter("get", key, StringUtil.EMPTY, loader, null));
 
         // record
@@ -84,6 +85,7 @@ public String load(String key) throws Exception {
                 return "test";
             }
         };
+        Mockito.when(CacheLoaderUtil.needRecordOrReplay(loader)).thenReturn(true);
 
         MockedStatic<MockUtils> mockedStatic = Mockito.mockStatic(MockUtils.class);
         Object result = "test";
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
@@ -1,7 +1,11 @@
 package io.arex.inst.cache.util;
 
 import com.google.common.cache.CacheLoader;
+import io.arex.inst.runtime.config.Config;
+import io.arex.inst.runtime.model.ArexConstants;
 import org.junit.jupiter.api.Test;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
 
 import java.lang.reflect.Field;
 import java.util.Map;
@@ -56,7 +60,29 @@ void getLocatedClass() throws Exception {
         assertEquals("io.arex.inst.cache.util.CacheLoaderUtilTest$AbstractCache2$1", locatedClass);
         assertEquals(2, NO_REFERENCE_MAP.size());
         assertEquals(1, REFERENCE_FIELD_MAP.size());
+    }
 
+    @Test
+    void needRecordOrReplay() {
+        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"});
+            // loader is null
+            assertFalse(CacheLoaderUtil.needRecordOrReplay(null));
+            // loader is in coverage packages
+            assertTrue(CacheLoaderUtil.needRecordOrReplay(cacheLoader));
+
+            // loader is not in coverage packages
+            Mockito.when(config.getCoveragePackages()).thenReturn(new String[]{});
+            assertFalse(CacheLoaderUtil.needRecordOrReplay(cacheLoader));
+
+            Mockito.when(config.getCoveragePackages()).thenReturn(new String[]{"io.arexs"});
+            assertFalse(CacheLoaderUtil.needRecordOrReplay(cacheLoader));
+
+            Mockito.when(config.getCoveragePackages()).thenReturn(new String[]{"io.arexs","io.arex.inst.cache.util"});
+            assertTrue(CacheLoaderUtil.needRecordOrReplay(cacheLoader));
+        }
     }
 
 

Original file line number	Diff line number	Diff line change
`@@ -105,4 +105,11 @@ void testBuildExcludeOperation() {`
`105`	`105`	`assertEquals(2, Config.get().excludeServiceOperations().size());`
`106`	`106`	`assertTrue(Config.get().excludeServiceOperations() instanceof ConcurrentHashSet);`
`107`	`107`	`}`
	`108`	`+`
	`109`	`+ @Test`
	`110`	`+ void buildCoveragePackages() {`
	`111`	`+ ConfigBuilder.create("mock").addProperty(ConfigConstants.COVERAGE_PACKAGES, "io.arex.inst").build();`
	`112`	`+ assertEquals(1, Config.get().getCoveragePackages().length);`
	`113`	`+ assertEquals("io.arex.inst", Config.get().getCoveragePackages()[0]);`
	`114`	`+ }`
`108`	`115`	`}`
Original file line number	Diff line number	Diff line change
`@@ -314,6 +314,6 @@ void appendCoveragePackages() throws Exception {`
`314`	`314`	`// defaultPackage is com.a.b`
`315`	`315`	`System.setProperty(ConfigConstants.COVERAGE_PACKAGES, "com.a.c");`
`316`	`316`	`appendCoveragePackages.invoke(configManager, "com.a.b");`
`317`		`- assertEquals("com.a.c,com.a.b", System.getProperty(ConfigConstants.COVERAGE_PACKAGES));`
	`317`	`+ assertEquals("com.a.b", System.getProperty(ConfigConstants.COVERAGE_PACKAGES));`
`318`	`318`	`}`
`319`	`319`	`}`
Original file line number	Diff line number	Diff line change
`@@ -44,7 +44,7 @@ public static boolean onEnter(@Advice.Argument(0) Object key,`
`44`	`44`	`if (ContextManager.needRecord()) {`
`45`	`45`	`RepeatedCollectManager.enter();`
`46`	`46`	`}`
`47`		`- if (ContextManager.needReplay()) {`
	`47`	`+ if (ContextManager.needReplay() && CacheLoaderUtil.needRecordOrReplay(cacheLoader)) {`
`48`	`48`	`String className = CacheLoaderUtil.getLocatedClass(cacheLoader);`
`49`	`49`	`DynamicClassExtractor extractor = new DynamicClassExtractor(className, methodName, new Object[]{key}, methodReturnType);`
`50`	`50`	`mockResult = extractor.replay();`
`@@ -68,7 +68,7 @@ public static void onExit(@Advice.Origin("#m") String methodName,`
`68`	`68`	`}`
`69`	`69`	`return;`
`70`	`70`	`}`
`71`		`- if (ContextManager.needRecord() && RepeatedCollectManager.exitAndValidate()) {`
	`71`	`+ if (ContextManager.needRecord() && RepeatedCollectManager.exitAndValidate() && CacheLoaderUtil.needRecordOrReplay(cacheLoader)) {`
`72`	`72`	`String className = CacheLoaderUtil.getLocatedClass(cacheLoader);`
`73`	`73`	`DynamicClassExtractor extractor = new DynamicClassExtractor(className, methodName, new Object[]{key}, methodReturnType);`
`74`	`74`	`extractor.recordResponse(throwable != null ? throwable : result);`
Original file line number	Diff line number	Diff line change
`@@ -43,7 +43,7 @@ public static boolean onEnter(@Advice.Argument(0) Object key,`
`43`	`43`	`if (ContextManager.needRecord()) {`
`44`	`44`	`RepeatedCollectManager.enter();`
`45`	`45`	`}`
`46`		`- if (ContextManager.needReplay()) {`
	`46`	`+ if (ContextManager.needReplay() && CacheLoaderUtil.needRecordOrReplay(cacheLoader)) {`
`47`	`47`	`String className = CacheLoaderUtil.getLocatedClass(cacheLoader);`
`48`	`48`	`DynamicClassExtractor extractor = new DynamicClassExtractor(className, methodName, new Object[]{key}, methodReturnType);`
`49`	`49`	`mockResult = extractor.replay();`
`@@ -67,7 +67,7 @@ public static void onExit(@Advice.Origin("#m") String methodName,`
`67`	`67`	`}`
`68`	`68`	`return;`
`69`	`69`	`}`
`70`		`- if (ContextManager.needRecord() && RepeatedCollectManager.exitAndValidate()) {`
	`70`	`+ if (ContextManager.needRecord() && RepeatedCollectManager.exitAndValidate() && CacheLoaderUtil.needRecordOrReplay(cacheLoader)) {`
`71`	`71`	`String className = CacheLoaderUtil.getLocatedClass(cacheLoader);`
`72`	`72`	`DynamicClassExtractor extractor = new DynamicClassExtractor(className, methodName, new Object[]{key}, methodReturnType);`
`73`	`73`	`extractor.recordResponse(throwable != null ? throwable : result);`