[pinpoint-apm#9575] Refactor AsyncTrace

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

@@ -18,6 +18,7 @@
 
 import com.navercorp.pinpoint.bootstrap.context.AsyncContext;
 import com.navercorp.pinpoint.bootstrap.context.AsyncState;
+import com.navercorp.pinpoint.profiler.context.id.LocalTraceRoot;
 import com.navercorp.pinpoint.profiler.context.id.TraceRoot;
 
 /**
@@ -31,4 +32,5 @@ public interface AsyncContextFactory {
 
     AsyncContext newAsyncContext(TraceRoot traceRoot, AsyncId asyncId, boolean canSampled, AsyncState asyncState);
 
+    AsyncContext newDisableAsyncContext(LocalTraceRoot traceRoot);
 }

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

@@ -17,6 +17,7 @@
 package com.navercorp.pinpoint.profiler.context;
 
 import com.navercorp.pinpoint.bootstrap.context.Trace;
+import com.navercorp.pinpoint.profiler.context.id.LocalTraceRoot;
 import com.navercorp.pinpoint.profiler.context.id.TraceRoot;
 
 
@@ -25,17 +26,8 @@
  */
 public interface AsyncTraceContext {
 
-//    Reference<Trace> continueAsyncTraceObject(TraceRoot traceRoot, int asyncId, short asyncSequence);
-
-    Reference<Trace> continueAsyncContextTraceObject(TraceRoot traceRoot, LocalAsyncId localAsyncId, boolean canSampled);
-
-    Trace newAsyncContextTraceObject(TraceRoot traceRoot, LocalAsyncId localAsyncId, boolean canSampled);
-
-    Reference<Trace> currentRawTraceObject();
-
-    Reference<Trace> currentTraceObject();
-
-    void removeTraceObject();
+    Trace continueAsyncContextTraceObject(TraceRoot traceRoot, LocalAsyncId localAsyncId);
 
+    Trace continueDisableAsyncContextTraceObject(LocalTraceRoot traceRoot);
 
 }

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

@@ -19,6 +19,7 @@
 import com.navercorp.pinpoint.bootstrap.context.Trace;
 import com.navercorp.pinpoint.bootstrap.context.TraceId;
 import com.navercorp.pinpoint.common.annotations.InterfaceAudience;
+import com.navercorp.pinpoint.profiler.context.id.LocalTraceRoot;
 import com.navercorp.pinpoint.profiler.context.id.TraceRoot;
 
 /**
@@ -34,7 +35,9 @@ public interface BaseTraceFactory {
     @InterfaceAudience.LimitedPrivate("vert.x")
     Trace continueAsyncTraceObject(TraceId traceId);
 
-    Trace continueAsyncContextTraceObject(TraceRoot traceRoot, LocalAsyncId localAsyncId, boolean canSampled);
+    Trace continueAsyncContextTraceObject(TraceRoot traceRoot, LocalAsyncId localAsyncId);
+
+    Trace continueDisableAsyncContextTraceObject(LocalTraceRoot traceRoot);
 
     Trace newTraceObject();
 

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

@@ -23,8 +23,8 @@
 import com.navercorp.pinpoint.common.trace.ServiceType;
 import com.navercorp.pinpoint.common.util.Assert;
 import com.navercorp.pinpoint.profiler.context.id.TraceRoot;
-import org.apache.logging.log4j.Logger;
 import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 
 import java.util.Objects;
 
@@ -38,21 +38,23 @@ public class DefaultAsyncContext implements AsyncContext {
 
     private final TraceRoot traceRoot;
     private final AsyncId asyncId;
-    private final boolean canSampled;
 
     private final AsyncTraceContext asyncTraceContext;
+    private final Binder<Trace> binder;
 
     private final int asyncMethodApiId;
 
 
-    public DefaultAsyncContext(AsyncTraceContext asyncTraceContext, TraceRoot traceRoot, AsyncId asyncId, int asyncMethodApiId, boolean canSampled) {
+    public DefaultAsyncContext(AsyncTraceContext asyncTraceContext,
+                               Binder<Trace> binder,
+                               TraceRoot traceRoot,
+                               AsyncId asyncId, int asyncMethodApiId) {
         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.canSampled = canSampled;
     }
 
 
@@ -63,7 +65,7 @@ public TraceRoot getTraceRoot() {
     @Override
     public Trace continueAsyncTraceObject() {
 
-        final Reference<Trace> reference = asyncTraceContext.currentRawTraceObject();
+        final Reference<Trace> reference = binder.get();
         final Trace nestedTrace = reference.get();
         if (nestedTrace != null) {
             // return Nested Trace Object?
@@ -80,8 +82,7 @@ private Trace newAsyncContextTrace(Reference<Trace> reference) {
 //        final int asyncId = this.asyncId.getAsyncId();
 //        final short asyncSequence = this.asyncId.nextAsyncSequence();
         final LocalAsyncId localAsyncId = this.asyncId.nextLocalAsyncId();
-        final Trace asyncTrace = asyncTraceContext.newAsyncContextTraceObject(traceRoot, localAsyncId, canSampled);
-
+        final Trace asyncTrace = asyncTraceContext.continueAsyncContextTraceObject(traceRoot, localAsyncId);
 
         bind(reference, asyncTrace);
 
@@ -127,14 +128,27 @@ private void bind(Reference<Trace> reference, Trace asyncTrace) {
 
     @Override
     public Trace currentAsyncTraceObject() {
-        final Reference<Trace> reference = asyncTraceContext.currentTraceObject();
+        final Reference<Trace> reference = currentTraceObject();
         return reference.get();
     }
 
 
+    private Reference<Trace> currentTraceObject() {
+        final Reference<Trace> reference = binder.get();
+        final Trace trace = reference.get();
+        if (trace == null) {
+            return DefaultReference.emptyReference();
+        }
+        if (trace.canSampled()) {
+            return reference;
+        }
+        return DefaultReference.emptyReference();
+    }
+
+
     @Override
     public void close() {
-        asyncTraceContext.removeTraceObject();
+        binder.remove();
     }
 
     @Override

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

@@ -19,7 +19,9 @@
 import com.navercorp.pinpoint.bootstrap.context.AsyncContext;
 import com.navercorp.pinpoint.bootstrap.context.AsyncState;
 import com.navercorp.pinpoint.bootstrap.context.MethodDescriptor;
+import com.navercorp.pinpoint.bootstrap.context.Trace;
 import com.navercorp.pinpoint.profiler.context.id.AsyncIdGenerator;
+import com.navercorp.pinpoint.profiler.context.id.LocalTraceRoot;
 import com.navercorp.pinpoint.profiler.context.id.TraceRoot;
 import com.navercorp.pinpoint.profiler.context.method.PredefinedMethodDescriptorRegistry;
 
@@ -31,13 +33,18 @@
 public class DefaultAsyncContextFactory implements AsyncContextFactory {
 
     private final AsyncTraceContext asyncTraceContext;
+    private final Binder<Trace> binder;
     private final AsyncIdGenerator asyncIdGenerator;
     private final PredefinedMethodDescriptorRegistry predefinedMethodDescriptorRegistry;
     private final int asyncMethodApiId;
 
-    public DefaultAsyncContextFactory(AsyncTraceContext asyncTraceContext, AsyncIdGenerator asyncIdGenerator, PredefinedMethodDescriptorRegistry predefinedMethodDescriptorRegistry) {
+    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");
 
         this.predefinedMethodDescriptorRegistry = Objects.requireNonNull(predefinedMethodDescriptorRegistry, "predefinedMethodDescriptorRegistry");
 
@@ -59,7 +66,11 @@ public AsyncContext newAsyncContext(TraceRoot traceRoot, AsyncId asyncId, boolea
         Objects.requireNonNull(traceRoot, "traceRoot");
         Objects.requireNonNull(asyncId, "asyncId");
 
-        return new DefaultAsyncContext(asyncTraceContext, traceRoot, asyncId, this.asyncMethodApiId, canSampled);
+        if (canSampled) {
+            return new DefaultAsyncContext(asyncTraceContext, binder, traceRoot, asyncId, this.asyncMethodApiId);
+        } else {
+            return newDisableAsyncContext(traceRoot);
+        }
     }
 
     @Override
@@ -68,8 +79,18 @@ public AsyncContext newAsyncContext(TraceRoot traceRoot, AsyncId asyncId, boolea
         Objects.requireNonNull(asyncId, "asyncId");
         Objects.requireNonNull(asyncState, "asyncState");
 
-        return new StatefulAsyncContext(asyncTraceContext, traceRoot, asyncId, asyncMethodApiId, asyncState, canSampled);
+        if (canSampled) {
+            return new StatefulAsyncContext(asyncTraceContext, binder, traceRoot, asyncId, asyncMethodApiId, asyncState);
+        } else {
+            // TODO
+            return new StatefulDisableAsyncContext(asyncTraceContext, binder, traceRoot, asyncState);
+        }
+
     }
 
+    @Override
+    public AsyncContext newDisableAsyncContext(LocalTraceRoot traceRoot) {
+        return new DisableAsyncContext(asyncTraceContext, binder, traceRoot);
+    }
 
 }

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

@@ -18,87 +18,34 @@
 
 import com.google.inject.Provider;
 import com.navercorp.pinpoint.bootstrap.context.Trace;
-import java.util.Objects;
-import com.navercorp.pinpoint.exception.PinpointException;
+import com.navercorp.pinpoint.profiler.context.id.LocalTraceRoot;
 import com.navercorp.pinpoint.profiler.context.id.TraceRoot;
-import org.apache.logging.log4j.Logger;
-import org.apache.logging.log4j.LogManager;
+
+import java.util.Objects;
 
 
 /**
  * @author Woonduk Kang(emeroad)
  */
 public class DefaultAsyncTraceContext implements AsyncTraceContext {
 
-    private final Logger logger = LogManager.getLogger(this.getClass());
-
-    private static final Reference<Trace> EMPTY = DefaultReference.emptyReference();
-
     private final Provider<BaseTraceFactory> baseTraceFactoryProvider;
-    private final Binder<Trace> binder;
 
-    public DefaultAsyncTraceContext(Provider<BaseTraceFactory> baseTraceFactoryProvider, Binder<Trace> binder) {
+    public DefaultAsyncTraceContext(Provider<BaseTraceFactory> baseTraceFactoryProvider) {
         this.baseTraceFactoryProvider = Objects.requireNonNull(baseTraceFactoryProvider, "baseTraceFactoryProvider");
-        this.binder = Objects.requireNonNull(binder, "binder");
-    }
-
-    @Override
-    public Reference<Trace> continueAsyncContextTraceObject(TraceRoot traceRoot, LocalAsyncId localAsyncId, boolean canSampled) {
-        final Reference<Trace> reference = checkAndGet();
-
-        final Trace trace = newAsyncContextTraceObject(traceRoot, localAsyncId, canSampled);
-
-        bind(reference, trace);
-        return reference;
     }
 
     @Override
-    public Trace newAsyncContextTraceObject(TraceRoot traceRoot, LocalAsyncId localAsyncId, boolean canSampled) {
+    public Trace continueAsyncContextTraceObject(TraceRoot traceRoot, LocalAsyncId localAsyncId) {
         final BaseTraceFactory baseTraceFactory = baseTraceFactoryProvider.get();
-        return baseTraceFactory.continueAsyncContextTraceObject(traceRoot, localAsyncId, canSampled);
+        return baseTraceFactory.continueAsyncContextTraceObject(traceRoot, localAsyncId);
     }
 
-
     @Override
-    public Reference<Trace> currentRawTraceObject() {
-        return binder.get();
-    }
-
-    @Override
-    public Reference<Trace> currentTraceObject() {
-        final Reference<Trace> reference = binder.get();
-        final Trace trace = reference.get();
-        if (trace == null) {
-            return EMPTY;
-        }
-        if (trace.canSampled()) {
-            return reference;
-        }
-        return EMPTY;
-    }
-
-
-    @Override
-    public void removeTraceObject() {
-        binder.remove();
-    }
-
-
-    private Reference<Trace> checkAndGet() {
-        final Reference<Trace> reference = this.binder.get();
-        final Trace old = reference.get();
-        if (old != null) {
-            final PinpointException exception = new PinpointException("already Trace Object exist.");
-            if (logger.isWarnEnabled()) {
-                logger.warn("beforeTrace:{}", old, exception);
-            }
-            throw exception;
-        }
-        return reference;
+    public Trace continueDisableAsyncContextTraceObject(LocalTraceRoot traceRoot) {
+        final BaseTraceFactory baseTraceFactory = baseTraceFactoryProvider.get();
+        return baseTraceFactory.continueDisableAsyncContextTraceObject(traceRoot);
     }
 
-    private void bind(Reference<Trace> reference, Trace trace) {
-        reference.set(trace);
-    }
 
 }

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

@@ -143,23 +143,23 @@ Trace newTraceObject(TraceSampler.State state) {
 
     // internal async trace.
     @Override
-    public Trace continueAsyncContextTraceObject(TraceRoot traceRoot, LocalAsyncId localAsyncId, boolean sampling) {
-        if (sampling) {
-            final SpanChunkFactory spanChunkFactory = new AsyncSpanChunkFactory(traceRoot, localAsyncId);
-            final Storage storage = storageFactory.createStorage(spanChunkFactory);
+    public Trace continueAsyncContextTraceObject(TraceRoot traceRoot, LocalAsyncId localAsyncId) {
+        final SpanChunkFactory spanChunkFactory = new AsyncSpanChunkFactory(traceRoot, localAsyncId);
+        final Storage storage = storageFactory.createStorage(spanChunkFactory);
 
-            final CallStack<SpanEvent> callStack = callStackFactory.newCallStack();
+        final CallStack<SpanEvent> callStack = callStackFactory.newCallStack();
 
-            final SpanRecorder spanRecorder = recorderFactory.newTraceRootSpanRecorder(traceRoot);
+        final SpanRecorder spanRecorder = recorderFactory.newTraceRootSpanRecorder(traceRoot);
 
-            final WrappedSpanEventRecorder wrappedSpanEventRecorder = recorderFactory.newWrappedSpanEventRecorder(traceRoot);
+        final WrappedSpanEventRecorder wrappedSpanEventRecorder = recorderFactory.newWrappedSpanEventRecorder(traceRoot);
 
-            return new AsyncChildTrace(traceRoot, callStack, storage, spanRecorder, wrappedSpanEventRecorder, localAsyncId);
-        } else {
-            return new DisableAsyncChildTrace(traceRoot, localAsyncId);
-        }
+        return new AsyncChildTrace(traceRoot, callStack, storage, spanRecorder, wrappedSpanEventRecorder, localAsyncId);
     }
 
+    @Override
+    public Trace continueDisableAsyncContextTraceObject(LocalTraceRoot traceRoot) {
+        return new DisableAsyncChildTrace(traceRoot);
+    }
 
     // entry point async trace.
     @InterfaceAudience.LimitedPrivate("vert.x")

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

@@ -34,11 +34,9 @@ public class DisableAsyncChildTrace implements Trace {
     private DefaultTraceScopePool scopePool;
 
     private final LocalTraceRoot traceRoot;
-    private final LocalAsyncId localAsyncId;
 
-    public DisableAsyncChildTrace(final LocalTraceRoot traceRoot, final LocalAsyncId localAsyncId) {
+    public DisableAsyncChildTrace(final LocalTraceRoot traceRoot) {
         this.traceRoot = Objects.requireNonNull(traceRoot, "traceRoot");
-        this.localAsyncId = Objects.requireNonNull(localAsyncId, "localAsyncId");
     }
 
     @Override
@@ -149,7 +147,6 @@ public TraceScope addScope(String name) {
     public String toString() {
         return "DisableAsyncChildTrace{" +
             "traceRoot=" + getTraceRoot() +
-            ", localAsyncId=" + localAsyncId +
             '}';
     }
 }