From db7a14d167c76fb04323b49c7b31cb6d9cb46150 Mon Sep 17 00:00:00 2001
From: mrproliu <741550557@qq.com>
Date: Wed, 29 Mar 2023 15:31:28 +0800
Subject: [PATCH] Fix the profiling segment count not correct (#491)

---
 .../apm/agent/core/profile/ProfileStatusContext.java       | 7 +++++++
 .../agent/core/profile/ProfileTaskExecutionContext.java    | 6 ++++--
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/profile/ProfileStatusContext.java b/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/profile/ProfileStatusContext.java
index 671216d4d..d480f28a8 100644
--- a/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/profile/ProfileStatusContext.java
+++ b/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/profile/ProfileStatusContext.java
@@ -31,11 +31,13 @@ import java.util.concurrent.atomic.AtomicInteger;
 public class ProfileStatusContext {
 
     private volatile ProfileStatus status;
+    private volatile boolean fromFirstSegment;
     private volatile long firstSegmentCreateTime;
     private volatile AtomicInteger subThreadProfilingCount;
 
     private ProfileStatusContext(ProfileStatus status, long firstSegmentCreateTime, AtomicInteger subThreadProfilingCount) {
         this.status = status;
+        this.fromFirstSegment = true;
         this.firstSegmentCreateTime = firstSegmentCreateTime;
         this.subThreadProfilingCount = subThreadProfilingCount;
     }
@@ -62,6 +64,10 @@ public class ProfileStatusContext {
         return this.firstSegmentCreateTime;
     }
 
+    public boolean isFromFirstSegment() {
+        return fromFirstSegment;
+    }
+
     /**
      * The profile monitoring is watching, wait for some profile conditions.
      */
@@ -83,6 +89,7 @@ public class ProfileStatusContext {
      */
     public boolean continued(ContextSnapshot snapshot) {
         this.status = snapshot.getProfileStatusContext().get();
+        this.fromFirstSegment = false;
         this.firstSegmentCreateTime = snapshot.getProfileStatusContext().firstSegmentCreateTime();
         this.subThreadProfilingCount = snapshot.getProfileStatusContext().subThreadProfilingCount;
         return this.isBeingWatched() &&
diff --git a/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/profile/ProfileTaskExecutionContext.java b/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/profile/ProfileTaskExecutionContext.java
index e53090b5f..929fab0d6 100644
--- a/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/profile/ProfileTaskExecutionContext.java
+++ b/apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/profile/ProfileTaskExecutionContext.java
@@ -131,7 +131,7 @@ public class ProfileTaskExecutionContext {
 
         // update profiling status
         tracingContext.profileStatus()
-                      .updateStatus(attemptProfiling(tracingContext, traceSegmentId, firstSpanOPName));
+            .updateStatus(attemptProfiling(tracingContext, traceSegmentId, firstSpanOPName));
     }
 
     /**
@@ -147,7 +147,9 @@ public class ProfileTaskExecutionContext {
 
                 // setting stop running
                 currentProfiler.stopProfiling();
-                currentEndpointProfilingCount.addAndGet(-1);
+                if (currentProfiler.profilingStatus().isFromFirstSegment()) {
+                    currentEndpointProfilingCount.addAndGet(-1);
+                }
                 break;
             }
         }