feat: init class before replay (#569)

YongwuHe · web-flow · commit a9e2d9ddd689 · 2024-09-26T17:03:28.000+08:00
diff --git a/arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/constants/ConfigConstants.java b/arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/constants/ConfigConstants.java
@@ -46,4 +46,5 @@ private ConfigConstants() {
     public static final String API_TOKEN = "arex.api.token";
     public static final String MOCKER_TAGS = "arex.mocker.tags";
     public static final String LOG_PATH = "arex.log.path";
+    public static final String AREX_STATIC_CLASS_INIT = "arex.static.class.init";
 }
diff --git a/arex-instrumentation-api/src/main/java/io/arex/inst/runtime/listener/EventProcessor.java b/arex-instrumentation-api/src/main/java/io/arex/inst/runtime/listener/EventProcessor.java
@@ -1,11 +1,13 @@
 package io.arex.inst.runtime.listener;
 
 import io.arex.agent.bootstrap.cache.TimeCache;
+import io.arex.agent.bootstrap.constants.ConfigConstants;
 import io.arex.agent.bootstrap.model.Mocker;
 import io.arex.agent.bootstrap.util.AdviceClassesCollector;
 import io.arex.agent.bootstrap.util.NumberUtil;
 import io.arex.agent.bootstrap.util.StringUtil;
 import io.arex.agent.bootstrap.util.ServiceLoader;
+import io.arex.inst.runtime.config.Config;
 import io.arex.inst.runtime.model.InitializeEnum;
 import io.arex.inst.runtime.request.RequestHandlerManager;
 import io.arex.inst.runtime.log.LogManager;
@@ -18,6 +20,7 @@
 
 import java.util.List;
 
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -119,6 +122,7 @@ public static void onRequest(){
             ClassLoader contextClassLoader = Thread.currentThread().getContextClassLoader();
             // https://bugs.openjdk.org/browse/JDK-8172726
             CompletableFuture.runAsync(() -> {
+                initClass(contextClassLoader);
                 initSerializer(contextClassLoader);
                 initLog(contextClassLoader);
                 RequestHandlerManager.init(contextClassLoader);
@@ -129,6 +133,42 @@ public static void onRequest(){
         ContextManager.remove();
     }
 
+    /**
+     * To prevent class initialization failure during request replay,
+     * it is advisable to perform early initialization of classes that rely on external dependencies.
+     * This helps to avoid triggering the class initialization logic when replaying requests.
+     */
+    private static void initClass(ClassLoader contextClassLoader) {
+        boolean disableReplay = Config.get().getBoolean(ConfigConstants.DISABLE_REPLAY, false);
+        if (disableReplay) {
+            return;
+        }
+        String arexStaticClassInit = Config.get().getString(ConfigConstants.AREX_STATIC_CLASS_INIT);
+        if (StringUtil.isBlank(arexStaticClassInit)) {
+            return;
+        }
+        Set<String> initClassNameSet = StringUtil.splitToSet(arexStaticClassInit, ',');
+        for (String initClass : initClassNameSet) {
+            try {
+                if (isClassLoaded(initClass, contextClassLoader)) {
+                    continue;
+                }
+                Class.forName(initClass, true, contextClassLoader);
+            } catch (ClassNotFoundException e) {
+                LOGGER.warn(LogManager.buildTitle("init.class"), e);
+            }
+        }
+    }
+
+    private static boolean isClassLoaded(String className, ClassLoader classLoader) {
+        try {
+            Class<?> clazz = classLoader.loadClass(className);
+            return clazz != null;
+        } catch (ClassNotFoundException e) {
+            return false;
+        }
+    }
+
     private static void initLog(ClassLoader contextClassLoader) {
         List<Logger> extensionLoggerList = ServiceLoader.load(Logger.class, contextClassLoader);
         LogManager.build(extensionLoggerList);
diff --git a/arex-instrumentation-api/src/test/java/io/arex/inst/runtime/listener/EventProcessorTest.java b/arex-instrumentation-api/src/test/java/io/arex/inst/runtime/listener/EventProcessorTest.java
@@ -7,17 +7,19 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
+import io.arex.agent.bootstrap.constants.ConfigConstants;
 import io.arex.agent.bootstrap.util.StringUtil;
+import io.arex.inst.runtime.config.ConfigBuilder;
 import io.arex.inst.runtime.context.ArexContext;
 import io.arex.inst.runtime.context.ContextManager;
 import io.arex.inst.runtime.log.LogManager;
 import io.arex.inst.runtime.log.Logger;
-import io.arex.inst.runtime.serializer.Serializer;
 import io.arex.inst.runtime.serializer.StringSerializable;
 import io.arex.agent.bootstrap.util.ServiceLoader;
+
+import java.lang.reflect.Method;
 import java.lang.reflect.Type;
-import java.util.Arrays;
-import java.util.Collections;
+
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
@@ -33,12 +35,16 @@ public class EventProcessorTest {
     static Logger logger = null;
     static MockedStatic<LogManager> mockedStatic = null;
     static MockedStatic<ContextManager> contextMockedStatic = null;
+    static ConfigBuilder configBuilder;
     @BeforeAll
     static void setUp() {
         contextMockedStatic = Mockito.mockStatic(ContextManager.class);
         Mockito.mockStatic(ServiceLoader.class);
         logger = Mockito.mock(Logger.class);
         mockedStatic = Mockito.mockStatic(LogManager.class);
+        configBuilder = ConfigBuilder.create("testInitClass");
+        // disPlay
+        configBuilder.build();
     }
 
     @AfterAll
@@ -47,6 +53,8 @@ static void tearDown() {
         mockedStatic = null;
         contextMockedStatic = null;
         Mockito.clearAllCaches();
+        configBuilder.build();
+        configBuilder = null;
     }
 
     @Test
@@ -80,6 +88,27 @@ void testInit() throws InterruptedException {
         contextMockedStatic.verify(() -> ContextManager.remove(), Mockito.times(1));
     }
 
+    @Test
+    @Order(3)
+    void testInitClass() throws Exception {
+        Method initClassMethod = EventProcessor.class.getDeclaredMethod("initClass", ClassLoader.class);
+        initClassMethod.setAccessible(true);
+        // disPlay
+        initClassMethod.invoke(null, Thread.currentThread().getContextClassLoader());
+        mockedStatic.verify(() -> LogManager.buildTitle("init.class"), Mockito.never());
+        // not display but no initClass
+        configBuilder.addProperty(ConfigConstants.DISABLE_REPLAY, "true");
+        configBuilder.build();
+        initClassMethod.invoke(null, Thread.currentThread().getContextClassLoader());
+        mockedStatic.verify(() -> LogManager.buildTitle("init.class"), Mockito.never());
+        // no display and initClass
+        configBuilder.addProperty(ConfigConstants.DISABLE_REPLAY, "false");
+        configBuilder.addProperty(ConfigConstants.AREX_STATIC_CLASS_INIT, "test1,io.arex.inst.runtime.listener.EventProcessorTest");
+        configBuilder.build();
+        initClassMethod.invoke(null, Thread.currentThread().getContextClassLoader());
+        mockedStatic.verify(() -> LogManager.buildTitle("init.class"), Mockito.times(1));
+    }
+
     @Test
     void onExit() {
         Assertions.assertDoesNotThrow(EventProcessor::onExit);
diff --git a/arex-instrumentation-foundation/src/test/java/io/arex/foundation/logger/AgentLoggerFactoryTest.java b/arex-instrumentation-foundation/src/test/java/io/arex/foundation/logger/AgentLoggerFactoryTest.java
@@ -4,6 +4,9 @@
 
 import io.arex.agent.bootstrap.constants.ConfigConstants;
 import java.io.IOException;
+import java.io.PrintStream;
+import java.lang.reflect.Field;
+import java.lang.reflect.Method;
 import java.nio.file.Files;
 import java.nio.file.Paths;
 import java.time.LocalDate;
@@ -16,8 +19,11 @@
 class AgentLoggerFactoryTest {
 
     @Test
-    void getAgentLogger() throws IOException {
+    void getAgentLogger() throws Exception {
         System.setProperty(ConfigConstants.LOG_PATH, "/var/tmp");
+        Method getPrintStream = AgentLoggerFactory.class.getDeclaredMethod("getPrintStream", String.class);
+        getPrintStream.setAccessible(true);
+        getPrintStream.invoke(null, "startup");
         AgentLogger agentLogger = AgentLoggerFactory.getAgentLogger(AgentLoggerFactoryTest.class);
         assertEquals(AgentLoggerFactoryTest.class.getName(), agentLogger.getName());
 

Original file line number	Diff line number	Diff line change
`@@ -46,4 +46,5 @@ private ConfigConstants() {`
`46`	`46`	`public static final String API_TOKEN = "arex.api.token";`
`47`	`47`	`public static final String MOCKER_TAGS = "arex.mocker.tags";`
`48`	`48`	`public static final String LOG_PATH = "arex.log.path";`
	`49`	`+ public static final String AREX_STATIC_CLASS_INIT = "arex.static.class.init";`
`49`	`50`	`}`