feat: forester: add proper error handling and V2 tree support to (#1821)

sergeytimoshin · web-flow · commit 046d452ab51e · 2025-06-16T21:07:51.000+02:00
`status` cmd
diff --git a/forester/src/forester_status.rs b/forester/src/forester_status.rs
@@ -19,7 +19,7 @@ use crate::{
     ForesterConfig,
 };
 
-pub async fn fetch_forester_status(args: &StatusArgs) {
+pub async fn fetch_forester_status(args: &StatusArgs) -> crate::Result<()> {
     let commitment_config = CommitmentConfig::confirmed();
 
     let client = solana_client::rpc_client::RpcClient::new_with_commitment(
@@ -34,7 +34,7 @@ pub async fn fetch_forester_status(args: &StatusArgs) {
     let mut epoch_pdas = vec![];
     let mut protocol_config_pdas = vec![];
     for (_, account) in registry_accounts {
-        match account.data()[0..8].try_into().unwrap() {
+        match account.data()[0..8].try_into()? {
             ForesterEpochPda::DISCRIMINATOR => {
                 let forester_epoch_pda =
                     ForesterEpochPda::try_deserialize_unchecked(&mut account.data())
@@ -63,12 +63,10 @@ pub async fn fetch_forester_status(args: &StatusArgs) {
 
     let current_active_epoch = protocol_config_pdas[0]
         .config
-        .get_current_active_epoch(slot)
-        .unwrap();
+        .get_current_active_epoch(slot)?;
     let current_registration_epoch = protocol_config_pdas[0]
         .config
-        .get_latest_register_epoch(slot)
-        .unwrap();
+        .get_latest_register_epoch(slot)?;
     println!("Current active epoch: {:?}", current_active_epoch);
 
     println!(
@@ -165,7 +163,7 @@ pub async fn fetch_forester_status(args: &StatusArgs) {
         println!("protocol config: {:?}", protocol_config_pdas[0]);
     }
 
-    let config = Arc::new(ForesterConfig::new_for_status(args).unwrap());
+    let config = Arc::new(ForesterConfig::new_for_status(args)?);
 
     if config.general_config.enable_metrics {
         register_metrics();
@@ -179,19 +177,20 @@ pub async fn fetch_forester_status(args: &StatusArgs) {
         fetch_active_tree: false,
         with_indexer: false,
     })
-    .await
-    .unwrap();
-    let trees = fetch_trees(&rpc).await.unwrap();
+    .await?;
+    let trees = fetch_trees(&rpc).await?;
     if trees.is_empty() {
         warn!("No trees found. Exiting.");
     }
-    run_queue_info(config.clone(), trees.clone(), TreeType::StateV1).await;
-    run_queue_info(config.clone(), trees.clone(), TreeType::AddressV1).await;
+    run_queue_info(config.clone(), trees.clone(), TreeType::StateV1).await?;
+    run_queue_info(config.clone(), trees.clone(), TreeType::AddressV1).await?;
+
+    run_queue_info(config.clone(), trees.clone(), TreeType::StateV2).await?;
+    run_queue_info(config.clone(), trees.clone(), TreeType::AddressV2).await?;
+
     for tree in &trees {
         let tree_type = format!("[{}]", tree.tree_type);
-        let tree_info = get_tree_fullness(&mut rpc, tree.merkle_tree, tree.tree_type)
-            .await
-            .unwrap();
+        let tree_info = get_tree_fullness(&mut rpc, tree.merkle_tree, tree.tree_type).await?;
         let fullness_percentage = tree_info.fullness * 100.0;
         println!(
             "{} Tree {}: Fullness: {:.4}% | Next Index: {} | Threshold: {}",
@@ -255,9 +254,9 @@ pub async fn fetch_forester_status(args: &StatusArgs) {
         println!("No active foresters found for the current epoch.");
     }
 
-    push_metrics(&config.external_services.pushgateway_url)
-        .await
-        .unwrap();
+    push_metrics(&config.external_services.pushgateway_url).await?;
+
+    Ok(())
 }
 
 fn print_current_forester_assignments(
diff --git a/forester/src/lib.rs b/forester/src/lib.rs
@@ -41,7 +41,9 @@ use crate::{
     indexer_type::IndexerType,
     metrics::QUEUE_LENGTH,
     processor::tx_cache::ProcessedHashCache,
-    queue_helpers::fetch_queue_item_data,
+    queue_helpers::{
+        fetch_address_v2_queue_length, fetch_queue_item_data, fetch_state_v2_queue_length,
+    },
     slot_tracker::SlotTracker,
     utils::get_protocol_config,
 };
@@ -50,7 +52,7 @@ pub async fn run_queue_info(
     config: Arc<ForesterConfig>,
     trees: Vec<TreeAccounts>,
     queue_type: TreeType,
-) {
+) -> Result<()> {
     let mut rpc = LightClient::new(LightClientConfig {
         url: config.external_services.rpc_url.to_string(),
         commitment_config: None,
@@ -67,28 +69,47 @@ pub async fn run_queue_info(
         .collect();
 
     for tree_data in trees {
-        if tree_data.tree_type == TreeType::StateV2 {
-            continue;
-        }
-
-        let length = if tree_data.tree_type == TreeType::StateV1 {
-            STATE_NULLIFIER_QUEUE_VALUES
-        } else {
-            ADDRESS_QUEUE_VALUES
+        let queue_length = match tree_data.tree_type {
+            TreeType::StateV1 => fetch_queue_item_data(
+                &mut rpc,
+                &tree_data.queue,
+                0,
+                STATE_NULLIFIER_QUEUE_VALUES,
+                STATE_NULLIFIER_QUEUE_VALUES,
+            )
+            .await?
+            .len(),
+            TreeType::AddressV1 => fetch_queue_item_data(
+                &mut rpc,
+                &tree_data.queue,
+                0,
+                ADDRESS_QUEUE_VALUES,
+                ADDRESS_QUEUE_VALUES,
+            )
+            .await?
+            .len(),
+            TreeType::StateV2 => fetch_state_v2_queue_length(&mut rpc, &tree_data.queue).await?,
+            TreeType::AddressV2 => {
+                fetch_address_v2_queue_length(&mut rpc, &tree_data.merkle_tree).await?
+            }
         };
 
-        let queue_length = fetch_queue_item_data(&mut rpc, &tree_data.queue, 0, length, length)
-            .await
-            .unwrap()
-            .len();
         QUEUE_LENGTH
             .with_label_values(&[&*queue_type.to_string(), &tree_data.merkle_tree.to_string()])
             .set(queue_length as i64);
+
+        let queue_identifier = if tree_data.tree_type == TreeType::AddressV2 {
+            tree_data.merkle_tree.to_string()
+        } else {
+            tree_data.queue.to_string()
+        };
+
         println!(
             "{:?} queue {} length: {}",
-            queue_type, tree_data.queue, queue_length
+            queue_type, queue_identifier, queue_length
         );
     }
+    Ok(())
 }
 
 pub async fn run_pipeline<R: Rpc, I: Indexer + IndexerType<R> + 'static>(
diff --git a/forester/src/main.rs b/forester/src/main.rs
@@ -76,7 +76,7 @@ async fn main() -> Result<(), ForesterError> {
             .await?
         }
         Commands::Status(args) => {
-            forester_status::fetch_forester_status(args).await;
+            forester_status::fetch_forester_status(args).await?;
         }
     }
     Ok(())
diff --git a/forester/src/queue_helpers.rs b/forester/src/queue_helpers.rs
@@ -1,4 +1,7 @@
 use account_compression::QueueAccount;
+use light_batched_merkle_tree::{
+    merkle_tree::BatchedMerkleTreeAccount, queue::BatchedQueueAccount,
+};
 use light_client::rpc::Rpc;
 use light_hash_set::HashSet;
 use solana_sdk::pubkey::Pubkey;
@@ -40,6 +43,41 @@ pub async fn fetch_queue_item_data<R: Rpc>(
     Ok(filtered_queue)
 }
 
+pub async fn fetch_state_v2_queue_length<R: Rpc>(
+    rpc: &mut R,
+    output_queue_pubkey: &Pubkey,
+) -> Result<usize> {
+    trace!(
+        "Fetching StateV2 queue length for {:?}",
+        output_queue_pubkey
+    );
+    if let Some(mut account) = rpc.get_account(*output_queue_pubkey).await? {
+        let output_queue = BatchedQueueAccount::output_from_bytes(account.data.as_mut_slice())?;
+        Ok(output_queue.get_metadata().batch_metadata.next_index as usize)
+    } else {
+        Err(anyhow::anyhow!("account not found"))
+    }
+}
+
+pub async fn fetch_address_v2_queue_length<R: Rpc>(
+    rpc: &mut R,
+    merkle_tree_pubkey: &Pubkey,
+) -> Result<usize> {
+    trace!(
+        "Fetching AddressV2 queue length for {:?}",
+        merkle_tree_pubkey
+    );
+    if let Some(mut account) = rpc.get_account(*merkle_tree_pubkey).await? {
+        let merkle_tree = BatchedMerkleTreeAccount::address_from_bytes(
+            account.data.as_mut_slice(),
+            &(*merkle_tree_pubkey).into(),
+        )?;
+        Ok(merkle_tree.queue_batches.next_index as usize)
+    } else {
+        Err(anyhow::anyhow!("account not found"))
+    }
+}
+
 #[derive(Debug)]
 pub struct QueueUpdate {
     pub pubkey: Pubkey,

Original file line number	Diff line number	Diff line change
`@@ -76,7 +76,7 @@ async fn main() -> Result<(), ForesterError> {`
`76`	`76`	`.await?`
`77`	`77`	`}`
`78`	`78`	`Commands::Status(args) => {`
`79`		`- forester_status::fetch_forester_status(args).await;`
	`79`	`+ forester_status::fetch_forester_status(args).await?;`
`80`	`80`	`}`
`81`	`81`	`}`
`82`	`82`	`Ok(())`