From a1af2c7150b16e554fc8c5d4d222b7581203de93 Mon Sep 17 00:00:00 2001
From: Christian Schwarz <christian@neon.tech>
Date: Mon, 29 Jan 2024 19:22:55 +0000
Subject: [PATCH] broken impl of a permit pool to shave off its allocations

---
 pageserver/src/page_cache.rs             | 18 ++++---
 pageserver/src/page_cache/permit_pool.rs | 69 ++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 8 deletions(-)
 create mode 100644 pageserver/src/page_cache/permit_pool.rs

diff --git a/pageserver/src/page_cache.rs b/pageserver/src/page_cache.rs
index 708f2b02df..b229bddd82 100644
--- a/pageserver/src/page_cache.rs
+++ b/pageserver/src/page_cache.rs
@@ -76,7 +76,7 @@ use std::{
     convert::TryInto,
     sync::{
         atomic::{AtomicU64, AtomicU8, AtomicUsize, Ordering},
-        Arc, Weak,
+        Arc,
     },
     time::Duration,
 };
@@ -164,7 +164,7 @@ struct Slot {
 struct SlotInner {
     key: Option<CacheKey>,
     // for `coalesce_readers_permit`
-    permit: std::sync::Mutex<Weak<PinnedSlotsPermit>>,
+    permit: std::sync::Mutex<permit_pool::PooledWeak>,
     buf: &'static mut SlotContents,
 }
 
@@ -217,17 +217,19 @@ impl Slot {
     }
 }
 
+mod permit_pool;
+
 impl SlotInner {
     /// If there is aready a reader, drop our permit and share its permit, just like we share read access.
-    fn coalesce_readers_permit(&self, permit: PinnedSlotsPermit) -> Arc<PinnedSlotsPermit> {
+    fn coalesce_readers_permit(&self, permit: PinnedSlotsPermit) -> permit_pool::Pooled {
         let mut guard = self.permit.lock().unwrap();
         if let Some(existing_permit) = guard.upgrade() {
             drop(guard);
             drop(permit);
             existing_permit
         } else {
-            let permit = Arc::new(permit);
-            *guard = Arc::downgrade(&permit);
+            let permit = permit_pool::get(permit);
+            *guard = permit_pool::Pooled::downgrade(&permit);
             permit
         }
     }
@@ -255,7 +257,7 @@ struct PinnedSlotsPermit(tokio::sync::OwnedSemaphorePermit);
 /// until the guard is dropped.
 ///
 pub struct PageReadGuard<'i> {
-    _permit: Arc<PinnedSlotsPermit>,
+    _permit: permit_pool::Pooled,
     slot_guard: tokio::sync::RwLockReadGuard<'i, SlotInner>,
 }
 
@@ -321,7 +323,7 @@ impl<'a> PageWriteGuard<'a> {
             PageWriteGuardState::Invalid { inner, _permit } => {
                 assert!(inner.key.is_some());
                 PageReadGuard {
-                    _permit: Arc::new(_permit),
+                    _permit: permit_pool::get(_permit),
                     slot_guard: inner.downgrade(),
                 }
             }
@@ -695,7 +697,7 @@ impl PageCache {
                 inner: tokio::sync::RwLock::new(SlotInner {
                     key: None,
                     buf: slot_contents,
-                    permit: std::sync::Mutex::new(Weak::new()),
+                    permit: std::sync::Mutex::new(permit_pool::PooledWeak::new()),
                 }),
                 usage_count: AtomicU8::new(0),
             })
diff --git a/pageserver/src/page_cache/permit_pool.rs b/pageserver/src/page_cache/permit_pool.rs
new file mode 100644
index 0000000000..1512eb547b
--- /dev/null
+++ b/pageserver/src/page_cache/permit_pool.rs
@@ -0,0 +1,69 @@
+use std::{
+    cell::RefCell,
+    sync::{Arc, Mutex, Weak},
+};
+
+use super::PinnedSlotsPermit;
+
+// Thread-local list of re-usable buffers.
+thread_local! {
+    static POOL: RefCell<Vec<Arc<Mutex<Option<super::PinnedSlotsPermit>>>>> = RefCell::new(Vec::new());
+}
+
+pub(crate) struct Pooled {
+    // Always Some() except when dropping
+    strong: Option<Arc<Mutex<Option<super::PinnedSlotsPermit>>>>,
+}
+
+pub(crate) fn get(permit: PinnedSlotsPermit) -> Pooled {
+    let maybe = POOL.with(|rc| rc.borrow_mut().pop());
+    match maybe {
+        Some(arc) => {
+            let mut inner = arc.lock().unwrap();
+            let prev = inner.replace(permit);
+            assert!(prev.is_none(), "we set it to None() on Pooled::drop()");
+            drop(inner);
+            Pooled { strong: Some(arc) }
+        }
+        None => Pooled {
+            strong: Some(Arc::new(Mutex::new(Some(permit)))),
+        },
+    }
+}
+
+impl Pooled {
+    pub(crate) fn downgrade(this: &Self) -> PooledWeak {
+        PooledWeak {
+            weak: Arc::downgrade(this.strong.as_ref().unwrap()),
+        }
+    }
+}
+
+pub(crate) struct PooledWeak {
+    weak: Weak<Mutex<Option<super::PinnedSlotsPermit>>>,
+}
+
+impl PooledWeak {
+    pub(crate) fn new() -> Self {
+        PooledWeak { weak: Weak::new() }
+    }
+    pub(crate) fn upgrade(&self) -> Option<Pooled> {
+        let arc = self.weak.upgrade()?;
+        todo!("this is broken here, another clone of the now-ugpraded Pooled can be returned");
+        Some(Pooled { strong: arc })
+    }
+}
+
+impl Drop for Pooled {
+    fn drop(&mut self) {
+        let arc = self.strong.take().unwrap();
+        let mut inner = arc.lock().unwrap();
+        let permit: super::PinnedSlotsPermit = inner
+            .take()
+            .expect("we handed it out as Some(), should get it back as Some()");
+        drop(permit);
+        assert!(inner.is_none());
+        drop(inner);
+        POOL.with(|rc| rc.borrow_mut().push(arc))
+    }
+}