fix(spinlock): bound OptionalSharedLockGuard spin in signal handler paths (#572)

Author: jbachorik

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

@@ -114,7 +114,10 @@
   X(JVMTI_STACKS_REQUESTED, "jvmti_stacks_requested")                         \
   X(JVMTI_STACKS_FAILED_WRONG_PHASE, "jvmti_stacks_failed_wrong_phase")       \
   X(JVMTI_STACKS_FAILED_OTHER, "jvmti_stacks_failed_other")                  \
-  X(JVMTI_STACKS_DROPPED_LOCK, "jvmti_stacks_dropped_lock")
+  /* Delegated stacks dropped at slot-lock. Rec-lock drops from all recording  \
+   * paths (delegated and direct) go into SAMPLES_DROPPED_REC_LOCK. */         \
+  X(JVMTI_STACKS_DROPPED_LOCK, "jvmti_stacks_dropped_lock")                   \
+  X(SAMPLES_DROPPED_REC_LOCK, "samples_dropped_rec_lock")
 #define X_ENUM(a, b) a,
 typedef enum CounterId : int {
   DD_COUNTER_TABLE(X_ENUM) DD_NUM_COUNTERS

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

@@ -2084,6 +2084,8 @@ void FlightRecorder::recordEvent(int lock_index, int tid, u64 call_trace_id,
         rec->flushIfNeeded(buf);
         rec->addThread(lock_index, tid);
       }
+  } else {
+    Counters::increment(SAMPLES_DROPPED_REC_LOCK);
   }
 }
 
@@ -2111,6 +2113,8 @@ void FlightRecorder::recordEventDelegated(int lock_index, int tid,
       rec->flushIfNeeded(buf);
       rec->addThread(lock_index, tid);
     }
+  } else {
+    Counters::increment(SAMPLES_DROPPED_REC_LOCK);
   }
 }
 

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

@@ -18,6 +18,7 @@
 #define _SPINLOCK_H
 
 #include "arch.h"
+#include <cassert>
 
 // Cannot use regular mutexes inside signal handler.
 // This lock is based on CAS busy loop. GCC atomic builtins imply full barrier.
@@ -34,7 +35,7 @@ class alignas(DEFAULT_CACHE_LINE_SIZE) SpinLock {
     static_assert(sizeof(SpinLock) == DEFAULT_CACHE_LINE_SIZE);
   }
 
-  void reset() { _lock = 0; }
+  void reset() { __atomic_store_n(&_lock, 0, __ATOMIC_RELAXED); }
 
   bool tryLock() { return __sync_bool_compare_and_swap(&_lock, 0, 1); }
 
@@ -44,7 +45,15 @@ class alignas(DEFAULT_CACHE_LINE_SIZE) SpinLock {
     }
   }
 
