fix: fix emission problems

Author: kaahos

ddprof-lib/src/main/cpp/profiler.cpp

@@ -1437,11 +1437,11 @@ Error Profiler::init() {
 
 Error Profiler::start(Arguments &args, bool reset) {
   MutexLocker ml(_state_lock);
-  _task_block_enabled.store(false, std::memory_order_release);
   Error error = checkState();
   if (error) {
     return error;
   }
+  _task_block_enabled.store(false, std::memory_order_release);
 
   error = checkJvmCapabilities();
   if (error) {

ddprof-lib/src/main/cpp/threadFilter.cpp

@@ -673,8 +673,20 @@ bool ThreadFilter::snapshotAndExitBlockedRun(SlotID slot_id, u64 generation,
     return true;
 }
 
+bool ThreadFilter::activeOwnedBlockGeneration(const ThreadEntry& entry,
+                                              u64& generation) const {
+    return ownedBlockGeneration(entry, generation, false);
+}
+
 bool ThreadFilter::isOwnedBlockSuppressionCandidate(
     const ThreadEntry& entry) const {
+    u64 generation = 0;
+    return ownedBlockGeneration(entry, generation, true);
+}
+
+bool ThreadFilter::ownedBlockGeneration(const ThreadEntry& entry,
+                                        u64& generation,
+                                        bool require_sampled) const {
     Slot* slot = entry.slot;
     if (!unfilteredWallTrackingActive() || slot == nullptr ||
         slot->nativeTid() != entry.tid ||
@@ -701,7 +713,11 @@ bool ThreadFilter::isOwnedBlockSuppressionCandidate(
 
     u64 block_generation = slot->blockGeneration();
     BlockRunOwner owner = slot->activeBlockOwner();
-    if (owner == BlockRunOwner::NONE) return false;
+    if (owner == BlockRunOwner::NONE ||
+        (require_sampled &&
+         slot->sampledBlockGeneration() != block_generation)) {
+        return false;
+    }
 
 #ifdef UNIT_TEST
     if (_suppression_snapshot_hook != nullptr) {
@@ -714,13 +730,16 @@ bool ThreadFilter::isOwnedBlockSuppressionCandidate(
     if (slot->activeBlockOwner() != owner ||
         slot->blockGeneration() != block_generation ||
         slot->activeBlockState() != state || slot->nativeTid() != entry.tid ||
-        slot->lifecycleGeneration() != entry.lifecycle_generation) {
+        slot->lifecycleGeneration() != entry.lifecycle_generation ||
+        (require_sampled &&
+         slot->sampledBlockGeneration() != block_generation)) {
         return false;
     }
     if (recordingEpoch() != epoch || slot->recordingEpoch() != epoch ||
         !slot->activeBlockRemainedOutsideContextWindow()) {
         return false;
     }
+    generation = block_generation;
     return true;
 }
 

ddprof-lib/src/main/cpp/threadFilter.h

@@ -85,6 +85,10 @@ class ThreadFilter {
         std::atomic<u64>           recording_epoch{0};
         std::atomic<u64>           active_block_context_epoch{0};
         std::atomic<u64>           block_generation{0};
+        // The most recent owned-block generation for which a MethodSample was
+        // recorded successfully. Generations are monotonic, so a delayed
+        // completion from an older run cannot mark a newer run as sampled.
+        std::atomic<u64>           sampled_block_generation{0};
         std::atomic<OSThreadState> unowned_blocked_state{OSThreadState::UNKNOWN};
         // Native identity and context-window membership are independent so an
         // unfiltered wall recording can retain lifecycle metadata without
@@ -103,6 +107,7 @@ class ThreadFilter {
                      - sizeof(std::atomic<u64>)
                      - sizeof(std::atomic<u64>)
                      - sizeof(std::atomic<u64>)
+                     - sizeof(std::atomic<u64>)
                      - sizeof(std::atomic<OSThreadState>)
                      - sizeof(std::atomic<int>)
                      - sizeof(std::atomic<int>)
@@ -160,6 +165,20 @@ class ThreadFilter {
         inline u64 blockGeneration() const {
             return block_generation.load(std::memory_order_acquire);
         }
+        inline u64 sampledBlockGeneration() const {
+            return sampled_block_generation.load(std::memory_order_acquire);
+        }
+        inline void markBlockGenerationSampled(u64 generation) {
+            u64 sampled = sampled_block_generation.load(std::memory_order_relaxed);
+            while (sampled < generation &&
+                   !sampled_block_generation.compare_exchange_weak(
+                       sampled, generation, std::memory_order_release,
+                       std::memory_order_relaxed)) {
+            }
+        }
+        inline void resetSampledBlockGeneration() {
+            sampled_block_generation.store(0, std::memory_order_relaxed);
+        }
         inline bool unownedBlockedFallbackEnabled() const {
             return unowned_blocked_fallback_enabled.load(std::memory_order_acquire) != 0;
         }
@@ -236,6 +255,7 @@ class ThreadFilter {
             generation++;
             block_generation.store(generation, std::memory_order_relaxed);
             active_block_context_epoch.store(context_state >> 1, std::memory_order_relaxed);
+            resetSampledBlockGeneration();
             disableUnownedBlockedFallback();
             *generation_out = generation;
             return true;
@@ -245,6 +265,7 @@ class ThreadFilter {
         }
         inline void clearActiveBlockRun(OSThreadState) {
             active_block_state.store(OSThreadState::UNKNOWN, std::memory_order_release);
+            resetSampledBlockGeneration();
             resetUnownedBlockedSampling();
             active_block_owner.store(static_cast<int>(BlockRunOwner::NONE), std::memory_order_release);
         }
@@ -299,6 +320,8 @@ class ThreadFilter {
     bool exitBlockedRun(SlotID slot_id, u64 generation);
     bool snapshotAndExitBlockedRun(SlotID slot_id, u64 generation,
                                    BlockRunSnapshot* snapshot);
+    bool activeOwnedBlockGeneration(const ThreadEntry& entry,
+                                    u64& generation) const;
     bool isOwnedBlockSuppressionCandidate(const ThreadEntry& entry) const;
 
 #ifdef UNIT_TEST
@@ -356,6 +379,8 @@ class ThreadFilter {
     SlotID slotIdByTid(int tid) const { return lookupSlotIdByTid(tid); }
 
 private:
+    bool ownedBlockGeneration(const ThreadEntry& entry, u64& generation,
+                              bool require_sampled) const;
 
     // Lock-free free list using a stack-like structure
     struct FreeListNode {

ddprof-lib/src/main/cpp/wallClock.cpp

@@ -59,6 +59,8 @@ static inline bool hasKnownActiveTraceContext(ProfiledThread* thread) {
 
 struct WallPrecheckResult {
   bool suppress = false;
+  ThreadFilter::Slot* owned_block_slot = nullptr;
+  u64 owned_block_generation = 0;
   OSThreadState observed_state = OSThreadState::UNKNOWN;
   bool observed_state_valid = false;
   ThreadFilter::Slot* unowned_weight_slot = nullptr;
@@ -101,9 +103,9 @@ static inline WallPrecheckResult prepareWallPrecheck(ProfiledThread* current,
     return result;
   }
 
-  // TaskBlock replaces signals only for threads that unfiltered wall-clock
-  // profiling observes outside the tracing context window. Context-scoped
-  // profiling must continue sampling its selected threads normally.
+  // Owned blocks replace repeated signals only after their current generation
+  // has produced one MethodSample. Context-scoped profiling must continue
+  // sampling its selected threads normally.
   if (!registry->unfilteredWallTrackingActive() || slot->inContextWindow()) {
     return result;
   }
@@ -115,6 +117,14 @@ static inline WallPrecheckResult prepareWallPrecheck(ProfiledThread* current,
     result.suppress = true;
     return result;
   }
+  u64 block_generation = 0;
+  if (registry->activeOwnedBlockGeneration(entry, block_generation)) {
+    // Arm only after recordSample succeeds. A skipped JFR write must leave the
+    // run eligible so the next signal retries instead of losing its only stack.
+    result.owned_block_slot = slot;
+    result.owned_block_generation = block_generation;
+    return result;
+  }
   if (!slot->unownedBlockedFallbackEnabled()) {
     return result;
   }
@@ -139,6 +149,10 @@ static inline WallPrecheckResult prepareWallPrecheck(ProfiledThread* current,
 static inline void finishWallPrecheck(const WallPrecheckResult& precheck,
                                       bool recorded,
                                       u64 recorded_call_trace_id = 0) {
+  if (recorded && precheck.owned_block_slot != nullptr) {
+    precheck.owned_block_slot->markBlockGenerationSampled(
+        precheck.owned_block_generation);
+  }
   if (!recorded && precheck.unowned_weight_slot != nullptr) {
     precheck.unowned_weight_slot->restoreUnownedBlockedWeight(
         precheck.unowned_weight);
@@ -238,12 +252,10 @@ void WallClockASGCT::signalHandler(int signo, siginfo_t *siginfo, void *ucontext
     current->tickInitWindow();
     return;
   }
-  // Once-per-run filter (wallprecheck=true): for untraced threads, exact
-  // suppression is only valid while an explicit lifecycle hook owns the blocked
-  // interval. Raw OS thread state is only an observation; it cannot distinguish
-  // one long sleep from several short sleeps separated by runnable gaps between
-  // signals. Unowned blocked observations therefore use weighted fallback
-  // sampling instead of arming sampled_this_run.
+  // Once-per-run filter (wallprecheck=true): an explicitly owned block keeps
+  // its first successful MethodSample and suppresses subsequent signals.
+  // Unowned blocked observations use weighted fallback sampling because raw OS
+  // state cannot distinguish one long sleep from several shorter runs.
   WallPrecheckResult precheck = prepareWallPrecheck(current, _precheck);
   if (precheck.suppress) {
     return;
@@ -379,9 +391,8 @@ void WallClockASGCT::timerLoop() {
         registry_lookups++;
         thread_filter->lookupThreadEntry(entry, recording_epoch);
       }
-      // Timer-thread fast path (wallprecheck=true): skip the kernel IPI while
-      // an explicit lifecycle hook owns a suppressible blocked run. Raw OS
-      // thread state cannot prove run boundaries for the target thread.
+      // Timer-thread fast path (wallprecheck=true): skip the kernel IPI only
+      // after an explicitly owned run has recorded its first MethodSample.
       if (_precheck && suppressOwnedBlock(entry)) {
         return WallClockCandidateOutcome::PRECHECK_REJECTED;
       }

ddprof-lib/src/test/cpp/park_state_ut.cpp

@@ -304,6 +304,7 @@ TEST(WallClockOncePerRunFilterTest, FilterHelpersManageActiveBlockState) {
 
   filter.exitBlockedRun(slot_id);
   EXPECT_EQ(OSThreadState::UNKNOWN, slot->activeBlockState());
+  EXPECT_EQ(0ULL, slot->sampledBlockGeneration());
 }
 
 TEST(WallClockOncePerRunFilterTest, ResetClearsOwnedBlockOnSlotReuse) {
@@ -314,8 +315,11 @@ TEST(WallClockOncePerRunFilterTest, ResetClearsOwnedBlockOnSlotReuse) {
   ThreadFilter::Slot *slot = filter.slotForId(slot_id);
   ASSERT_NE(nullptr, slot);
   EXPECT_EQ(OSThreadState::CONDVAR_WAIT, slot->activeBlockState());
+  slot->markBlockGenerationSampled(slot->blockGeneration());
+  ASSERT_EQ(slot->blockGeneration(), slot->sampledBlockGeneration());
 
   filter.resetSlotRunState(slot_id);
 
   EXPECT_EQ(OSThreadState::UNKNOWN, slot->activeBlockState());
+  EXPECT_EQ(0ULL, slot->sampledBlockGeneration());
 }

ddprof-lib/src/test/cpp/threadFilter_ut.cpp

@@ -631,7 +631,7 @@ TEST_F(ThreadFilterTest, SnapshotCapturesOwnedLifecycle) {
     EXPECT_FALSE(slot->snapshotBlockRun().active);
 }
 
-TEST_F(ThreadFilterTest, OwnedBlockSuppressesBeforeAnyWallSample) {
+TEST_F(ThreadFilterTest, OwnedBlockSuppressesOnlyAfterSuccessfulWallSample) {
     filter->init(nullptr, true);
     int slot_id = filter->registerThread(1234);
     ASSERT_GE(slot_id, 0);
@@ -642,6 +642,12 @@ TEST_F(ThreadFilterTest, OwnedBlockSuppressesBeforeAnyWallSample) {
 
     ThreadEntry entry{1234, slot, slot->lifecycleGeneration(),
                       slot->recordingEpoch()};
+    u64 generation = 0;
+    EXPECT_TRUE(filter->activeOwnedBlockGeneration(entry, generation));
+    EXPECT_EQ(ThreadFilter::tokenGeneration(token), generation);
+    EXPECT_FALSE(filter->isOwnedBlockSuppressionCandidate(entry));
+
+    slot->markBlockGenerationSampled(generation);
     EXPECT_TRUE(filter->isOwnedBlockSuppressionCandidate(entry));
     EXPECT_FALSE(filter->isOwnedBlockSuppressionCandidate(
         {1235, slot, slot->lifecycleGeneration(), slot->recordingEpoch()}));
@@ -654,6 +660,38 @@ TEST_F(ThreadFilterTest, OwnedBlockSuppressesBeforeAnyWallSample) {
     EXPECT_FALSE(filter->isOwnedBlockSuppressionCandidate(entry));
 }
 
+TEST_F(ThreadFilterTest, StaleSampleCompletionCannotSuppressNewBlockGeneration) {
+    filter->init(nullptr, true);
+    int slot_id = filter->registerThread(1234);
+    ASSERT_GE(slot_id, 0);
+    ThreadFilter::Slot* slot = filter->slotForId(slot_id);
+    ASSERT_NE(nullptr, slot);
+
+    u64 first_token = filter->enterBlockedRun(slot_id, OSThreadState::SLEEPING);
+    ASSERT_NE(0ULL, first_token);
+    u64 first_generation = ThreadFilter::tokenGeneration(first_token);
+    ASSERT_TRUE(filter->exitBlockedRun(slot_id, first_generation));
+
+    u64 second_token =
+        filter->enterBlockedRun(slot_id, OSThreadState::CONDVAR_WAIT);
+    ASSERT_NE(0ULL, second_token);
+    u64 second_generation = ThreadFilter::tokenGeneration(second_token);
+    ASSERT_GT(second_generation, first_generation);
+
+    ThreadEntry entry{1234, slot, slot->lifecycleGeneration(),
+                      slot->recordingEpoch()};
+    slot->markBlockGenerationSampled(first_generation);
+    EXPECT_FALSE(filter->isOwnedBlockSuppressionCandidate(entry));
+
+    slot->markBlockGenerationSampled(second_generation);
+    EXPECT_TRUE(filter->isOwnedBlockSuppressionCandidate(entry));
+
+    // A delayed completion from the first run must not overwrite the newer mark.
+    slot->markBlockGenerationSampled(first_generation);
+    EXPECT_EQ(second_generation, slot->sampledBlockGeneration());
+    EXPECT_TRUE(filter->isOwnedBlockSuppressionCandidate(entry));
+}
+
 TEST_F(ThreadFilterTest, ContextScopeNeverSuppressesOwnedBlock) {
     filter->init("0", false);
     int slot_id = filter->registerThread(1234);
@@ -666,6 +704,7 @@ TEST_F(ThreadFilterTest, ContextScopeNeverSuppressesOwnedBlock) {
 
     ThreadEntry entry{1234, slot, slot->lifecycleGeneration(),
                       slot->recordingEpoch()};
+    slot->markBlockGenerationSampled(slot->blockGeneration());
     EXPECT_FALSE(filter->isOwnedBlockSuppressionCandidate(entry));
 }
 
@@ -679,6 +718,7 @@ TEST_F(ThreadFilterTest, ContextEpochDisablesOwnedBlockSuppression) {
         slot_id, OSThreadState::CONDVAR_WAIT));
     ThreadEntry entry{1234, slot, slot->lifecycleGeneration(),
                       slot->recordingEpoch()};
+    slot->markBlockGenerationSampled(slot->blockGeneration());
     ASSERT_TRUE(filter->isOwnedBlockSuppressionCandidate(entry));
 
     filter->add(1234, slot_id);
@@ -868,6 +908,7 @@ TEST_F(ThreadRegistryTest, UnfilteredSuppressionValidatesIdentityAndLifecycle) {
 
     u64 token = registry.enterBlockedRun(slot_id, OSThreadState::SLEEPING);
     ASSERT_NE(0u, token);
+    slot->markBlockGenerationSampled(ThreadFilter::tokenGeneration(token));
     ThreadEntry entry{4444, slot, slot->lifecycleGeneration(),
                       slot->recordingEpoch()};
     EXPECT_TRUE(registry.isOwnedBlockSuppressionCandidate(entry));
@@ -905,7 +946,9 @@ TEST_F(ThreadRegistryTest, ConcurrentTidReuseInvalidatesSuppressionSnapshot) {
     ASSERT_GE(slot_id, 0);
     ThreadFilter::Slot* slot = registry.slotForId(slot_id);
     ASSERT_NE(nullptr, slot);
-    ASSERT_NE(0u, registry.enterBlockedRun(slot_id, OSThreadState::SLEEPING));
+    u64 token = registry.enterBlockedRun(slot_id, OSThreadState::SLEEPING);
+    ASSERT_NE(0u, token);
+    slot->markBlockGenerationSampled(ThreadFilter::tokenGeneration(token));
     ThreadEntry stale{tid, slot, slot->lifecycleGeneration(),
                       slot->recordingEpoch()};
 
@@ -996,6 +1039,7 @@ TEST_F(ThreadRegistryTest, NewUnfilteredRecordingReclaimsRetainedSlot) {
 
     u64 token = registry.enterBlockedRun(slot_id, OSThreadState::SLEEPING);
     ASSERT_NE(0u, token);
+    slot->markBlockGenerationSampled(ThreadFilter::tokenGeneration(token));
     ThreadEntry stale{tid, slot, slot->lifecycleGeneration(),
                       slot->recordingEpoch()};
     ASSERT_TRUE(registry.isOwnedBlockSuppressionCandidate(stale));

ddprof-test/src/test/java/com/datadoghq/profiler/wallclock/JavaProfilerTaskBlockApiTest.java

@@ -16,10 +16,13 @@
 import java.util.concurrent.atomic.AtomicReference;
 import org.junit.jupiter.api.Assumptions;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.openjdk.jmc.common.item.IItemCollection;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 /** End-to-end coverage for the paired synchronous TaskBlock API. */
@@ -41,6 +44,22 @@ public void pairedApiEmitsTaskBlockWithStack() throws Exception {
     TaskBlockAssertions.assertContainsObservedState(events, "SLEEPING");
   }
 
+  @ParameterizedTest
+  @ValueSource(strings = {"start", "resume"})
+  public void rejectedDuplicateStartOrResumePreservesActiveTaskBlockRecording(String action)
+      throws Exception {
+    IllegalStateException rejected = assertThrows(
+        IllegalStateException.class,
+        () -> profiler.execute(action + "," + getProfilerCommand()));
+    assertEquals("Profiler already started", rejected.getMessage());
+
+    assertTrue(runEligibleBlock(BLOCKER));
+    stopProfiler();
+
+    TaskBlockAssertions.assertContains(
+        verifyEvents("datadog.TaskBlock"), 0L, 0L, BLOCKER, UNBLOCKING_SPAN_ID);
+  }
+
   @Test
   public void invalidAndNestedTokensDoNotLoseCurrentOwner() throws Exception {
     AtomicBoolean recorded = new AtomicBoolean();
@@ -141,7 +160,7 @@ public void virtualThreadCannotMutateCarrierTaskBlockState() throws Exception {
   }
 
   @Test
-  public void liveDumpDoesNotRequireAnEntrySample() throws Exception {
+  public void liveDumpPreservesTaskBlockAfterEntrySample() throws Exception {
     CountDownLatch armed = new CountDownLatch(1);
     CountDownLatch release = new CountDownLatch(1);
     AtomicBoolean recorded = new AtomicBoolean();