From a79d49bce9a4eb46fea2dd39797cc9ac262e1d47 Mon Sep 17 00:00:00 2001
From: Michael Stetsyuk <michael.stetsyuk@clickhouse.com>
Date: Mon, 27 Jul 2026 13:46:56 +0000
Subject: [PATCH] Add mmap accounting hooks

---
 include/silk/fibers/fiber.h | 11 +++++++++++
 src/fibers/fiber.cpp        | 31 +++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/include/silk/fibers/fiber.h b/include/silk/fibers/fiber.h
index 756b3f7..3032010 100644
--- a/include/silk/fibers/fiber.h
+++ b/include/silk/fibers/fiber.h
@@ -8,6 +8,7 @@
 
 #include <atomic>
 #include <cerrno>
+#include <cstddef>
 #include <cstdint>
 #include <memory>
 #include <utility>
@@ -43,6 +44,11 @@ using FiberParametersDtor = void(void * parameters) noexcept;
  */
 using FiberSwitchCallback = void(Fiber * fiber) noexcept;
 
+/**
+ * Memory mapping callback.  Called when silk maps or unmaps memory outside the C++ heap.
+ */
+using MemoryMapCallback = void(void * ptr, size_t size) noexcept;
+
 /**
  * Packed fiber identity: [category:8 | cpu:10 | counter:46] stored as uint64_t.
  *  category: byte passed to run; the runtime treats it as opaque.
@@ -165,6 +171,11 @@ class FiberScheduler
         // across the OS thread the fiber borrows. Not invoked for proxy fibers.
         FiberSwitchCallback * fiberSuspend = nullptr;
         FiberSwitchCallback * fiberResume = nullptr;
+
+        // Optional hooks for silk memory maps/unmaps outside the heap: fiber stacks (guard pages
+        // excluded) and io_uring rings.
+        MemoryMapCallback * accountMemoryMapped = nullptr;
+        MemoryMapCallback * accountMemoryUnmapped = nullptr;
     };
 
     /**
diff --git a/src/fibers/fiber.cpp b/src/fibers/fiber.cpp
index dcbffc8..cb5fc83 100644
--- a/src/fibers/fiber.cpp
+++ b/src/fibers/fiber.cpp
@@ -309,6 +309,12 @@ Fiber::~Fiber() noexcept
     {
         int r = ::munmap(stack, FiberScheduler::getOptions().fiberStackSize + 2 * kPageSize);
         SILK_ASSERT(!r);
+
+        if (FiberScheduler::getOptions().accountMemoryUnmapped)
+        {
+            FiberScheduler::getOptions().accountMemoryUnmapped(
+                static_cast<uint8_t *>(stack) + kPageSize, FiberScheduler::getOptions().fiberStackSize);
+        }
     }
 }
 
@@ -344,6 +350,11 @@ bool Fiber::initialize(
 
         r = ::mprotect(static_cast<uint8_t *>(stack) + kPageSize + fiberStackSize, kPageSize, PROT_NONE);
         SILK_ASSERT(!r);
+
+        if (FiberScheduler::getOptions().accountMemoryMapped)
+        {
+            FiberScheduler::getOptions().accountMemoryMapped(static_cast<uint8_t *>(stack) + kPageSize, fiberStackSize);
+        }
     }
 
 #if defined(__SANITIZE_ADDRESS__)
@@ -688,6 +699,23 @@ struct FiberScheduler::ProcessorState
     BoundedQueue<Fiber *> readyQueue;
 };
 
+static void accountRingMemoryMappings(const io_uring & ring, MemoryMapCallback * callback) noexcept
+{
+    if (!callback)
+    {
+        return;
+    }
+
+    callback(ring.sq.sqes, ring.sq.sqes_sz);
+    callback(ring.sq.ring_ptr, ring.sq.ring_sz);
+
+    // The kernel serves both rings from one mapping given IORING_FEAT_SINGLE_MMAP.
+    if (ring.cq.ring_ptr != ring.sq.ring_ptr)
+    {
+        callback(ring.cq.ring_ptr, ring.cq.ring_sz);
+    }
+}
+
 void FiberScheduler::ProcessorState::initialize(uint16_t cpu) noexcept
 {
     SILK_ASSERT(cpu < kInvalidProcessorNumber);
@@ -709,6 +737,8 @@ void FiberScheduler::ProcessorState::initialize(uint16_t cpu) noexcept
     // postWakeup posts cross-ring doorbells with IOSQE_CQE_SKIP_SUCCESS to drop the send-side completion.
     SILK_ASSERT(params.features & IORING_FEAT_CQE_SKIP);
 
+    accountRingMemoryMappings(ring, options.accountMemoryMapped);
+
     // Arm the wakeup doorbell. The kernel can end the multishot poll on CQ overflow,
     // so handleCompletionQueueSlow re-arms it through the same path on F_MORE loss.
     enqueueDoorbell();
@@ -726,6 +756,7 @@ void FiberScheduler::ProcessorState::destroy() noexcept
 {
     if (eventFd >= 0)
     {
+        accountRingMemoryMappings(ring, options.accountMemoryUnmapped);
         ::io_uring_queue_exit(&ring);
         ::close(eventFd);
     }