From 9f087f93f8de5a939b79aeba1d3520762a5e4e7b Mon Sep 17 00:00:00 2001
From: Christian Schwarz <me@cschwarz.com>
Date: Tue, 29 Aug 2023 15:48:42 +0000
Subject: [PATCH] buffer pool impl: re-use allocations

---
 pageserver/src/buffer_pool.rs | 29 +++++++++++++++++++++++------
 1 file changed, 23 insertions(+), 6 deletions(-)

diff --git a/pageserver/src/buffer_pool.rs b/pageserver/src/buffer_pool.rs
index 47c6fc39e1..ecba0c561c 100644
--- a/pageserver/src/buffer_pool.rs
+++ b/pageserver/src/buffer_pool.rs
@@ -2,21 +2,38 @@ use std::cell::RefCell;
 
 use crate::tenant::disk_btree::PAGE_SZ;
 
-pub(crate) type Buffer = Box<[u8; PAGE_SZ]>;
+pub struct Buffer(Option<Box<[u8; PAGE_SZ]>>);
 
 // Thread-local list of re-usable buffers.
 thread_local! {
-    static POOL: RefCell<Vec<Buffer>> = RefCell::new(Vec::new());
+    static POOL: RefCell<Vec<Box<[u8; PAGE_SZ]>>> = RefCell::new(Vec::new());
 }
 
 pub(crate) fn get() -> Buffer {
     let maybe = POOL.with(|rc| rc.borrow_mut().pop());
     match maybe {
-        Some(buf) => buf,
-        Nonne => Box::new([0; PAGE_SZ]),
+        Some(buf) => Buffer(Some(buf)),
+        None => Buffer(Some(Box::new([0; PAGE_SZ]))),
     }
 }
 
-pub(crate) fn put(buf: Buffer) {
-    POOL.with(|rc| rc.borrow_mut().push(buf))
+impl Drop for Buffer {
+    fn drop(&mut self) {
+        let buf = self.0.take().unwrap();
+        POOL.with(|rc| rc.borrow_mut().push(buf))
+    }
+}
+
+impl std::ops::Deref for Buffer {
+    type Target = [u8; PAGE_SZ];
+
+    fn deref(&self) -> &Self::Target {
+        self.0.as_ref().unwrap().as_ref()
+    }
+}
+
+impl std::ops::DerefMut for Buffer {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.0.as_mut().unwrap().as_mut()
+    }
 }