-  void unlock() { __sync_fetch_and_sub(&_lock, 1); }
+  void unlock() {
+    assert(__atomic_load_n(&_lock, __ATOMIC_RELAXED) == 1);
+    __sync_fetch_and_sub(&_lock, 1);
+  }
+
+  // Spin budget for bounded shared acquisition in signal handlers.
+  // Bounds the number of CAS-retry iterations under reader contention;
+  // does NOT bound wall-clock latency (CAS stall time is hardware-dependent).
+  static constexpr int DEFAULT_SHARED_SPIN_BUDGET = 256;
 
   bool tryLockShared() {
     // Spins while no exclusive lock is held and the CAS to acquire a shared
@@ -60,6 +69,23 @@ class alignas(DEFAULT_CACHE_LINE_SIZE) SpinLock {
     return false;
   }
 
+  // Bounded variant for signal-handler paths. Returns false when an exclusive
+  // lock is observed OR the spin budget is exhausted under reader contention.
+  bool tryLockShared(int max_spins) {
+    int value;
+    int spins = 0;
+    while ((value = __atomic_load_n(&_lock, __ATOMIC_ACQUIRE)) <= 0) {
+      if (__sync_bool_compare_and_swap(&_lock, value, value - 1)) {
+        return true;
+      }
+      if (++spins >= max_spins) {
+        return false;
+      }
+      spinPause();
+    }
+    return false;
+  }
+
   void lockShared() {
     int value;
     while ((value = __atomic_load_n(&_lock, __ATOMIC_ACQUIRE)) > 0 ||
@@ -68,7 +94,10 @@ class alignas(DEFAULT_CACHE_LINE_SIZE) SpinLock {
     }
   }
 
-  void unlockShared() { __sync_fetch_and_add(&_lock, 1); }
+  void unlockShared() {
+    assert(__atomic_load_n(&_lock, __ATOMIC_RELAXED) < 0);
+    __sync_fetch_and_add(&_lock, 1);
+  }
 };
 
 // RAII guard classes for automatic lock management
@@ -89,12 +118,17 @@ class SharedLockGuard {
   SharedLockGuard& operator=(SharedLockGuard&&) = delete;
 };
 
+// Acquires a shared lock with a bounded CAS-retry budget. Returns without
+// acquiring (ownsLock() == false) when an exclusive lock is observed or
+// the spin budget is exhausted under reader contention. Safe to use in
+// signal-handler paths.
 class OptionalSharedLockGuard {
   SpinLock* _lock;
 public:
-  explicit OptionalSharedLockGuard(SpinLock* lock) : _lock(lock) {
-    if (!_lock->tryLockShared()) {
-      // Exclusive lock is held; no unlock needed. Only fails when an exclusive lock is observed.
+  explicit OptionalSharedLockGuard(
+      SpinLock* lock, int max_spins = SpinLock::DEFAULT_SHARED_SPIN_BUDGET)
+      : _lock(lock) {
+    if (!_lock->tryLockShared(max_spins)) {
       _lock = nullptr;
     }
   }
@@ -103,7 +137,7 @@ class OptionalSharedLockGuard {
       _lock->unlockShared();
     }
   }
-  bool ownsLock() { return _lock != nullptr; }
+  bool ownsLock() const { return _lock != nullptr; }
 
   // Non-copyable and non-movable
   OptionalSharedLockGuard(const OptionalSharedLockGuard&) = delete;

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

@@ -0,0 +1,116 @@
+/*
+ * 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 <gtest/gtest.h>
+#include "spinLock.h"
+#include "../../main/cpp/gtest_crash_handler.h"
+#include <atomic>
+#include <thread>
+
+static constexpr char SPINLOCK_TEST_NAME[] = "SpinLockTest";
+
+class SpinLockTest : public ::testing::Test {
+protected:
+    void SetUp() override {
+        installGtestCrashHandler<SPINLOCK_TEST_NAME>();
+    }
+    void TearDown() override {
+        restoreDefaultSignalHandlers();
+    }
+
+    SpinLock lock;
+};
+
+// OptionalSharedLockGuard acquires on a free lock and releases on destruction.
+TEST_F(SpinLockTest, OptionalGuard_UncontendedAcquire) {
+    {
+        OptionalSharedLockGuard g(&lock);
+        EXPECT_TRUE(g.ownsLock());
+    }
+    // After destruction the lock must be back to 0 (unlocked).
+    // Verify by taking an exclusive lock — would spin forever if still shared.
+    EXPECT_TRUE(lock.tryLock());
+    lock.unlock();
+}
+
+// When an exclusive lock is held, tryLockShared returns false immediately
+// (first load sees _lock == 1 > 0, exits without spinning).
+TEST_F(SpinLockTest, OptionalGuard_ExclusiveHeld_ImmediateReturn) {
+    lock.lock();
+    OptionalSharedLockGuard g(&lock, 1000000);
+    EXPECT_FALSE(g.ownsLock());
+    lock.unlock();
+}
+
+// The spin budget must be honoured: even with a tiny budget the constructor
+// returns (does not hang) when readers are continuously racing the CAS.
+TEST_F(SpinLockTest, OptionalGuard_SpinBudgetBound) {
+    // A background contender thread rapidly locks/unlocks shared to create
+    // CAS contention on _lock; the guard must return without hanging.
+    std::atomic<bool> stop{false};
+
+    // Background thread hammers shared lock/unlock to create CAS contention.
+    std::thread contender([&] {
+        while (!stop.load(std::memory_order_relaxed)) {
+            lock.lockShared();
+            lock.unlockShared();
+        }
+    });
+
+    // With a very small budget the guard must return promptly, not hang.
+    for (int i = 0; i < 1000; ++i) {
+        OptionalSharedLockGuard g(&lock, 8);
+        // ownsLock() may be true or false depending on timing — we only assert
+        // the constructor returned (i.e. we reach here without hanging).
+        (void)g.ownsLock();
+    }
+
+    stop.store(true, std::memory_order_relaxed);
+    contender.join();
+}
+
+// Verifies that the budget is enforced: with an exclusive lock held,
+// tryLockShared(N) must return false regardless of N.
+TEST_F(SpinLockTest, OptionalGuard_BudgetEnforced_ExclusivePath) {
+    lock.lock();
+    // With any budget, exclusive lock causes immediate false return.
+    EXPECT_FALSE(lock.tryLockShared(1));
+    EXPECT_FALSE(lock.tryLockShared(1000));
+    EXPECT_FALSE(lock.tryLockShared(SpinLock::DEFAULT_SHARED_SPIN_BUDGET));
+    lock.unlock();
+}
+
+// Multiple shared guards can be held simultaneously (readers don't starve each other).
+TEST_F(SpinLockTest, OptionalGuard_SharedReentrancy) {
+    OptionalSharedLockGuard g1(&lock);
+    OptionalSharedLockGuard g2(&lock);
+    EXPECT_TRUE(g1.ownsLock());
+    EXPECT_TRUE(g2.ownsLock());
+    // Exclusive try must fail while shared locks are held.
+    EXPECT_FALSE(lock.tryLock());
+}
+
+// tryLockShared() (unbounded) still works correctly alongside the bounded overload.
+TEST_F(SpinLockTest, TryLockShared_ExclusiveHeld_ReturnsFalse) {
+    lock.lock();
+    EXPECT_FALSE(lock.tryLockShared());
+    lock.unlock();
+}
+
+TEST_F(SpinLockTest, TryLockShared_Free_ReturnsTrue) {
+    EXPECT_TRUE(lock.tryLockShared());
+    lock.unlockShared();
+}