[pinpoint-apm#9595] Refactor AsyncContext

emeroad · Feb 15, 2023 · 27092eb · 27092eb
1 parent 2ff90ab
commit 27092eb
Show file tree

Hide file tree

Showing 12 changed files with 146 additions and 128 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,72 @@
+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;
+
+import java.util.Objects;
+
+/**
+ * @author Woonduk Kang(emeroad)
+ */
+public final class AsyncContexts {
+    private AsyncContexts() {
+    }
+
+    public static Remote remote(AsyncTraceContext asyncTraceContext,
+                                Binder<Trace> binder,
+                                int asyncMethodApiId) {
+        return new Remote(asyncTraceContext, binder, asyncMethodApiId);
+    }
+
+    public static Local local(AsyncTraceContext asyncTraceContext,
+                              Binder<Trace> binder) {
+        return new Local(asyncTraceContext, binder);
+    }
+
+    public static class Remote {
+        private final AsyncTraceContext asyncTraceContext;
+        private final Binder<Trace> binder;
+        private final int asyncMethodApiId;
+
+        public Remote(AsyncTraceContext asyncTraceContext, Binder<Trace> binder, int asyncMethodApiId) {
+            this.asyncTraceContext = Objects.requireNonNull(asyncTraceContext, "asyncTraceContext");
+            this.binder = Objects.requireNonNull(binder, "binder");
+            this.asyncMethodApiId = asyncMethodApiId;
+        }
+
+        public AsyncContext sync(TraceRoot traceRoot,
+                                 AsyncId asyncId) {
+            return new DefaultAsyncContext(asyncTraceContext, binder, traceRoot, asyncId, asyncMethodApiId, null);
+        }
+
+        public AsyncContext async(TraceRoot traceRoot,
+                                  AsyncState asyncState,
+                                  AsyncId asyncId) {
+            return new DefaultAsyncContext(asyncTraceContext, binder, traceRoot, asyncId, asyncMethodApiId, asyncState);
+        }
+    }
+
+    public static class Local {
+
+        private final AsyncTraceContext asyncTraceContext;
+        private final Binder<Trace> binder;
+
+        public Local(AsyncTraceContext asyncTraceContext, Binder<Trace> binder) {
+            this.asyncTraceContext = Objects.requireNonNull(asyncTraceContext, "asyncTraceContext");
+            this.binder = Objects.requireNonNull(binder, "binder");
+        }
+
+        public AsyncContext sync(LocalTraceRoot traceRoot) {
+            return new DisableAsyncContext(asyncTraceContext, binder, traceRoot, null);
+        }
+
+        public AsyncContext async(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
@@ -32,23 +32,18 @@
  */
 public class DefaultAsyncContextFactory implements AsyncContextFactory {
 
-    private final AsyncTraceContext asyncTraceContext;
-    private final Binder<Trace> binder;
+    private final AsyncContexts.Remote remote;
+    private final AsyncContexts.Local local;
     private final AsyncIdGenerator asyncIdGenerator;
-    private final PredefinedMethodDescriptorRegistry predefinedMethodDescriptorRegistry;
-    private final int asyncMethodApiId;
-
     public DefaultAsyncContextFactory(AsyncTraceContext asyncTraceContext,
                                       Binder<Trace> binder,
                                       AsyncIdGenerator asyncIdGenerator,
                                       PredefinedMethodDescriptorRegistry predefinedMethodDescriptorRegistry) {
-        this.asyncTraceContext = Objects.requireNonNull(asyncTraceContext, "traceFactoryProvider");
-        this.asyncIdGenerator = Objects.requireNonNull(asyncIdGenerator, "asyncIdGenerator");
-        this.binder = Objects.requireNonNull(binder, "binder");
+        int asyncMethodApiId = getAsyncMethodApiId(predefinedMethodDescriptorRegistry);
+        this.remote = AsyncContexts.remote(asyncTraceContext, binder, asyncMethodApiId);
+        this.local = AsyncContexts.local(asyncTraceContext, binder);
 
-        this.predefinedMethodDescriptorRegistry = Objects.requireNonNull(predefinedMethodDescriptorRegistry, "predefinedMethodDescriptorRegistry");
-
-        this.asyncMethodApiId = getAsyncMethodApiId(predefinedMethodDescriptorRegistry);
+        this.asyncIdGenerator = Objects.requireNonNull(asyncIdGenerator, "asyncIdGenerator");
     }
 
     private int getAsyncMethodApiId(PredefinedMethodDescriptorRegistry predefinedMethodDescriptorRegistry) {
@@ -67,7 +62,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 remote.sync(traceRoot, asyncId);
         } else {
             return newDisableAsyncContext(traceRoot);
         }
@@ -80,7 +75,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 remote.async(traceRoot, asyncState, asyncId);
         } else {
             // TODO
             return newDisableAsyncContext(traceRoot, asyncState);
@@ -90,12 +85,12 @@ public AsyncContext newAsyncContext(TraceRoot traceRoot, AsyncId asyncId, boolea
 
     @Override
     public AsyncContext newDisableAsyncContext(LocalTraceRoot traceRoot) {
-        return new DisableAsyncContext(asyncTraceContext, binder, traceRoot);
+        return local.sync(traceRoot);
     }
 
     @Override
     public AsyncContext newDisableAsyncContext(LocalTraceRoot traceRoot, AsyncState asyncState) {
-        return new StatefulDisableAsyncContext(asyncTraceContext, binder, traceRoot, asyncState);
+        return local.async(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