feat(sv-dune): Add max_blocks_to_store

Author: pedronauck

services/dune/src/cli.rs

@@ -31,4 +31,11 @@ pub struct Cli {
         help = "Type of storage to use. Options are 'S3' or 'File'."
     )]
     pub storage_type: String,
+
+    #[arg(
+        long,
+        value_name = "MAX_BLOCKS_TO_STORE",
+        env = "MAX_BLOCKS_TO_STORE"
+    )]
+    pub max_blocks_to_store: Option<u16>,
 }

services/dune/src/helpers/mod.rs

@@ -1,7 +1,9 @@
 mod avro;
 mod batch_calculator;
+mod store;
 mod test_helpers;
 
 pub use avro::*;
 pub use batch_calculator::*;
+pub use store::*;
 pub use test_helpers::*;

services/dune/src/helpers/store.rs

@@ -0,0 +1,87 @@
+use std::path::Path;
+
+use anyhow::Result;
+use bytes::Bytes;
+use fuel_data_parser::DataEncoder;
+use fuel_streams_domains::blocks::Block;
+use fuel_streams_types::BlockHeight;
+use surrealkv::Store as SurrealStore;
+
+pub struct Store {
+    inner: SurrealStore,
+}
+
+impl Store {
+    pub fn new(storage_dir: Option<&str>) -> Result<Self> {
+        let mut opts = surrealkv::Options::new();
+        let storage_dir = match storage_dir {
+            Some(file_dir) => Path::new(file_dir).to_path_buf(),
+            None => {
+                let manifest_dir = env!("CARGO_MANIFEST_DIR");
+                let manifest_dir =
+                    Path::new(manifest_dir).join("output/storage");
+                manifest_dir.to_path_buf()
+            }
+        };
+        opts.dir = storage_dir;
+        opts.disk_persistence = true;
+        let store = SurrealStore::new(opts)?;
+        Ok(Self { inner: store })
+    }
+
+    pub async fn get_last_block_saved(&self) -> Result<BlockHeight> {
+        let block_height = {
+            let mut txn = self.inner.begin()?;
+            let key = Bytes::from("last_block_saved");
+            match txn.get(&key)? {
+                Some(value) => serde_json::from_slice(&value)?,
+                None => BlockHeight::from(0),
+            }
+        };
+        Ok(block_height)
+    }
+
+    pub async fn save_last_block(&self, block: &Block) -> Result<()> {
+        let mut txn = self.inner.begin()?;
+        let key = Bytes::from("last_block_saved");
+        let value = block.height.encode_json()?;
+        txn.set(&key, &value)?;
+        txn.commit()?;
+        Ok(())
+    }
+
+    pub async fn get_total_blocks(&self) -> Result<u16> {
+        let total = {
+            let mut txn = self.inner.begin()?;
+            let key = Bytes::from("blocks_saved");
+            match txn.get(&key)? {
+                Some(value) => serde_json::from_slice(&value)?,
+                None => 0u16,
+            }
+        };
+        Ok(total)
+    }
+
+    pub async fn save_total_blocks(&self, blocks_num: u16) -> Result<()> {
+        let current = self.get_total_blocks().await?;
+        let mut txn = self.inner.begin()?;
+        let key = Bytes::from("blocks_saved");
+        let new_total = current + blocks_num;
+        let value = serde_json::to_vec(&new_total)?;
+        txn.set(&key, &value)?;
+        txn.commit()?;
+        Ok(())
+    }
+
+    pub async fn should_continue_processing(
+        &self,
+        max_blocks: Option<u16>,
+    ) -> Result<bool> {
+        if let Some(max) = max_blocks {
+            let total = self.get_total_blocks().await?;
+            Ok(total < max)
+        } else {
+            Ok(true)
+        }
+    }
+}

services/dune/src/main.rs

@@ -1,7 +1,6 @@
-use std::{path::Path, sync::Arc, time::Instant};
+use std::{sync::Arc, time::Instant};
 
 use anyhow::Result;
