rellocate code and add tests

Author: hanabi1224

docs/docs/users/reference/env_variables.md

@@ -47,7 +47,7 @@ process.
 | `FOREST_TRACE_FILTER_MAX_RESULT`                          | positive integer                 | 500                                            | 1000                                                          | Sets the maximum results returned per request by `trace_filter`                                                       |
 | `FOREST_CHAIN_INDEXER_ENABLED`                            | 1 or true                        | false                                          | 1                                                             | Whether or not to index the chain to support the Ethereum RPC API                                                     |
 | `FOREST_MESSAGES_IN_TIPSET_CACHE_SIZE`                    | positive integer                 | 100                                            | 42                                                            | The size of an internal cache of tipsets to messages                                                                  |
-| `FOREST_STATE_MIGRATION_DB_WRITE_BUFFER`                  | non-negative integer             | 10000                                          | 100000                                                        | The size of db write buffer for state migration (~10MB memory per 10k buffer)                                         |
+| `FOREST_STATE_MIGRATION_DB_WRITE_BUFFER`                  | non-negative integer             | 10000                                          | 100000                                                        | The size of db write buffer for state migration (`~10MB` RAM per `10k` buffer)                                        |
 
 ### `FOREST_F3_SIDECAR_FFI_BUILD_OPT_OUT`
 

src/db/blockstore_with_write_buffer.rs

