chore: allow run-time check of safety margin (#353)

Any program with --fiber_safety_margin=<uint> will check its margins upon fiber destruction. Moreover, it is possible to check stack programmatically by calling `ThisFiber::CheckSafetyMargin()` from places of interest. if --fiber_safety_margin is 0, this call is a noop. Signed-off-by: Roman Gershman <[email protected]>
romange · Dec 25, 2024 · 9bd97e6 · 9bd97e6
1 parent b508ad5
commit 9bd97e6
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 28 deletions.
diff --git a/util/fibers/CMakeLists.txt b/util/fibers/CMakeLists.txt
@@ -14,14 +14,6 @@ add_library(fibers2 fibers.cc proactor_base.cc synchronization.cc
             sliding_counter.cc varz.cc fiberqueue_threadpool.cc dns_resolve.cc
             ${FB_LINUX_SRCS})
 
-set(HELIO_STACK_CHECK "" CACHE  STRING
-    "Enables fiber stack overflow checks. If set, shoult be a positive integer, that enforces the minimum available margin in bytes.")
-
-if (HELIO_STACK_CHECK)
-  set_property(SOURCE detail/fiber_interface.cc APPEND PROPERTY COMPILE_DEFINITIONS
-               CHECK_FIBER_STACK_MARGIN=${HELIO_STACK_CHECK})
-endif()
-
 cxx_link(fibers2 base io ${FB_LINUX_LIBS} Boost::context Boost::headers TRDP::cares)
 
 cxx_test(fibers_test fibers2 LABELS CI)

diff --git a/util/fibers/detail/fiber_interface.cc b/util/fibers/detail/fiber_interface.cc
@@ -7,10 +7,15 @@
 
 #include <mutex>  // for g_scheduler_lock
 
+#include "base/flags.h"
 #include "base/logging.h"
 #include "util/fibers/detail/scheduler.h"
 #include "util/fibers/detail/utils.h"
 
+ABSL_FLAG(uint32_t, fiber_safety_margin, 0,
+          "If > 0, ensures the stack each fiber has at least this margin. "
+          "The check is done at fiber destruction time.");
+
 namespace util {
 namespace fb2 {
 using namespace std;
@@ -157,23 +162,7 @@ ctx::fiber_context FiberInterface::Terminate() {
   DCHECK(this == FiberActive());
   DCHECK(!list_hook.is_linked());
 
-#ifdef CHECK_FIBER_STACK_MARGIN
-  {
-    const uint8_t* ptr = stack_bottom_;
-    constexpr uint32_t kHardLimit = (CHECK_FIBER_STACK_MARGIN);
-    while (*ptr == 0xAB) {
-      ++ptr;
-    }
-    uint32_t margin = ptr - stack_bottom_;
-
-    CHECK_GE(margin, kHardLimit) << "Low stack margin for " << name_;
-
-    // Log if margins are within the the orange zone.
-    LOG_IF(INFO, margin < kHardLimit * 1.5)
-        << "Stack margin for " << name_ << ": " << margin << " bytes";
-  }
-#endif
-
+  CheckStackMargin();
   scheduler_->ScheduleTermination(this);
   DVLOG(2) << "Terminating " << name_;
 
@@ -196,12 +185,30 @@ ctx::fiber_context FiberInterface::Terminate() {
   return scheduler_->Preempt();
 }
 
+void FiberInterface::CheckStackMargin() {
+  uint32_t check_margin = absl::GetFlag(FLAGS_fiber_safety_margin);
+  if (check_margin == 0)
+    return;
+
+  const uint8_t* ptr = stack_bottom_;
+  while (*ptr == 0xAB) {
+    ++ptr;
+  }
+  uint32_t margin = ptr - stack_bottom_;
+
+  CHECK_GE(margin, check_margin) << "Low stack margin for " << name_;
+
+  // Log if margins are within the the orange zone.
+  LOG_IF(INFO, margin < check_margin * 1.5)
+      << "Stack margin for " << name_ << ": " << margin << " bytes";
+}
+
 void FiberInterface::InitStackBottom(uint8_t* stack_bottom, uint32_t stack_size) {
   stack_bottom_ = stack_bottom;
 
-#ifdef CHECK_FIBER_STACK_MARGIN
-  memset(stack_bottom_, 0xAB, stack_size);
-#endif
+  if (absl::GetFlag(FLAGS_fiber_safety_margin) > 0) {
+    memset(stack_bottom_, 0xAB, stack_size);
+  }
 }
 
 void FiberInterface::Start(Launch launch) {

diff --git a/util/fibers/detail/fiber_interface.h b/util/fibers/detail/fiber_interface.h
@@ -220,6 +220,8 @@ class FiberInterface {
     return reinterpret_cast<uintptr_t>(stack_bottom_);
   }
 
+  void CheckStackMargin();
+
  protected:
   static constexpr uint16_t kTerminatedBit = 0x1;
   static constexpr uint16_t kBusyBit = 0x2;

diff --git a/util/fibers/fibers.h b/util/fibers/fibers.h
@@ -175,6 +175,10 @@ inline std::string_view GetName() {
   return fb2::detail::FiberActive()->name();
 }
 
+inline void CheckSafetyMargin() {
+  fb2::detail::FiberActive()->CheckStackMargin();
+}
+
 class PrintLocalsCallback {
 public:
   template<typename Fn>