configurable buffer size

Author: hanabi1224

docs/docs/users/reference/env_variables.md

@@ -47,6 +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_F3_SIDECAR_FFI_BUILD_OPT_OUT`
 

src/state_migration/mod.rs

@@ -95,6 +95,12 @@ pub fn run_state_migrations<DB>(
 where
     DB: Blockstore + Send + Sync,
 {
+    // ~10MB memory per 10k buffer
+    let db_write_buffer = match std::env::var("FOREST_STATE_MIGRATION_DB_WRITE_BUFFER") {
+        Ok(v) => v.parse().ok(),
+        _ => None,
+    }
+    .unwrap_or(10000);
     let mappings = get_migrations(&chain_config.network);
 
     // Make sure bundle is defined.
@@ -118,7 +124,10 @@ where
         if epoch == chain_config.epoch(height) {
             tracing::info!("Running {height} migration at epoch {epoch}");
             let start_time = std::time::Instant::now();
-            let db = Arc::new(BlockstoreWithWriteBuffer::new(db.clone()));
+            let db = Arc::new(BlockstoreWithWriteBuffer::new_with_capacity(
+                db.clone(),
+                db_write_buffer,
+            ));
             let new_state = migrate(chain_config, &db, parent_state, epoch)?;
             let elapsed = start_time.elapsed();
             // `new_state_actors` is the Go state migration output, log for comparision
@@ -148,7 +157,7 @@ where
     Ok(None)
 }
 
-pub struct BlockstoreWithWriteBuffer<DB: Blockstore> {
+pub(crate) struct BlockstoreWithWriteBuffer<DB: Blockstore> {
     inner: DB,
     buffer: RwLock<HashMap<Cid, Vec<u8>>>,
     buffer_capacity: usize,
@@ -176,10 +185,6 @@ impl<DB: Blockstore> Blockstore for BlockstoreWithWriteBuffer<DB> {
 }
 
 impl<DB: Blockstore> BlockstoreWithWriteBuffer<DB> {
-    pub fn new(inner: DB) -> Self {
-        Self::new_with_capacity(inner, 10000)
-    }
-
     pub fn new_with_capacity(inner: DB, buffer_capacity: usize) -> Self {
         Self {
             inner,