From a6169da002d982ae8cffacd9aaa294173d1c6421 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dawid=20Ci=C4=99=C5=BCarkiewicz?= <dpc@dpc.pw>
Date: Tue, 12 Dec 2023 20:40:20 -0800
Subject: [PATCH] feat: ELECTRS_ROCKSDB_WRITE_BUFFER_SIZE

In testing setups large write buffer sizes lead to large
db files that potentially waste resources / interfere
with CIs limits.

This change adds an ability to overwrite the default behavior.
---
 src/new_index/db.rs | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/new_index/db.rs b/src/new_index/db.rs
index ef226d34b..9b0c01833 100644
--- a/src/new_index/db.rs
+++ b/src/new_index/db.rs
@@ -1,6 +1,7 @@
 use rocksdb;
 
 use std::path::Path;
+use std::str::FromStr;
 
 use crate::config::Config;
 use crate::util::Bytes;
@@ -80,6 +81,20 @@ pub enum DBFlush {
     Enable,
 }
 
+fn get_write_buffer_size() -> usize {
+    const VAR_NAME: &str = "ELECTRS_ROCKSDB_WRITE_BUFFER_SIZE";
+
+    if let Ok(var) = std::env::var(VAR_NAME) {
+        let size = FromStr::from_str(&var).expect("Could not parse {VAR_NAME}");
+
+        info!("Using custom write buffer size: {size}");
+
+        size
+    } else {
+        256 << 20
+    }
+}
+
 impl DB {
     pub fn open(path: &Path, config: &Config) -> DB {
         debug!("opening DB at {:?}", path);
@@ -89,7 +104,7 @@ impl DB {
         db_opts.set_compaction_style(rocksdb::DBCompactionStyle::Level);
         db_opts.set_compression_type(rocksdb::DBCompressionType::Snappy);
         db_opts.set_target_file_size_base(1_073_741_824);
-        db_opts.set_write_buffer_size(256 << 20);
+        db_opts.set_write_buffer_size(get_write_buffer_size());
         db_opts.set_disable_auto_compactions(true); // for initial bulk load
 
         // db_opts.set_advise_random_on_open(???);