@@ -0,0 +1,115 @@
+// Copyright 2019-2025 ChainSafe Systems
+// SPDX-License-Identifier: Apache-2.0, MIT
+
+use ahash::{HashMap, HashMapExt};
+use cid::Cid;
+use fvm_ipld_blockstore::Blockstore;
+use itertools::Itertools;
+use parking_lot::RwLock;
+
+pub struct BlockstoreWithWriteBuffer<DB: Blockstore> {
+    inner: DB,
+    buffer: RwLock<HashMap<Cid, Vec<u8>>>,
+    buffer_capacity: usize,
+}
+
+impl<DB: Blockstore> Blockstore for BlockstoreWithWriteBuffer<DB> {
+    fn get(&self, k: &Cid) -> anyhow::Result<Option<Vec<u8>>> {
+        if let Some(v) = self.buffer.read().get(k) {
+            return Ok(Some(v.clone()));
+        }
+        self.inner.get(k)
+    }
+
+    fn has(&self, k: &Cid) -> anyhow::Result<bool> {
+        Ok(self.buffer.read().contains_key(k) || self.inner.has(k)?)
+    }
+
+    fn put_keyed(&self, k: &Cid, block: &[u8]) -> anyhow::Result<()> {
+        {
+            let mut buffer = self.buffer.write();
+            buffer.insert(*k, block.to_vec());
+        }
+        self.flush_buffer_if_needed()
+    }
+}
+
+impl<DB: Blockstore> BlockstoreWithWriteBuffer<DB> {
+    pub fn new_with_capacity(inner: DB, buffer_capacity: usize) -> Self {
+        Self {
+            inner,
+            buffer_capacity,
+            buffer: RwLock::new(HashMap::with_capacity(buffer_capacity)),
+        }
+    }
+
+    fn flush_buffer(&self) -> anyhow::Result<()> {
+        let records = {
+            let mut buffer = self.buffer.write();
+            buffer.drain().collect_vec()
+        };
+        self.inner.put_many_keyed(records)
+    }
+
+    fn flush_buffer_if_needed(&self) -> anyhow::Result<()> {
+        if self.buffer.read().len() >= self.buffer_capacity {
+            self.flush_buffer()
+        } else {
+            Ok(())
+        }
+    }
+}
+
+impl<DB: Blockstore> Drop for BlockstoreWithWriteBuffer<DB> {
+    fn drop(&mut self) {
+        if let Err(e) = self.flush_buffer() {
+            tracing::warn!("{e}");
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{db::MemoryDB, utils::rand::forest_rng};
+    use fvm_ipld_encoding::DAG_CBOR;
+    use multihash_codetable::Code::Blake2b256;
+    use multihash_codetable::MultihashDigest as _;
+    use rand::Rng as _;
+    use std::sync::Arc;
+
+    #[test]
+    fn test_buffer_flush() {
+        const BUFFER_SIZE: usize = 10;
+        const N_RECORDS: usize = 15;
+        let mem_db = Arc::new(MemoryDB::default());
+        let buf_db = BlockstoreWithWriteBuffer::new_with_capacity(mem_db.clone(), BUFFER_SIZE);
+        let mut records = Vec::with_capacity(N_RECORDS);
+        for _ in 0..N_RECORDS {
+            let mut record = [0; 1024];
+            forest_rng().fill(&mut record);
+            let key = Cid::new_v1(DAG_CBOR, Blake2b256.digest(record.as_slice()));
+            records.push((key, record));
+        }
+
+        buf_db.put_many_keyed(records.clone()).unwrap();
+
+        for (i, (k, v)) in records.iter().enumerate() {
+            assert!(buf_db.has(k).unwrap());
+            assert_eq!(buf_db.get(k).unwrap().unwrap().as_slice(), v);
+            if i < BUFFER_SIZE {
+                assert!(mem_db.has(k).unwrap());
+                assert_eq!(mem_db.get(k).unwrap().unwrap().as_slice(), v);
+            } else {
+                assert!(!mem_db.has(k).unwrap());
+            }
+        }
+
+        drop(buf_db);
+
+        for (k, v) in records.iter() {
+            assert!(mem_db.has(k).unwrap());
+            assert_eq!(mem_db.get(k).unwrap().unwrap().as_slice(), v);
+        }
+    }
+}

src/db/mod.rs

@@ -1,13 +1,15 @@
 // Copyright 2019-2025 ChainSafe Systems
 // SPDX-License-Identifier: Apache-2.0, MIT
 
+mod blockstore_with_write_buffer;
 pub mod car;
 mod memory;
 pub mod parity_db;
 pub mod parity_db_config;
 
 mod gc;
 pub mod ttl;
+pub use blockstore_with_write_buffer::BlockstoreWithWriteBuffer;
 pub use gc::MarkAndSweep;
 pub use memory::MemoryDB;
 use setting_keys::ETH_MAPPING_UP_TO_DATE_KEY;

src/state_migration/mod.rs

@@ -6,15 +6,13 @@ use std::sync::{
     atomic::{self, AtomicBool},
 };
 
+use crate::db::BlockstoreWithWriteBuffer;
 use crate::networks::{ChainConfig, Height, NetworkChain};
 use crate::shim::clock::ChainEpoch;
 use crate::shim::state_tree::StateRoot;
-use ahash::{HashMap, HashMapExt};
 use cid::Cid;
 use fvm_ipld_blockstore::Blockstore;
 use fvm_ipld_encoding::CborStore;
-use itertools::Itertools;
-use parking_lot::RwLock;
 
 pub(in crate::state_migration) mod common;
 mod nv17;
@@ -95,7 +93,7 @@ pub fn run_state_migrations<DB>(
 where
     DB: Blockstore + Send + Sync,
 {
-    // ~10MB memory per 10k buffer
+    // ~10MB RAM per 10k buffer
     let db_write_buffer = match std::env::var("FOREST_STATE_MIGRATION_DB_WRITE_BUFFER") {
         Ok(v) => v.parse().ok(),
         _ => None,
@@ -157,66 +155,5 @@ where
     Ok(None)
 }
 
-pub(crate) struct BlockstoreWithWriteBuffer<DB: Blockstore> {
-    inner: DB,
-    buffer: RwLock<HashMap<Cid, Vec<u8>>>,
-    buffer_capacity: usize,
-}
-
-impl<DB: Blockstore> Blockstore for BlockstoreWithWriteBuffer<DB> {
-    fn get(&self, k: &Cid) -> anyhow::Result<Option<Vec<u8>>> {
-        if let Some(v) = self.buffer.read().get(k) {
-            return Ok(Some(v.clone()));
-        }
-        self.inner.get(k)
-    }
-
-    fn has(&self, k: &Cid) -> anyhow::Result<bool> {
-        Ok(self.buffer.read().contains_key(k) || self.inner.has(k)?)
-    }
-
-    fn put_keyed(&self, k: &Cid, block: &[u8]) -> anyhow::Result<()> {
-        {
-            let mut buffer = self.buffer.write();
-            buffer.insert(*k, block.to_vec());
-        }
-        self.flush_buffer_if_needed()
-    }
-}
-
-impl<DB: Blockstore> BlockstoreWithWriteBuffer<DB> {
-    pub fn new_with_capacity(inner: DB, buffer_capacity: usize) -> Self {
-        Self {
-            inner,
-            buffer_capacity,
-            buffer: RwLock::new(HashMap::with_capacity(buffer_capacity)),
-        }
-    }
-
-    fn flush_buffer(&self) -> anyhow::Result<()> {
-        let records = {
-            let mut buffer = self.buffer.write();
-            buffer.drain().collect_vec()
-        };
-        self.inner.put_many_keyed(records)
-    }
-
-    fn flush_buffer_if_needed(&self) -> anyhow::Result<()> {
-        if self.buffer.read().len() >= self.buffer_capacity {
-            self.flush_buffer()
-        } else {
-            Ok(())
-        }
-    }
-}
-
-impl<DB: Blockstore> Drop for BlockstoreWithWriteBuffer<DB> {
-    fn drop(&mut self) {
-        if let Err(e) = self.flush_buffer() {
-            tracing::warn!("{e}");
-        }
-    }
-}
-
 #[cfg(test)]
 mod tests;

src/tool/subcommands/shed_cmd/migration.rs

@@ -10,7 +10,7 @@ use clap::Args;
 use fvm_ipld_blockstore::Blockstore;
 use itertools::Itertools;
 
-use crate::state_migration::BlockstoreWithWriteBuffer;
+use crate::db::BlockstoreWithWriteBuffer;
 use crate::utils::db::CborStoreExt;
 use crate::{
     blocks::CachingBlockHeader,