-use bytes::Bytes;
 use clap::Parser;
 use fuel_data_parser::DataEncoder;
 use fuel_streams_domains::{
@@ -11,9 +10,8 @@ use fuel_streams_domains::{
 };
 use fuel_streams_types::BlockHeight;
 use fuel_web_utils::{shutdown::ShutdownController, tracing::init_tracing};
-use surrealkv::{Options, Store};
 use sv_dune::{
-    helpers::BatchCalculator,
+    helpers::{BatchCalculator, Store},
     processor::BlocksProcessor,
     Cli,
     DuneError,
@@ -32,7 +30,7 @@ async fn main() -> Result<()> {
 
     let opts = QueryOptions::default();
     let last_height_indexed = Block::find_last_block_height(&db, &opts).await?;
-    let last_height_saved = get_last_block_saved(&store).await?;
+    let last_height_saved = store.get_last_block_saved().await?;
     if *last_height_saved >= *last_height_indexed {
         tracing::info!("Last block saved is up to date");
         return Ok(());
@@ -71,19 +69,7 @@ async fn setup_db(db_url: &str) -> Result<Arc<Db>, DuneError> {
 }
 
 fn setup_store(cli: &Cli) -> Result<Store> {
-    let mut opts = Options::new();
-    let storage_dir = match cli.storage_file_dir.as_ref() {
-        Some(file_dir) => Path::new(file_dir).to_path_buf(),
-        None => {
-            let manifest_dir = env!("CARGO_MANIFEST_DIR");
-            let manifest_dir = Path::new(manifest_dir).join("output/storage");
-            manifest_dir.to_path_buf()
-        }
-    };
-    opts.dir = storage_dir.to_owned();
-    opts.disk_persistence = true;
-    let store = Store::new(opts)?;
-    Ok(store)
+    Store::new(cli.storage_file_dir.as_deref())
 }
 
 async fn process_blocks(
@@ -112,6 +98,18 @@ async fn process_blocks(
             break;
         }
 
+        // Check if we should continue processing based on max_blocks
+        if !store
+            .should_continue_processing(cli.max_blocks_to_store)
+            .await?
+        {
+            tracing::info!(
+                "Reached maximum blocks to store ({}), stopping processing",
+                cli.max_blocks_to_store.unwrap()
+            );
+            break;
+        }
+
         let start_time = Instant::now();
         let batch_end = std::cmp::min(
             current_start + (BATCH_SIZE - 1) as u32,
@@ -131,7 +129,8 @@ async fn process_blocks(
 
         process_batch(&processor, &blocks_and_txs, shutdown).await?;
         if let Some((block, _)) = blocks_and_txs.last() {
-            save_last_block(store, block).await?;
+            store.save_last_block(block).await?;
+            store.save_total_blocks(blocks_and_txs.len() as u16).await?;
         }
         tracing::info!("Batch processed in {:?}", start_time.elapsed());
         current_start = batch_end + 1;
@@ -230,31 +229,6 @@ async fn get_initial_height(
     }
 }
 
-async fn get_last_block_saved(store: &Store) -> anyhow::Result<BlockHeight> {
-    let block_height = {
-        let mut txn = store.begin()?;
-        let key = Bytes::from("last_block_saved");
-        match txn.get(&key)? {
-            Some(value) => {
-                Ok::<BlockHeight, DuneError>(serde_json::from_slice(&value)?)
-            }
-            None => Ok(BlockHeight::from(0)),
-        }
-    }?;
-    Ok(block_height)
-}
-
-async fn save_last_block(store: &Store, block: &Block) -> Result<()> {
-    {
-        let mut txn = store.begin()?;
-        let key = Bytes::from("last_block_saved");
-        let value = block.height.encode_json()?;
-        txn.set(&key, &value)?;
-        txn.commit()?;
-    }
-    Ok(())
-}
-
 async fn get_blocks_and_transactions(
     db: &Arc<Db>,
     start_height: BlockHeight,