fix: failed to record while result type is mono (#343)

Author: DanLi39

arex-instrumentation/dynamic/arex-dynamic-common/src/main/java/io/arex/inst/dynamic/common/DynamicClassExtractor.java

@@ -9,6 +9,7 @@
 import io.arex.agent.bootstrap.util.StringUtil;
 import io.arex.agent.thirdparty.util.time.DateFormatUtils;
 import io.arex.inst.dynamic.common.listener.ListenableFutureAdapter;
+import io.arex.inst.dynamic.common.listener.MonoConsumer;
 import io.arex.inst.dynamic.common.listener.ResponseConsumer;
 import io.arex.inst.runtime.config.Config;
 import io.arex.inst.runtime.context.ArexContext;
@@ -93,7 +94,7 @@ public Object recordResponse(Object response) {
         }
         // Compatible with not import package reactor-core
         if (MONO.equals(methodReturnType) && response instanceof Mono<?>) {
-            return this.resetMonoResponse((Mono<?>) response);
+            return new MonoConsumer(this).accept((Mono<?>) response);
         }
         this.result = response;
         if (needRecord()) {
@@ -155,12 +156,6 @@ void setFutureResponse(Future<?> result) {
         }
     }
 
-    public Mono<?> resetMonoResponse(Mono<?> result) {
-        return result
-            .doOnError(this::recordResponse)
-            .doOnSuccess((Consumer<Object>) this::recordResponse);
-    }
-
     String buildResultClazz(String resultClazz) {
         if (StringUtil.isEmpty(resultClazz) || resultClazz.contains(TypeUtil.HORIZONTAL_LINE_STR)) {
             return resultClazz;

arex-instrumentation/dynamic/arex-dynamic-common/src/main/java/io/arex/inst/dynamic/common/listener/MonoConsumer.java

@@ -0,0 +1,30 @@
+package io.arex.inst.dynamic.common.listener;
+
+import io.arex.agent.bootstrap.ctx.TraceTransmitter;
+import io.arex.inst.dynamic.common.DynamicClassExtractor;
+import reactor.core.publisher.Mono;
+
+public class MonoConsumer {
+
+    private final TraceTransmitter traceTransmitter;
+    private final DynamicClassExtractor extractor;
+
+    public MonoConsumer(DynamicClassExtractor extractor) {
+        this.traceTransmitter = TraceTransmitter.create();
+        this.extractor = extractor;
+    }
+
+    public Mono<?> accept(Mono<?> responseMono) {
+        return responseMono
+            .doOnSuccess(o -> {
+                try (TraceTransmitter tm = traceTransmitter.transmit()) {
+                    extractor.recordResponse(o);
+                }
+            })
+            .doOnError(o -> {
+                try (TraceTransmitter tm = traceTransmitter.transmit()) {
+                    extractor.recordResponse(o);
+                }
+            });
+    }
+}

arex-instrumentation/dynamic/arex-dynamic-common/src/test/java/io/arex/inst/dynamic/common/DynamicClassExtractorTest.java

@@ -5,6 +5,7 @@
 import io.arex.agent.bootstrap.model.ArexMocker;
 import io.arex.agent.bootstrap.model.Mocker.Target;
 import io.arex.agent.thirdparty.util.time.DateFormatUtils;
+import io.arex.inst.dynamic.common.listener.MonoConsumer;
 import io.arex.inst.runtime.config.ConfigBuilder;
 import io.arex.inst.runtime.context.ArexContext;
 import io.arex.inst.runtime.context.ContextManager;
@@ -16,25 +17,15 @@
 import io.arex.inst.runtime.util.MockUtils;
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
-import java.sql.Timestamp;
-import java.text.SimpleDateFormat;
-import java.time.LocalDate;
 import java.time.LocalDateTime;
 import java.time.LocalTime;
-import java.time.format.DateTimeFormatter;
-import java.time.temporal.TemporalAccessor;
 import java.util.Calendar;
 import java.util.Date;
-import java.util.GregorianCalendar;
 import java.util.HashSet;
-import java.util.Locale;
 import java.util.Set;
 import java.util.TimeZone;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
-import java.util.function.Consumer;
-
-import org.joda.time.format.DateTimeFormat;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
@@ -55,8 +46,6 @@
 import org.mockito.stubbing.Answer;
 import reactor.core.publisher.Mono;
 
-import javax.xml.datatype.XMLGregorianCalendar;
-
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNull;
@@ -66,7 +55,6 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.mockStatic;
-import static org.mockito.Mockito.times;
 
 @ExtendWith(MockitoExtension.class)
 class DynamicClassExtractorTest {
@@ -119,15 +107,16 @@ void resetMonoResponse() {
             final DynamicClassExtractor extractor = new DynamicClassExtractor(testWithArexMock, args);
             final Predicate<Object> nonNull = Objects::nonNull;
 
-            //exception
+            // exception
             Mono<?> result = monoExceptionTest();
-            result = extractor.resetMonoResponse(result);
+            MonoConsumer monoConsumer = new MonoConsumer(extractor);
+            result = monoConsumer.accept(result);
             result.subscribe();
             assertTrue(nonNull.test(result));
 
-            //nomal
+            // normal
             result = monoTest();
-            result = extractor.resetMonoResponse(result);
+            result = monoConsumer.accept(result);
             result.subscribe();
             assertTrue(nonNull.test(result));
         } catch (NoSuchMethodException e) {