[#9595] Refactor AsyncState

Author: emeroad

bootstraps/bootstrap-core/src/main/java/com/navercorp/pinpoint/bootstrap/context/AsyncContextUtils.java

@@ -7,8 +7,7 @@ private AsyncContextUtils() {
     public static boolean asyncStateFinish(final AsyncContext asyncContext) {
         if (asyncContext instanceof AsyncStateSupport) {
             final AsyncStateSupport asyncStateSupport = (AsyncStateSupport) asyncContext;
-            AsyncState asyncState = asyncStateSupport.getAsyncState();
-            asyncState.finish();
+            asyncStateSupport.finish();
             return true;
         }
         return false;

bootstraps/bootstrap-core/src/main/java/com/navercorp/pinpoint/bootstrap/context/AsyncStateSupport.java

@@ -23,5 +23,9 @@
  */
 @InterfaceAudience.LimitedPrivate("vert.x")
 public interface AsyncStateSupport {
+    @Deprecated
     AsyncState getAsyncState();
+
+    void finish();
+
 }

profiler/src/main/java/com/navercorp/pinpoint/profiler/context/StatefulAsyncContext.java

@@ -44,6 +44,11 @@ public AsyncState getAsyncState() {
         return asyncState;
     }
 
+    @Override
+    public void finish() {
+        this.asyncState.finish();
+    }
+
     @Override
     public String toString() {
         return "StatefulAsyncContext{" +

profiler/src/main/java/com/navercorp/pinpoint/profiler/context/StatefulDisableAsyncContext.java

@@ -22,6 +22,11 @@ public AsyncState getAsyncState() {
         return asyncState;
     }
 
+    @Override
+    public void finish() {
+        this.asyncState.finish();
+    }
+
     @Override
     public String toString() {
         return "StatefulDisableAsyncContext{" +

profiler/src/main/java/com/navercorp/pinpoint/profiler/context/recorder/DefaultRecorderFactory.java

@@ -51,30 +51,39 @@ public DefaultRecorderFactory(Provider<AsyncContextFactory> asyncContextFactoryP
 
     @Override
     public SpanRecorder newSpanRecorder(Span span) {
+        Objects.requireNonNull(span, "span");
+
         return new DefaultSpanRecorder(span, stringMetaDataService, sqlMetaDataService, errorHandler);
     }
 
     @Override
     public SpanRecorder newTraceRootSpanRecorder(TraceRoot traceRoot) {
+        Objects.requireNonNull(traceRoot, "traceRoot");
+
         return new TraceRootSpanRecorder(traceRoot);
     }
 
     @Override
     public SpanRecorder newDisableSpanRecorder(LocalTraceRoot traceRoot) {
+        Objects.requireNonNull(traceRoot, "traceRoot");
+
         return new DisableSpanRecorder(traceRoot, errorHandler);
     }
 
     @Override
     public WrappedSpanEventRecorder newWrappedSpanEventRecorder(TraceRoot traceRoot) {
+        Objects.requireNonNull(traceRoot, "traceRoot");
+
         final AsyncContextFactory asyncContextFactory = asyncContextFactoryProvider.get();
         return new WrappedSpanEventRecorder(traceRoot, asyncContextFactory, stringMetaDataService, sqlMetaDataService, errorHandler);
     }
 
     @Override
     public WrappedSpanEventRecorder newWrappedSpanEventRecorder(TraceRoot traceRoot, AsyncState asyncState) {
+        Objects.requireNonNull(traceRoot, "traceRoot");
         Objects.requireNonNull(asyncState, "asyncState");
 
         final AsyncContextFactory asyncContextFactory = asyncContextFactoryProvider.get();
-        return new WrappedAsyncSpanEventRecorder(traceRoot, asyncContextFactory, stringMetaDataService, sqlMetaDataService, errorHandler, asyncState);
+        return new WrappedSpanEventRecorder(traceRoot, asyncContextFactory, asyncState, stringMetaDataService, sqlMetaDataService, errorHandler);
     }
 }

profiler/src/main/java/com/navercorp/pinpoint/profiler/context/recorder/WrappedAsyncSpanEventRecorder.java

@@ -17,6 +17,7 @@
 
 
 import com.navercorp.pinpoint.bootstrap.context.AsyncContext;
+import com.navercorp.pinpoint.bootstrap.context.AsyncState;
 import com.navercorp.pinpoint.bootstrap.context.ParsingResult;
 import com.navercorp.pinpoint.bootstrap.context.SpanEventRecorder;
 import com.navercorp.pinpoint.common.trace.AnnotationKey;
@@ -38,6 +39,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
+import javax.annotation.Nullable;
 import java.util.Objects;
 
 /**
@@ -49,18 +51,34 @@ public class WrappedSpanEventRecorder extends AbstractRecorder implements SpanEv
     private static final Logger logger = LogManager.getLogger(DefaultTrace.class.getName());
     private static final boolean isDebug = logger.isDebugEnabled();
 
-    protected final TraceRoot traceRoot;
-    protected final AsyncContextFactory asyncContextFactory;
+    private final TraceRoot traceRoot;
+    private final AsyncContextFactory asyncContextFactory;
+    @Nullable
+    private final AsyncState asyncState;
 
     private SpanEvent spanEvent;
 
-    public WrappedSpanEventRecorder(TraceRoot traceRoot, AsyncContextFactory asyncContextFactory,
-                                    final StringMetaDataService stringMetaDataService, final SqlMetaDataService sqlMetaCacheService,
+    public WrappedSpanEventRecorder(TraceRoot traceRoot,
+                                    AsyncContextFactory asyncContextFactory,
+                                    StringMetaDataService stringMetaDataService,
+                                    SqlMetaDataService sqlMetaDataService,
+                                    IgnoreErrorHandler ignoreErrorHandler) {
+
+        this(traceRoot, asyncContextFactory, null, stringMetaDataService, sqlMetaDataService, ignoreErrorHandler);
+    }
+
+    public WrappedSpanEventRecorder(TraceRoot traceRoot,
+                                    AsyncContextFactory asyncContextFactory,
+                                    @Nullable
+                                    final AsyncState asyncState,
+                                    final StringMetaDataService stringMetaDataService,
+                                    final SqlMetaDataService sqlMetaCacheService,
                                     final IgnoreErrorHandler errorHandler) {
         super(stringMetaDataService, sqlMetaCacheService, errorHandler);
         this.traceRoot = Objects.requireNonNull(traceRoot, "traceRoot");
 
         this.asyncContextFactory = Objects.requireNonNull(asyncContextFactory, "asyncContextFactory");
+        this.asyncState = asyncState;
     }
 
     public void setWrapped(final SpanEvent spanEvent) {
@@ -137,6 +155,15 @@ protected boolean isOverflowState() {
 
     @Override
     public AsyncContext recordNextAsyncContext(boolean asyncStateSupport) {
+        if (asyncStateSupport && asyncState != null) {
+            final TraceRoot traceRoot = this.traceRoot;
+            final AsyncId asyncIdObject = getNextAsyncId();
+            final boolean isDisabled = isOverflowState();
+
+            final AsyncState asyncState = this.asyncState;
+            asyncState.setup();
+            return asyncContextFactory.newAsyncContext(traceRoot, asyncIdObject, isDisabled, asyncState);
+        }
         return recordNextAsyncContext();
     }