From 486aacfa47febaf440f16526728ab31a5713bf8f Mon Sep 17 00:00:00 2001
From: Amos Wenger <amos@bearcove.net>
Date: Wed, 28 Aug 2024 21:26:21 +0200
Subject: [PATCH] Allow customizing number of buffers

---
 crates/buffet/src/bufpool.rs | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/crates/buffet/src/bufpool.rs b/crates/buffet/src/bufpool.rs
index 7509f898..b69c5be5 100644
--- a/crates/buffet/src/bufpool.rs
+++ b/crates/buffet/src/bufpool.rs
@@ -19,7 +19,20 @@ pub fn initialize_allocator() -> Result<()> {
     #[cfg(feature = "miri")]
     let default_num_bufs = 1024;
 
-    initialize_allocator_with_num_bufs(default_num_bufs)
+    let mut num_bufs = default_num_bufs;
+
+    if let Ok(env_num_bufs) = std::env::var("BUFFET_NUM_BUFS") {
+        if let Ok(parsed_num_bufs) = env_num_bufs.parse::<usize>() {
+            num_bufs = parsed_num_bufs;
+        }
+    }
+
+    let mem_usage_in_mb: f64 = num_bufs as f64 * (BUF_SIZE as usize) as f64 / 1024.0 / 1024.0;
+    eprintln!(
+        "===== Initializing buffer pool with {} buffers, will use {:.2} MiB (override with $BUFFET_NUM_BUFS)",
+        num_bufs, mem_usage_in_mb
+    );
+    initialize_allocator_with_num_bufs(default_num_bufs as _)
 }
 
 impl BufMut {