[PROF-15201] fix(profiler): close two TOCTOU races between SIGPROF handler and JFR lifecycle (#614)

fix(profiler): drain in-flight signal handlers before JFR teardown

Close two TOCTOU races between profiling signal handlers and the JFR
lifecycle that could cause use-after-free or hangs during recording-cycle
rotation.

   - _enabled: unify __ATOMIC_RELEASE/ACQUIRE across CTimer, CTimerJvmti,
     ITimer, ITimerJvmti, PerfEvents, WallClockASGCT, WallClockJvmti so
     arm64 handlers cannot observe stale _enabled=true after
     disableEngines() has returned.
   - Ordering: Profiler::stop() now runs disableEngines() → drain →
     per-engine stop() → _jfr.stop(); Profiler::start() enables wall
     before its pthread starts and CPU only after _jfr.start() completes.
   - Drain: new SignalInflight module + InflightGuard RAII counter on a
     dedicated cache line; every JFR-writing signal handler participates.
     drain() spins to 0 with a 200 ms deadline; on timeout, stop() returns
     Error and leaves _state == RUNNING so retry is safe (no non-idempotent
     engine teardown has run).

co-authored by claude
Great reviews from both @rkennke, @jbachorik

Author: r1viollet

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

@@ -25,9 +25,15 @@
 #include <signal.h>
 
 class CTimer : public Engine {
+private:
+  // cppcheck-suppress unusedPrivateFunction
+  static void signalHandler(int signo, siginfo_t *siginfo, void *ucontext);
+
 protected:
-  // This is accessed from signal handlers, so must be async-signal-safe
+  // Accessed from signal handlers (including CTimerJvmti subclass), so must
+  // be async-signal-safe. Mutated via enableEvents().
   static bool _enabled;
+
   static long _interval;
   static CStack _cstack;
   static int _signal;
@@ -38,10 +44,6 @@ class CTimer : public Engine {
   int registerThread(int tid);
   void unregisterThread(int tid);
 
-private:
-  // cppcheck-suppress unusedPrivateFunction
-  static void signalHandler(int signo, siginfo_t *siginfo, void *ucontext);
-
 public:
   const char *units() { return "ns"; }
 

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

@@ -20,6 +20,7 @@
 #include "counters.h"
 #include "guards.h"
 #include "ctimer.h"
+#include "signalInflight.h"
 #include "debugSupport.h"
 #include "jvmThread.h"
 #include "libraries.h"
@@ -180,6 +181,8 @@ void CTimer::stop() {
   for (int i = 0; i < _max_timers; i++) {
     unregisterThread(i);
   }
+  // Note: SignalInflight::drain() is called by Profiler::stop() after
+  // disableEngines(); see signalInflight.h.
 }
 
 Error CTimerJvmti::check(Arguments &args) {
@@ -210,6 +213,8 @@ void CTimerJvmti::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
   }
   Counters::increment(CTIMER_SIGNAL_OWN);
 
+  InflightGuard inflight;
+
   CriticalSection cs;
   if (!cs.entered()) {
     return;
@@ -261,6 +266,8 @@ void CTimer::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
   }
   Counters::increment(CTIMER_SIGNAL_OWN);
 
+  InflightGuard inflight;
+
   // Atomically try to enter critical section - prevents all reentrancy races
   CriticalSection cs;
   if (!cs.entered()) {
@@ -270,8 +277,9 @@ void CTimer::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
   int saved_errno = errno;
   // we want to ensure memory order because of the possibility the instance gets
   // cleared
-  if (!__atomic_load_n(&_enabled, __ATOMIC_ACQUIRE))
+  if (!__atomic_load_n(&_enabled, __ATOMIC_ACQUIRE)) {
     return;
+  }
   int tid = 0;
   ProfiledThread *current = ProfiledThread::currentSignalSafe();
   assert(current == nullptr || !current->isDeepCrashHandler());

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

@@ -20,13 +20,14 @@
 #include "jvmThread.h"
 #include "os.h"
 #include "profiler.h"
+#include "signalInflight.h"
 #include "stackWalker.h"
 #include "thread.h"
 #include "threadState.inline.h"
 #include "guards.h"
 #include <sys/time.h>
 
-volatile bool ITimer::_enabled = false;
+bool ITimer::_enabled = false;
 long ITimer::_interval;
 CStack ITimer::_cstack;
 
@@ -38,7 +39,8 @@ void ITimer::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
   // is therefore vulnerable to the foreign-SIGPROF deadlock scenario this
   // feature addresses. Use CTimer (the default) when signal-origin
   // validation is required.
-  if (!_enabled)
+  InflightGuard inflight;
+  if (!__atomic_load_n(&_enabled, __ATOMIC_ACQUIRE))
     return;
 
   // Atomically try to enter critical section - prevents all reentrancy races
@@ -99,11 +101,12 @@ void ITimer::stop() {
   setitimer(ITIMER_PROF, &tv, NULL);
 }
 
-volatile bool ITimerJvmti::_enabled = false;
+bool ITimerJvmti::_enabled = false;
 long ITimerJvmti::_interval = 0;
 
 void ITimerJvmti::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
   SIGNAL_HANDLER_GUARD();
+  InflightGuard inflight;
   CriticalSection cs;
   if (!cs.entered()) {
     return;

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

@@ -22,7 +22,7 @@
 
 class ITimer : public Engine {
 private:
-  static volatile bool _enabled;
+  static bool _enabled;
   static long _interval;
   static CStack _cstack;
 
@@ -39,7 +39,9 @@ class ITimer : public Engine {
   Error start(Arguments &args);
   void stop();
 
-  inline void enableEvents(bool enabled) { _enabled = enabled; }
+  inline void enableEvents(bool enabled) {
+    __atomic_store_n(&_enabled, enabled, __ATOMIC_RELEASE);
+  }
 };
 
 // CPU-time engine identical to ITimer in its timer mechanism (process-wide
@@ -51,7 +53,7 @@ class ITimer : public Engine {
 // stack walking rather than relying on the signal-frame PC.
 class ITimerJvmti : public Engine {
 private:
-  static volatile bool _enabled;
+  static bool _enabled;
   static long _interval;
 
   static void signalHandler(int signo, siginfo_t *siginfo, void *ucontext);

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

@@ -30,7 +30,7 @@ class StackContext;
 
 class PerfEvents : public Engine {
 private:
-  static volatile bool _enabled;
+  static bool _enabled;
   static int _max_events;
   static PerfEvent *_events;
   static PerfEventType *_event_type;
@@ -62,7 +62,9 @@ class PerfEvents : public Engine {
 
   static const char *getEventName(int event_id);
 
-  inline void enableEvents(bool enabled) { _enabled = enabled; }
+  inline void enableEvents(bool enabled) {
+    __atomic_store_n(&_enabled, enabled, __ATOMIC_RELEASE);
+  }
 };
 
 #else

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

@@ -29,6 +29,7 @@
 #include "os.h"
 #include "perfEvents.h"
 #include "profiler.h"
+#include "signalInflight.h"
 #include "spinLock.h"
 #include "stackFrame.h"
 #include "stackWalker.h"
@@ -560,7 +561,7 @@ class PerfEvent : public SpinLock {
   friend class PerfEvents;
 };
 
-volatile bool PerfEvents::_enabled = false;
+bool PerfEvents::_enabled = false;
 int PerfEvents::_max_events = -1;
 PerfEvent *PerfEvents::_events = NULL;
 PerfEventType *PerfEvents::_event_type = NULL;
@@ -735,6 +736,7 @@ void PerfEvents::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
     // Looks like an external signal; don't treat as a profiling event
     return;
   }
+  InflightGuard inflight;
   // Atomically try to enter critical section - prevents all reentrancy races
   CriticalSection cs;
   if (!cs.entered()) {
@@ -745,7 +747,7 @@ void PerfEvents::signalHandler(int signo, siginfo_t *siginfo, void *ucontext) {
     current->noteCPUSample(Profiler::instance()->recordingEpoch());
   }
   int tid = current != NULL ? current->tid() : OS::threadId();
-  if (_enabled) {
+  if (__atomic_load_n(&_enabled, __ATOMIC_ACQUIRE)) {
     Shims::instance().setSighandlerTid(tid);
 
     u64 counter = readCounter(siginfo, ucontext);

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

@@ -14,6 +14,7 @@
 #include "common.h"
 #include "counters.h"
 #include "ctimer.h"
+#include "signalInflight.h"
 #include "dwarf.h"
 #include "flightRecorder.h"
 #include "itimer.h"
@@ -866,11 +867,6 @@ void Profiler::switchLibraryTrap(bool enable) {
   __atomic_store_n(_dlopen_entry, impl, __ATOMIC_RELEASE);
 }
 
-void Profiler::enableEngines() {
-  _cpu_engine->enableEvents(true);
-  _wall_engine->enableEvents(true);
-}
-
 void Profiler::disableEngines() {
   _cpu_engine->enableEvents(false);
   _wall_engine->enableEvents(false);
@@ -1402,8 +1398,6 @@ Error Profiler::start(Arguments &args, bool reset) {
     _libs->updateBuildIds();
   }
 
-  enableEngines();
-
   // Refresher must be running before the trap fires: dlopen_hook's
   // signal-context branch only marks dirty and relies on the refresher
   // to call refresh() within REFRESH_INTERVAL_NS (500 ms).
@@ -1417,7 +1411,6 @@ Error Profiler::start(Arguments &args, bool reset) {
   _num_context_attributes = args._context_attributes.size();
   error = _jfr.start(args, reset);
   if (error) {
-    disableEngines();
     switchLibraryTrap(false);
     _libs->stopRefresher();
     return error;
@@ -1434,9 +1427,16 @@ Error Profiler::start(Arguments &args, bool reset) {
     }
   }
   if ((_event_mask & EM_WALL) && _wall_engine != &noop_engine) {
+    // Enable wall clock BEFORE starting its pthread: timerLoopCommon() checks
+    // _enabled once at startup and exits early if false, so the pthread would
+    // otherwise die immediately and wall profiling would be silently dead for
+    // the recording. JFR is already initialized by this point (_jfr.start()
+    // ran above).
+    _wall_engine->enableEvents(true);
     error = _wall_engine->start(args);
     if (error) {
       Log::warn("%s", error.message());
+      _wall_engine->enableEvents(false);
       error = Error::OK; // recoverable
     } else {
       activated |= EM_WALL;
@@ -1492,6 +1492,15 @@ Error Profiler::start(Arguments &args, bool reset) {
     // TODO: find a better way to resolve the thread name.
     onThreadStart(nullptr, nullptr, nullptr);
 
+    // Enable CPU profiling last. CPU SIGPROF handlers write JFR events, so
+    // _enabled must not become true until _jfr.start() has completed; we also
+    // wait until after _cpu_engine->start() so the very first signal a thread
+    // ever sees finds the engine fully wired. Unlike wall clock, CPU handlers
+    // re-check _enabled on every signal, so enabling here (rather than before
+    // start()) does not lose samples.
+    // Paired with drainInflight() on the stop side.
+    _cpu_engine->enableEvents(true);
+
     _state.store(RUNNING, std::memory_order_release);
     _start_time = time(NULL);
     __atomic_add_fetch(&_epoch, 1, __ATOMIC_RELAXED);
@@ -1516,8 +1525,22 @@ Error Profiler::stop() {
     return Error("Profiler is not active");
   }
 
+  // Order matters: disable engines first so the _enabled check inside signal
+  // handlers will fail for any new signal delivered from now on. drain() then
+  // waits for handlers that already passed the check to leave their JFR write
+  // path. Only once those are gone is it safe to run the rest of the teardown
+  // (engine stops, JFR teardown) without risking use-after-free.
   disableEngines();
 
+  if (!SignalInflight::drain()) {
+    // Signal handlers stuck past the timeout. Leave state == RUNNING so the
+    // caller cannot start() against half-torn-down JFR and so engine stops
+    // (notably BaseWallClock::stop()'s pthread_join) are not double-invoked.
+    // The operation is idempotent on retry: disableEngines() above is an atomic
+    // store, and no other engine stop has run yet.
+    return Error("signal handlers did not drain; teardown skipped, retry stop()");
+  }
+
   if (_event_mask & EM_ALLOC)
     _alloc_engine->stop();
   if (_event_mask & EM_NATIVEMEM)

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

@@ -127,7 +127,6 @@ class alignas(alignof(SpinLock)) Profiler {
   void switchLibraryTrap(bool enable);
   static void prewarmUnwinder();
 
-  void enableEngines();
   void disableEngines();
 
   void onThreadStart(jvmtiEnv *jvmti, JNIEnv *jni, jthread thread);

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

@@ -0,0 +1,66 @@
+/*
+ * Copyright 2026, Datadog, Inc.
+ *
+ * Licensed 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.
+ */
+
+#include "signalInflight.h"
+#include "log.h"
+
+#include <errno.h>
+#include <sched.h>
+#include <time.h>
+
+alignas(64) int SignalInflight::_counter = 0;
+
+// 200 ms: long enough for any legitimate signal handler to finish,
+// short enough to avoid a perceptible hang if a handler is somehow stuck.
+static const long DRAIN_TIMEOUT_NS = 200000000L;
+
+bool SignalInflight::drain() {
+  if (__atomic_load_n(&_counter, __ATOMIC_ACQUIRE) == 0) {
+    return true; // fast path: nothing in flight
+  }
+
+  struct timespec deadline;
+  if (clock_gettime(CLOCK_MONOTONIC, &deadline) != 0) {
+    Log::error("SignalInflight::drain: clock_gettime(CLOCK_MONOTONIC) failed (errno=%d). "
+               "Skipping JFR teardown to prevent use-after-free.", errno);
+    return false;
+  }
+  deadline.tv_nsec += DRAIN_TIMEOUT_NS;
+  if (deadline.tv_nsec >= 1000000000L) {
+    deadline.tv_sec += 1;
+    deadline.tv_nsec -= 1000000000L;
+  }
+
+  while (__atomic_load_n(&_counter, __ATOMIC_ACQUIRE) > 0) {
+    struct timespec now;
+    if (clock_gettime(CLOCK_MONOTONIC, &now) != 0) {
+      Log::error("SignalInflight::drain: clock_gettime(CLOCK_MONOTONIC) failed (errno=%d). "
+                 "Skipping JFR teardown to prevent use-after-free.", errno);
+      return false;
+    }
+    if (now.tv_sec > deadline.tv_sec ||
+        (now.tv_sec == deadline.tv_sec && now.tv_nsec >= deadline.tv_nsec)) {
+      int remaining = __atomic_load_n(&_counter, __ATOMIC_ACQUIRE);
+      Log::error("SignalInflight::drain: timed out after %ldms waiting for "
+                 "%d in-flight signal handler(s). Skipping JFR teardown to "
+                 "prevent use-after-free. This indicates a stuck signal handler.",
+                 DRAIN_TIMEOUT_NS / 1000000L, remaining);
+      return false;
+    }
+    sched_yield();
+  }
+  return true;
+}