feat: remove the trace of case on service exit (#232)

Author: mr3

arex-instrumentation-api/src/main/java/io/arex/inst/runtime/listener/CaseEvent.java

@@ -3,8 +3,22 @@
 import java.util.EventObject;
 
 public class CaseEvent extends EventObject {
-    public enum Action { ENTER, CREATE, DESTROY }
-    CaseEvent.Action action;
+    public enum Action {
+        /*
+         * Mainly used to clean up the previous case on service entrance
+         */
+        ENTER,
+        /*
+         * Mainly used to create a new trace for the current case
+         */
+        CREATE,
+        /*
+         * Mainly used to clean up the trace of the current case on service exit
+         */
+        EXIT
+    }
+
+    private final CaseEvent.Action action;
 
     private CaseEvent(EventSource source, CaseEvent.Action action) {
         super(source);
@@ -19,7 +33,11 @@ public static CaseEvent ofCreateEvent(EventSource source) {
         return new CaseEvent(source, Action.CREATE);
     }
 
-    public static CaseEvent ofDestroyEvent() {
-        return new CaseEvent(EventSource.empty(), CaseEvent.Action.DESTROY);
+    public static CaseEvent ofExitEvent() {
+        return new CaseEvent(EventSource.empty(), CaseEvent.Action.EXIT);
+    }
+
+    public CaseEvent.Action getAction() {
+        return action;
     }
 }

arex-instrumentation-api/src/main/java/io/arex/inst/runtime/listener/CaseEventDispatcher.java

@@ -2,15 +2,15 @@
 
 public class CaseEventDispatcher {
 
-    public static void onEvent(CaseEvent e) {
-        switch (e.action) {
+    public static void onEvent(CaseEvent event) {
+        switch (event.getAction()) {
             case ENTER:
                 EventProcessor.onRequest();
                  break;
             case CREATE:
-                EventProcessor.onCreate((EventSource) e.getSource());
+                EventProcessor.onCreate((EventSource) event.getSource());
                 break;
-            case DESTROY:
+            case EXIT:
                 EventProcessor.onExit();
                 break;
             default:

arex-instrumentation-api/src/main/java/io/arex/inst/runtime/listener/EventProcessor.java

@@ -88,6 +88,8 @@ private static void initClock(){
     }
 
     public static void onExit(){
+        TimeCache.remove();
+        ContextManager.remove();
     }
 
     /**

arex-instrumentation-api/src/test/java/io/arex/inst/runtime/listener/CaseEventDispatcherTest.java

@@ -0,0 +1,36 @@
+package io.arex.inst.runtime.listener;
+
+import static org.mockito.Mockito.times;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+
+class CaseEventDispatcherTest {
+
+    @BeforeEach
+    void setUp() {
+    }
+
+    @AfterEach
+    void tearDown() {
+    }
+
+    @Test
+    void onEvent() {
+        try (MockedStatic<EventProcessor> mockedStatic = Mockito.mockStatic(EventProcessor.class)) {
+            // test enter event
+            CaseEventDispatcher.onEvent(CaseEvent.ofEnterEvent());
+            mockedStatic.verify(EventProcessor::onRequest, times(1));
+
+            // test create event
+            CaseEventDispatcher.onEvent(CaseEvent.ofCreateEvent(EventSource.empty()));
+            mockedStatic.verify(() -> EventProcessor.onCreate(EventSource.empty()), times(1));
+
+            // test exit event
+            CaseEventDispatcher.onEvent(CaseEvent.ofExitEvent());
+            mockedStatic.verify(EventProcessor::onExit, times(1));
+        }
+    }
+}

arex-instrumentation-api/src/test/java/io/arex/inst/runtime/listener/EventProcessorTest.java

@@ -92,6 +92,11 @@ void testInit() {
         Assertions.assertNotNull(Serializer.getINSTANCE());
     }
 
+    @Test
+    void onExit() {
+        Assertions.assertDoesNotThrow(EventProcessor::onExit);
+    }
+
     public static class TestStringSerializable implements StringSerializable {
 
         @Override
@@ -152,4 +157,4 @@ public StringSerializable reCreateSerializer() {
             return null;
         }
     }
-}
+}

arex-instrumentation/dubbo/arex-dubbo-alibaba/src/main/java/io/arex/inst/dubbo/alibaba/DubboProviderExtractor.java

@@ -36,6 +36,7 @@ public static void onServiceExit(Invoker<?> invoker, Invocation invocation, Resu
         RequestHandlerManager.postHandle(invocation.getAttachments(), result != null ? result.getAttachments() : null,
                 MockCategoryType.DUBBO_PROVIDER.getName());
         adapter.execute(result, makeMocker(adapter));
+        CaseEventDispatcher.onEvent(CaseEvent.ofExitEvent());
     }
     private static Mocker makeMocker(DubboAdapter adapter) {
         Mocker mocker = MockUtils.createDubboProvider(adapter.getServiceOperation());

arex-instrumentation/dubbo/arex-dubbo-apache-v2/src/main/java/io/arex/inst/dubbo/apache/v2/DubboProviderExtractor.java

@@ -39,6 +39,7 @@ public static void onServiceExit(Invoker<?> invoker, Invocation invocation, Resu
         RequestHandlerManager.postHandle(invocation.getAttachments(), RpcContext.getServerContext().getAttachments(),
                 MockCategoryType.DUBBO_PROVIDER.getName());
         adapter.execute(result, makeMocker(adapter));
+        CaseEventDispatcher.onEvent(CaseEvent.ofExitEvent());
     }
     private static Mocker makeMocker(DubboAdapter adapter) {
         Mocker mocker = MockUtils.createDubboProvider(adapter.getServiceOperation());

arex-instrumentation/dubbo/arex-dubbo-apache-v3/src/main/java/io/arex/inst/dubbo/apache/v3/DubboProviderExtractor.java

@@ -36,6 +36,7 @@ public static void onServiceExit(Invoker<?> invoker, Invocation invocation, Resu
         RequestHandlerManager.postHandle(invocation.getAttachments(), RpcContext.getServerAttachment().getAttachments(),
                 MockCategoryType.DUBBO_PROVIDER.getName());
         adapter.execute(result, makeMocker(adapter));
+        CaseEventDispatcher.onEvent(CaseEvent.ofExitEvent());
     }
     private static Mocker makeMocker(DubboAdapter adapter) {
         Mocker mocker = MockUtils.createDubboProvider(adapter.getServiceOperation());

arex-instrumentation/netty/arex-netty-v4/src/main/java/io/arex/inst/netty/v4/server/RequestTracingHandler.java

@@ -25,6 +25,7 @@ public class RequestTracingHandler extends ChannelInboundHandlerAdapter {
     public void channelRead(ChannelHandlerContext ctx, Object msg) {
         EventProcessor.onRequest();
         if (msg instanceof HttpRequest) {
+            CaseEventDispatcher.onEvent(CaseEvent.ofEnterEvent());
             HttpRequest request = (HttpRequest) msg;
             String caseId = request.headers().get(ArexConstants.RECORD_ID);
             if (shouldSkip(request, caseId)) {

arex-instrumentation/netty/arex-netty-v4/src/main/java/io/arex/inst/netty/v4/server/ResponseTracingHandler.java

@@ -86,6 +86,6 @@ private void invoke(final Channel channel, final String content) {
             MockUtils.recordMocker(mocker);
         }
 
-        CaseEventDispatcher.onEvent(CaseEvent.ofEnterEvent());
+        CaseEventDispatcher.onEvent(CaseEvent.ofExitEvent());
     }
 }