Support keep trace profiling when cross-thread (#479)

Author: mrproliu

CHANGES.md

@@ -20,6 +20,7 @@ Release Notes.
 * Adapt Armeria's plugins to the latest version 1.22.x
 * Fix tomcat-10x-plugin and add test case to support tomcat7.x-8.x-9.x.
 * Fix thrift plugin generate duplicate traceid when `sendBase` error occurs
+* Support keep trace profiling when cross-thread.
 
 #### Documentation
 * Update docs of Tracing APIs, reorganize the API docs into six parts.

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/conf/Config.java

@@ -241,10 +241,15 @@ public static class Profile {
         public static boolean ACTIVE = true;
 
         /**
-         * Parallel monitor segment count
+         * Parallel monitor endpoint thread count
          */
         public static int MAX_PARALLEL = 5;
 
+        /**
+         * Max monitoring sub-tasks count of one single endpoint access
+         */
+        public static int MAX_ACCEPT_SUB_PARALLEL = 5;
+
         /**
          * Max monitor segment time(minutes), if current segment monitor time out of limit, then stop it.
          */

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/ContextSnapshot.java

@@ -20,6 +20,7 @@
 
 import lombok.Getter;
 import org.apache.skywalking.apm.agent.core.context.ids.DistributedTraceId;
+import org.apache.skywalking.apm.agent.core.profile.ProfileStatusContext;
 
 /**
  * The <code>ContextSnapshot</code> is a snapshot for current context. The snapshot carries the info for building
@@ -34,19 +35,22 @@ public class ContextSnapshot {
 
     private CorrelationContext correlationContext;
     private ExtensionContext extensionContext;
+    private ProfileStatusContext profileStatusContext;
 
     ContextSnapshot(String traceSegmentId,
                     int spanId,
                     DistributedTraceId primaryTraceId,
                     String parentEndpoint,
                     CorrelationContext correlationContext,
-                    ExtensionContext extensionContext) {
+                    ExtensionContext extensionContext,
+                    ProfileStatusContext profileStatusContext) {
         this.traceSegmentId = traceSegmentId;
         this.spanId = spanId;
         this.traceId = primaryTraceId;
         this.parentEndpoint = parentEndpoint;
         this.correlationContext = correlationContext.clone();
         this.extensionContext = extensionContext.clone();
+        this.profileStatusContext = profileStatusContext.clone();
     }
 
     public boolean isFromCurrent() {

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/IgnoredTracerContext.java

@@ -22,6 +22,7 @@
 import java.util.List;
 import org.apache.skywalking.apm.agent.core.context.trace.AbstractSpan;
 import org.apache.skywalking.apm.agent.core.context.trace.NoopSpan;
+import org.apache.skywalking.apm.agent.core.profile.ProfileStatusContext;
 
 /**
  * The <code>IgnoredTracerContext</code> represent a context should be ignored. So it just maintains the stack with an
@@ -35,13 +36,15 @@ public class IgnoredTracerContext implements AbstractTracerContext {
 
     private final CorrelationContext correlationContext;
     private final ExtensionContext extensionContext;
+    private final ProfileStatusContext profileStatusContext;
 
     private int stackDepth;
 
     public IgnoredTracerContext() {
         this.stackDepth = 0;
         this.correlationContext = new CorrelationContext();
         this.extensionContext = new ExtensionContext();
+        this.profileStatusContext = ProfileStatusContext.createWithNone();
     }
 
     @Override
@@ -56,7 +59,7 @@ public void extract(ContextCarrier carrier) {
 
     @Override
     public ContextSnapshot capture() {
-        return new ContextSnapshot(null, -1, null, null, correlationContext, extensionContext);
+        return new ContextSnapshot(null, -1, null, null, correlationContext, extensionContext, profileStatusContext);
     }
 
     @Override

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/TracingContext.java

@@ -41,7 +41,7 @@
 import org.apache.skywalking.apm.agent.core.context.trace.TraceSegmentRef;
 import org.apache.skywalking.apm.agent.core.logging.api.ILog;
 import org.apache.skywalking.apm.agent.core.logging.api.LogManager;
-import org.apache.skywalking.apm.agent.core.profile.ProfileStatusReference;
+import org.apache.skywalking.apm.agent.core.profile.ProfileStatusContext;
 import org.apache.skywalking.apm.agent.core.profile.ProfileTaskExecutionService;
 import org.apache.skywalking.apm.util.StringUtil;
 
@@ -106,7 +106,7 @@ public class TracingContext implements AbstractTracerContext {
     /**
      * profile status
      */
-    private final ProfileStatusReference profileStatus;
+    private final ProfileStatusContext profileStatus;
     @Getter(AccessLevel.PACKAGE)
     private final CorrelationContext correlationContext;
     @Getter(AccessLevel.PACKAGE)
@@ -213,7 +213,8 @@ public ContextSnapshot capture() {
             getPrimaryTraceId(),
             primaryEndpoint.getName(),
             this.correlationContext,
-            this.extensionContext
+            this.extensionContext,
+            this.profileStatus
         );
 
         return snapshot;
@@ -234,6 +235,9 @@ public void continued(ContextSnapshot snapshot) {
             this.correlationContext.continued(snapshot);
             this.extensionContext.continued(snapshot);
             this.extensionContext.handle(this.activeSpan());
+            if (this.profileStatus.continued(snapshot)) {
+                PROFILE_TASK_EXECUTION_SERVICE.continueProfiling(this, this.segment.getTraceSegmentId());
+            }
         }
     }
 
@@ -571,7 +575,7 @@ public long createTime() {
         return this.createTime;
     }
 
-    public ProfileStatusReference profileStatus() {
+    public ProfileStatusContext profileStatus() {
         return this.profileStatus;
     }
 

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/profile/ProfileStatusContext.java

@@ -0,0 +1,110 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package org.apache.skywalking.apm.agent.core.profile;
+
+import org.apache.skywalking.apm.agent.core.conf.Config;
+import org.apache.skywalking.apm.agent.core.context.ContextSnapshot;
+import org.apache.skywalking.apm.agent.core.context.TracingContext;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * Wrapper {@link ProfileStatus}, make sure {@link org.apache.skywalking.apm.agent.core.context.TracingContext} with {@link ThreadProfiler} have same reference with {@link ProfileStatus},
+ * And only the profile module could change the status
+ */
+public class ProfileStatusContext {
+
+    private volatile ProfileStatus status;
+    private volatile long firstSegmentCreateTime;
+    private volatile AtomicInteger subThreadProfilingCount;
+
+    private ProfileStatusContext(ProfileStatus status, long firstSegmentCreateTime, AtomicInteger subThreadProfilingCount) {
+        this.status = status;
+        this.firstSegmentCreateTime = firstSegmentCreateTime;
+        this.subThreadProfilingCount = subThreadProfilingCount;
+    }
+
+    /**
+     * Create with not watching
+     */
+    public static ProfileStatusContext createWithNone() {
+        return new ProfileStatusContext(ProfileStatus.NONE, 0, null);
+    }
+
+    /**
+     * Create with pending to profile
+     */
+    public static ProfileStatusContext createWithPending(long firstSegmentCreateTime) {
+        return new ProfileStatusContext(ProfileStatus.PENDING, firstSegmentCreateTime, new AtomicInteger(0));
+    }
+
+    public ProfileStatus get() {
+        return this.status;
+    }
+
+    public long firstSegmentCreateTime() {
+        return this.firstSegmentCreateTime;
+    }
+
+    /**
+     * The profile monitoring is watching, wait for some profile conditions.
+     */
+    public boolean isBeingWatched() {
+        return this.status != ProfileStatus.NONE;
+    }
+
+    public boolean isProfiling() {
+        return this.status == ProfileStatus.PROFILING;
+    }
+
+    public ProfileStatusContext clone() {
+        return new ProfileStatusContext(this.status, this.firstSegmentCreateTime, this.subThreadProfilingCount);
+    }
+
+    /**
+     * Continued profile status context
+     * @return is needs to keep profile
+     */
+    public boolean continued(ContextSnapshot snapshot) {
+        this.status = snapshot.getProfileStatusContext().get();
+        this.firstSegmentCreateTime = snapshot.getProfileStatusContext().firstSegmentCreateTime();
+        this.subThreadProfilingCount = snapshot.getProfileStatusContext().subThreadProfilingCount;
+        return this.isBeingWatched() &&
+            // validate is reach the count of sub-thread
+            this.subThreadProfilingCount != null &&
+            this.subThreadProfilingCount.incrementAndGet() <= Config.Profile.MAX_ACCEPT_SUB_PARALLEL;
+    }
+
+    /**
+     * Update status, only access with profile module
+     */
+    void updateStatus(ProfileStatus status, TracingContext tracingContext) {
+        this.status = status;
+        if (this.firstSegmentCreateTime == 0 && tracingContext != null) {
+            this.firstSegmentCreateTime = tracingContext.createTime();
+        }
+    }
+
+    void updateStatus(ProfileStatusContext statusContext) {
+        this.status = statusContext.get();
+        this.firstSegmentCreateTime = statusContext.firstSegmentCreateTime();
+        this.subThreadProfilingCount = statusContext.subThreadProfilingCount;
+    }
+
+}