From fdd3ae4a13fb089be67d83c1c5edaa2095b343bb Mon Sep 17 00:00:00 2001
From: Stepan Bagritsevich <stefan@dragonflydb.io>
Date: Wed, 15 Jan 2025 17:45:52 +0100
Subject: [PATCH] feat(fibers): Add GetStackMargin method

---
 util/fibers/detail/fiber_interface.cc |  4 ++++
 util/fibers/detail/fiber_interface.h  |  4 +++-
 util/fibers/fibers.h                  | 11 ++++++++---
 3 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/util/fibers/detail/fiber_interface.cc b/util/fibers/detail/fiber_interface.cc
index d5c47cce..371cc9dc 100644
--- a/util/fibers/detail/fiber_interface.cc
+++ b/util/fibers/detail/fiber_interface.cc
@@ -185,6 +185,10 @@ ctx::fiber_context FiberInterface::Terminate() {
   return scheduler_->Preempt();
 }
 
+size_t FiberInterface::GetStackMargin(const void* stack_address) const {
+  return static_cast<const uint8_t*>(stack_address) - stack_bottom_;
+}
+
 void FiberInterface::CheckStackMargin() {
   uint32_t check_margin = absl::GetFlag(FLAGS_fiber_safety_margin);
   if (check_margin == 0)
diff --git a/util/fibers/detail/fiber_interface.h b/util/fibers/detail/fiber_interface.h
index 892e8b1d..28df4678 100644
--- a/util/fibers/detail/fiber_interface.h
+++ b/util/fibers/detail/fiber_interface.h
@@ -59,7 +59,6 @@ using FI_SleepHook =
 
 class Scheduler;
 
-
 class FiberInterface {
   friend class Scheduler;
 
@@ -222,6 +221,8 @@ class FiberInterface {
 
   void CheckStackMargin();
 
+  size_t GetStackMargin(const void* stack_address) const;
+
  protected:
   static constexpr uint16_t kTerminatedBit = 0x1;
   static constexpr uint16_t kBusyBit = 0x2;
@@ -260,6 +261,7 @@ class FiberInterface {
   char name_[24];
   uint32_t stack_size_ = 0;
   uint8_t* stack_bottom_ = nullptr;
+
  private:
 #ifndef NDEBUG
   std::function<std::string()> stacktrace_print_cb_;
diff --git a/util/fibers/fibers.h b/util/fibers/fibers.h
index c8b613c8..c28c8419 100644
--- a/util/fibers/fibers.h
+++ b/util/fibers/fibers.h
@@ -175,14 +175,19 @@ inline std::string_view GetName() {
   return fb2::detail::FiberActive()->name();
 }
 
+/* Returns the margin between the provided stack address
+   and the bottom of the fiber's stack. */
+inline uint32_t GetStackMargin(const void* stack_address) {
+  return fb2::detail::FiberActive()->GetStackMargin(stack_address);
+}
+
 inline void CheckSafetyMargin() {
   fb2::detail::FiberActive()->CheckStackMargin();
 }
 
 class PrintLocalsCallback {
-public:
-  template<typename Fn>
-  PrintLocalsCallback(Fn&& fn) {
+ public:
+  template <typename Fn> PrintLocalsCallback(Fn&& fn) {
     fb2::detail::FiberActive()->SetPrintStacktraceCb(std::forward<Fn>(fn));
   }