[#9595] Refactor AsyncState

pinpoint-apm · Feb 14, 2023 · a7adf9f · a7adf9f
1 parent 2ff90ab
commit a7adf9f
Show file tree

Hide file tree

Showing 11 changed files with 126 additions and 117 deletions.
diff --git a/...s/bootstrap-core/src/main/java/com/navercorp/pinpoint/bootstrap/context/AsyncContext.java b/...s/bootstrap-core/src/main/java/com/navercorp/pinpoint/bootstrap/context/AsyncContext.java
@@ -33,6 +33,8 @@ public interface AsyncContext {
 
     void close();
 
+    boolean finish();
+
 //    void setAttribute(String name, Object o);
 //
 //    Object getAttribute(String name);

diff --git a/...tstrap-core/src/main/java/com/navercorp/pinpoint/bootstrap/context/AsyncContextUtils.java b/...tstrap-core/src/main/java/com/navercorp/pinpoint/bootstrap/context/AsyncContextUtils.java
@@ -5,10 +5,8 @@ private AsyncContextUtils() {
     }
 
     public static boolean asyncStateFinish(final AsyncContext asyncContext) {
-        if (asyncContext instanceof AsyncStateSupport) {
-            final AsyncStateSupport asyncStateSupport = (AsyncStateSupport) asyncContext;
-            asyncStateSupport.finish();
-            return true;
+        if (asyncContext != null) {
+            return asyncContext.finish();
         }
         return false;
     }

diff --git a/...ap-core/src/test/java/com/navercorp/pinpoint/bootstrap/context/AsyncContextUtilsTest.java b/...ap-core/src/test/java/com/navercorp/pinpoint/bootstrap/context/AsyncContextUtilsTest.java
@@ -10,21 +10,20 @@ class AsyncContextUtilsTest {
 
     @Test
     void finish_success() {
-        AsyncState state = mock(AsyncState.class);
-        TestAsyncContext context = mock(TestAsyncContext.class);
-        when(context.getAsyncState()).thenReturn(state);
+        AsyncContext context = mock(AsyncContext.class);
+        when(context.finish()).thenReturn(true);
 
         Assertions.assertTrue(AsyncContextUtils.asyncStateFinish(context));
+
     }
 
     @Test
     void finish_fail() {
         AsyncContext context = mock(AsyncContext.class);
 
         Assertions.assertFalse(AsyncContextUtils.asyncStateFinish(context));
+        Assertions.assertFalse(AsyncContextUtils.asyncStateFinish(null));
     }
 
-    interface TestAsyncContext extends AsyncContext, AsyncStateSupport {
-    }
 
 }
diff --git a/profiler/src/main/java/com/navercorp/pinpoint/profiler/context/AsyncContexts.java b/profiler/src/main/java/com/navercorp/pinpoint/profiler/context/AsyncContexts.java
@@ -0,0 +1,63 @@
+package com.navercorp.pinpoint.profiler.context;
+
+import com.navercorp.pinpoint.bootstrap.context.AsyncContext;
+import com.navercorp.pinpoint.bootstrap.context.AsyncState;
+import com.navercorp.pinpoint.bootstrap.context.Trace;
+import com.navercorp.pinpoint.profiler.context.id.LocalTraceRoot;
+import com.navercorp.pinpoint.profiler.context.id.TraceRoot;
+
+/**
+ * @author Woonduk Kang(emeroad)
+ */
+public final class AsyncContexts {
+    private AsyncContexts() {
+    }
+
+    private final static Remote remote = new Remote();
+    private final static Local local = new Local();
+
+    public static Remote remote() {
+        return remote;
+    }
+
+    public static Local local() {
+        return local;
+    }
+
+    public static class Remote {
+
+        public AsyncContext sync(AsyncTraceContext asyncTraceContext,
+                                 Binder<Trace> binder,
+                                 TraceRoot traceRoot,
+                                 AsyncId asyncId,
+                                 int asyncMethodApiId) {
+            return new DefaultAsyncContext(asyncTraceContext, binder, traceRoot, asyncId, asyncMethodApiId, null);
+        }
+
+        public AsyncContext async(AsyncTraceContext asyncTraceContext,
+                                  Binder<Trace> binder,
+                                  TraceRoot traceRoot,
+                                  AsyncState asyncState,
+                                  AsyncId asyncId,
+                                  int asyncMethodApiId) {
+            return new DefaultAsyncContext(asyncTraceContext, binder, traceRoot, asyncId, asyncMethodApiId, asyncState);
+        }
+    }
+
+    public static class Local {
+
+        public AsyncContext sync(AsyncTraceContext asyncTraceContext,
+                                 Binder<Trace> binder,
+                                 LocalTraceRoot traceRoot) {
+            return new DisableAsyncContext(asyncTraceContext, binder, traceRoot, null);
+        }
+
+        public AsyncContext async(AsyncTraceContext asyncTraceContext,
+                                  Binder<Trace> binder,
+                                  LocalTraceRoot traceRoot,
+                                  AsyncState asyncState) {
+            return new DisableAsyncContext(asyncTraceContext, binder, traceRoot, asyncState);
+        }
+    }
+
+}
diff --git a/profiler/src/main/java/com/navercorp/pinpoint/profiler/context/DefaultAsyncContext.java b/profiler/src/main/java/com/navercorp/pinpoint/profiler/context/DefaultAsyncContext.java
@@ -17,6 +17,7 @@
 package com.navercorp.pinpoint.profiler.context;
 
 import com.navercorp.pinpoint.bootstrap.context.AsyncContext;
+import com.navercorp.pinpoint.bootstrap.context.AsyncState;
 import com.navercorp.pinpoint.bootstrap.context.SpanEventRecorder;
 import com.navercorp.pinpoint.bootstrap.context.Trace;
 import com.navercorp.pinpoint.common.trace.ServiceType;
@@ -25,6 +26,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
+import javax.annotation.Nullable;
 import java.util.Objects;
 
 /**
@@ -43,17 +45,22 @@ public class DefaultAsyncContext implements AsyncContext {
 
     private final int asyncMethodApiId;
 
+    @Nullable
+    private final AsyncState asyncState;
 
-    public DefaultAsyncContext(AsyncTraceContext asyncTraceContext,
-                               Binder<Trace> binder,
-                               TraceRoot traceRoot,
-                               AsyncId asyncId, int asyncMethodApiId) {
+    DefaultAsyncContext(AsyncTraceContext asyncTraceContext,
+                        Binder<Trace> binder,
+                        TraceRoot traceRoot,
+                        AsyncId asyncId,
+                        int asyncMethodApiId,
+                        @Nullable AsyncState asyncState) {
         this.asyncTraceContext = Objects.requireNonNull(asyncTraceContext, "asyncTraceContext");
         this.binder = Objects.requireNonNull(binder, "binder");
         this.traceRoot = Objects.requireNonNull(traceRoot, "traceRoot");
         this.asyncId = Objects.requireNonNull(asyncId, "asyncId");
 
         this.asyncMethodApiId = asyncMethodApiId;
+        this.asyncState = asyncState;
     }
 
 
@@ -129,6 +136,21 @@ public void close() {
         binder.remove();
     }
 
+    @Override
+    public boolean finish() {
+        final AsyncState copy = this.asyncState;
+        if (copy != null) {
+            copy.finish();
+            return true;
+        }
+        return false;
+    }
+
+    @Nullable
+    public AsyncState getAsyncState() {
+        return asyncState;
+    }
+
     @Override
     public String toString() {
         return "DefaultAsyncContext{" +

diff --git a/...ler/src/main/java/com/navercorp/pinpoint/profiler/context/DefaultAsyncContextFactory.java b/...ler/src/main/java/com/navercorp/pinpoint/profiler/context/DefaultAsyncContextFactory.java
@@ -67,7 +67,7 @@ public AsyncContext newAsyncContext(TraceRoot traceRoot, AsyncId asyncId, boolea
         Objects.requireNonNull(asyncId, "asyncId");
 
         if (canSampled) {
-            return new DefaultAsyncContext(asyncTraceContext, binder, traceRoot, asyncId, this.asyncMethodApiId);
+            return AsyncContexts.remote().sync(asyncTraceContext, binder, traceRoot, asyncId, this.asyncMethodApiId);
         } else {
             return newDisableAsyncContext(traceRoot);
         }
@@ -80,7 +80,7 @@ public AsyncContext newAsyncContext(TraceRoot traceRoot, AsyncId asyncId, boolea
         Objects.requireNonNull(asyncState, "asyncState");
 
         if (canSampled) {
-            return new StatefulAsyncContext(asyncTraceContext, binder, traceRoot, asyncId, asyncMethodApiId, asyncState);
+            return AsyncContexts.remote().async(asyncTraceContext, binder, traceRoot, asyncState, asyncId, asyncMethodApiId);
         } else {
             // TODO
             return newDisableAsyncContext(traceRoot, asyncState);
@@ -90,12 +90,12 @@ public AsyncContext newAsyncContext(TraceRoot traceRoot, AsyncId asyncId, boolea
 
     @Override
     public AsyncContext newDisableAsyncContext(LocalTraceRoot traceRoot) {
-        return new DisableAsyncContext(asyncTraceContext, binder, traceRoot);
+        return AsyncContexts.local().sync(asyncTraceContext, binder, traceRoot);
     }
 
     @Override
     public AsyncContext newDisableAsyncContext(LocalTraceRoot traceRoot, AsyncState asyncState) {
-        return new StatefulDisableAsyncContext(asyncTraceContext, binder, traceRoot, asyncState);
+        return AsyncContexts.local().async(asyncTraceContext, binder, traceRoot, asyncState);
     }
 
 }
diff --git a/profiler/src/main/java/com/navercorp/pinpoint/profiler/context/DisableAsyncContext.java b/profiler/src/main/java/com/navercorp/pinpoint/profiler/context/DisableAsyncContext.java
@@ -1,12 +1,14 @@
 package com.navercorp.pinpoint.profiler.context;
 
 import com.navercorp.pinpoint.bootstrap.context.AsyncContext;
+import com.navercorp.pinpoint.bootstrap.context.AsyncState;
 import com.navercorp.pinpoint.bootstrap.context.Trace;
 import com.navercorp.pinpoint.common.util.Assert;
 import com.navercorp.pinpoint.profiler.context.id.LocalTraceRoot;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
+import javax.annotation.Nullable;
 import java.util.Objects;
 
 /**
@@ -17,11 +19,17 @@ public class DisableAsyncContext implements AsyncContext {
     private final LocalTraceRoot traceRoot;
     private final AsyncTraceContext asyncTraceContext;
     private final Binder<Trace> binder;
+    @Nullable
+    private final AsyncState asyncState;
 
-    public DisableAsyncContext(AsyncTraceContext asyncTraceContext, Binder<Trace> binder, LocalTraceRoot traceRoot) {
+    DisableAsyncContext(AsyncTraceContext asyncTraceContext,
+                               Binder<Trace> binder,
+                               LocalTraceRoot traceRoot,
+                               @Nullable AsyncState asyncState) {
         this.asyncTraceContext = Objects.requireNonNull(asyncTraceContext, "asyncTraceContext");
         this.binder = Objects.requireNonNull(binder, "binder");
         this.traceRoot = Objects.requireNonNull(traceRoot, "traceRoot");
+        this.asyncState = asyncState;
     }
 
     @Override
@@ -66,4 +74,15 @@ public Trace currentAsyncTraceObject() {
     public void close() {
         binder.remove();
     }
+
+
+    @Override
+    public boolean finish() {
+        final AsyncState copy = this.asyncState;
+        if (copy != null) {
+            copy.finish();
+            return true;
+        }
+        return false;
+    }
 }
diff --git a/profiler/src/main/java/com/navercorp/pinpoint/profiler/context/StatefulAsyncContext.java b/profiler/src/main/java/com/navercorp/pinpoint/profiler/context/StatefulAsyncContext.java
diff --git a/...er/src/main/java/com/navercorp/pinpoint/profiler/context/StatefulDisableAsyncContext.java b/...er/src/main/java/com/navercorp/pinpoint/profiler/context/StatefulDisableAsyncContext.java
diff --git a/profiler/src/test/java/com/navercorp/pinpoint/profiler/context/DefaultAsyncContextTest.java b/profiler/src/test/java/com/navercorp/pinpoint/profiler/context/DefaultAsyncContextTest.java
@@ -9,9 +9,9 @@ AsyncContext newAsyncContext(boolean canSampled) {
         Binder<Trace> binder = new ThreadLocalBinder<>();
         AsyncTraceContext asyncTraceContext = newAsyncTraceContext();
         if (canSampled) {
-            return new DefaultAsyncContext(asyncTraceContext, binder, traceRoot, asyncId, 0);
+            return AsyncContexts.remote().sync(asyncTraceContext, binder, traceRoot, asyncId, 0);
         } else {
-            return new DisableAsyncContext(asyncTraceContext, binder, traceRoot);
+            return AsyncContexts.local().sync(asyncTraceContext, binder, traceRoot);
         }
     }
 }
diff --git a/profiler/src/test/java/com/navercorp/pinpoint/profiler/context/StatefulAsyncContextTest.java b/profiler/src/test/java/com/navercorp/pinpoint/profiler/context/StatefulAsyncContextTest.java
@@ -19,16 +19,16 @@ AsyncContext newAsyncContext(boolean canSampled) {
         Binder<Trace> binder = new ThreadLocalBinder<>();
         AsyncTraceContext asyncTraceContext = newAsyncTraceContext();
         if (canSampled) {
-            return new StatefulAsyncContext(asyncTraceContext, binder, traceRoot, asyncId, 0, asyncState);
+            return AsyncContexts.remote().async(asyncTraceContext, binder, traceRoot, asyncState, asyncId, 0);
         } else {
-            return new DisableAsyncContext(asyncTraceContext, binder, traceRoot);
+            return AsyncContexts.local().sync(asyncTraceContext, binder, traceRoot);
         }
     }
 
     @Test
     @MockitoSettings(strictness = Strictness.LENIENT)
     public void testGetAsyncState() {
-        StatefulAsyncContext asyncContext = (StatefulAsyncContext) newAsyncContext(true);
+        DefaultAsyncContext asyncContext = (DefaultAsyncContext) newAsyncContext(true);
 
         assertEquals(asyncState, asyncContext.getAsyncState());
     }