refactor: remove unused worker count and availability tracking from multiproof manager

- Eliminated the worker count variable and associated logic from the multiproof manager.
- Updated multiproof spawning to check for available account workers instead of a general worker count.
- Cleaned up related comments and code for clarity and maintainability.

Author: yongkangc

crates/engine/tree/src/tree/payload_processor/mod.rs

@@ -210,7 +210,6 @@ where
         );
         let storage_worker_count = config.storage_worker_count();
         let account_worker_count = config.account_worker_count();
-        let worker_count = storage_worker_count + account_worker_count;
         let proof_handle = ProofWorkerHandle::new(
             self.executor.handle().clone(),
             consistent_view,
@@ -223,7 +222,6 @@ where
             state_root_config,
             self.executor.clone(),
             proof_handle.clone(),
-            worker_count,
             to_sparse_trie,
             config.multiproof_chunking_enabled().then_some(config.multiproof_chunk_size()),
         );

crates/engine/tree/src/tree/payload_processor/multiproof.rs

@@ -27,15 +27,13 @@ use std::{
     collections::BTreeMap,
     ops::DerefMut,
     sync::{
-        atomic::{AtomicUsize, Ordering},
         mpsc::{channel, Receiver, Sender},
         Arc,
     },
     time::{Duration, Instant},
 };
 use tracing::{debug, error, instrument, trace};
 
-
 /// A trie update that can be applied to sparse trie alongside the proofs for touched parts of the
 /// state.
 #[derive(Default, Debug)]
@@ -355,9 +353,6 @@ pub struct MultiproofManager {
     /// a big account change into different chunks, which may repeatedly
     /// revisit missed leaves.
     missed_leaves_storage_roots: Arc<DashMap<B256, B256>>,
-    /// Counter tracking available workers. Workers decrement when starting work,
-    /// increment when finishing. Used to determine whether to chunk multiproofs.
-    available_workers: Arc<AtomicUsize>,
     /// Metrics
     metrics: MultiProofTaskMetrics,
 }
@@ -368,23 +363,16 @@ impl MultiproofManager {
         executor: WorkloadExecutor,
         metrics: MultiProofTaskMetrics,
         proof_worker_handle: ProofWorkerHandle,
-        worker_count: usize,
     ) -> Self {
         Self {
             inflight: 0,
             executor,
             metrics,
             proof_worker_handle,
             missed_leaves_storage_roots: Default::default(),
-            available_workers: Arc::new(AtomicUsize::new(worker_count)),
         }
     }
 
-    /// Returns true if there are available workers to process tasks.
-    fn has_available_workers(&self) -> bool {
-        self.available_workers.load(Ordering::Relaxed) > 0
-    }
-
     /// Spawns a new multiproof calculation.
     fn spawn(&mut self, input: PendingMultiproofTask) {
         // If there are no proof targets, we can just send an empty multiproof back immediately
@@ -428,15 +416,8 @@ impl MultiproofManager {
 
         let storage_proof_worker_handle = self.proof_worker_handle.clone();
         let missed_leaves_storage_roots = self.missed_leaves_storage_roots.clone();
-        let available_workers = self.available_workers.clone();
 
         self.executor.spawn_blocking(move || {
-            let decremented = available_workers
-                .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
-                    current.checked_sub(1)
-                })
-                .is_ok();
-
             let storage_targets = proof_targets.len();
 
             trace!(
@@ -488,11 +469,6 @@ impl MultiproofManager {
                         .send(MultiProofMessage::ProofCalculationError(error.into()));
                 }
             }
-
-            // Increment at the end - worker is now available again
-            if decremented {
-                available_workers.fetch_add(1, Ordering::Relaxed);
-            }
         });
 
         self.inflight += 1;
@@ -518,15 +494,8 @@ impl MultiproofManager {
         } = multiproof_input;
         let account_proof_worker_handle = self.proof_worker_handle.clone();
         let missed_leaves_storage_roots = self.missed_leaves_storage_roots.clone();
-        let available_workers = self.available_workers.clone();
 
         self.executor.spawn_blocking(move || {
-            let decremented = available_workers
-                .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
-                    current.checked_sub(1)
-                })
-                .is_ok();
-
             let account_targets = proof_targets.len();
             let storage_targets = proof_targets.values().map(|slots| slots.len()).sum::<usize>();
 
@@ -596,11 +565,6 @@ impl MultiproofManager {
                         .send(MultiProofMessage::ProofCalculationError(error.into()));
                 }
             }
-
-            // Increment available workers after no more calculations are expected. 
-            if decremented {
-                available_workers.fetch_add(1, Ordering::Relaxed);
-            }
         });
 
         self.inflight += 1;
@@ -690,7 +654,6 @@ impl MultiProofTask {
         config: MultiProofConfig,
         executor: WorkloadExecutor,
         proof_worker_handle: ProofWorkerHandle,
-        worker_count: usize,
         to_sparse_trie: Sender<SparseTrieUpdate>,
         chunk_size: Option<usize>,
     ) -> Self {
@@ -710,7 +673,6 @@ impl MultiProofTask {
                 executor,
                 metrics.clone(),
                 proof_worker_handle,
-                worker_count,
             ),
             metrics,
         }
@@ -745,8 +707,9 @@ impl MultiProofTask {
         // Process proof targets in chunks.
         let mut chunks = 0;
 
-        // Only chunk if workers are available to take advantage of parallelism.
-        let should_chunk = self.multiproof_manager.has_available_workers();
+        // Only chunk if account workers are available to take advantage of parallelism.
+        let should_chunk =
+            self.multiproof_manager.proof_worker_handle.has_available_account_workers();
 
         let mut spawn = |proof_targets| {
             self.multiproof_manager.spawn(
@@ -884,8 +847,9 @@ impl MultiProofTask {
 
         let mut spawned_proof_targets = MultiProofTargets::default();
 
-        // Only chunk if workers are available to take advantage of parallelism.
-        let should_chunk = self.multiproof_manager.has_available_workers();
+        // Only chunk if account workers are available to take advantage of parallelism.
+        let should_chunk =
+            self.multiproof_manager.proof_worker_handle.has_available_account_workers();
 
         let mut spawn = |hashed_state_update| {
             let proof_targets = get_proof_targets(
@@ -1249,7 +1213,7 @@ mod tests {
             ProofWorkerHandle::new(executor.handle().clone(), consistent_view, task_ctx, 1, 1);
         let channel = channel();
 
-        MultiProofTask::new(config, executor, proof_handle, 1, channel.0, Some(1))
+        MultiProofTask::new(config, executor, proof_handle, channel.0, Some(1))
     }
 
     #[test]

crates/trie/parallel/src/proof_task.rs

@@ -51,6 +51,7 @@ use reth_trie_db::{DatabaseHashedCursorFactory, DatabaseTrieCursorFactory};
 use reth_trie_sparse::provider::{RevealedNode, TrieNodeProvider, TrieNodeProviderFactory};
 use std::{
     sync::{
+        atomic::{AtomicUsize, Ordering},
         mpsc::{channel, Receiver, Sender},
         Arc,
     },
@@ -116,6 +117,7 @@ fn storage_worker_loop<Factory>(
     task_ctx: ProofTaskCtx,
     work_rx: CrossbeamReceiver<StorageWorkerJob>,
     worker_id: usize,
+    available_workers: Arc<AtomicUsize>,
     #[cfg(feature = "metrics")] metrics: ProofTaskTrieMetrics,
 ) where
     Factory: DatabaseProviderFactory<Provider: BlockReader>,
@@ -145,6 +147,12 @@ fn storage_worker_loop<Factory>(
     let mut storage_nodes_processed = 0u64;
 
     while let Ok(job) = work_rx.recv() {
+        // Reserve a worker slot with saturating decrement to prevent underflow.
+        // Only decrement if counter is above zero.
+        let decremented = available_workers
+            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| current.checked_sub(1))
+            .is_ok();
+
         match job {
             StorageWorkerJob::StorageProof { input, result_sender } => {
                 let hashed_address = input.hashed_address;
@@ -186,6 +194,12 @@ fn storage_worker_loop<Factory>(
                     total_processed = storage_proofs_processed,
                     "Storage proof completed"
                 );
+
+                // Release worker slot only if we successfully reserved it.
+                // Prevents counter from wrapping to usize::MAX on underflow.
+                if decremented {
+                    available_workers.fetch_add(1, Ordering::Relaxed);
+                }
             }
 
             StorageWorkerJob::BlindedStorageNode { account, path, result_sender } => {
@@ -224,6 +238,12 @@ fn storage_worker_loop<Factory>(
                     total_processed = storage_nodes_processed,
                     "Blinded storage node completed"
                 );
+
+                // Release worker slot only if we successfully reserved it.
+                // Prevents counter from wrapping to usize::MAX on underflow.
+                if decremented {
+                    available_workers.fetch_add(1, Ordering::Relaxed);
+                }
             }
         }
     }
@@ -246,9 +266,11 @@ fn storage_worker_loop<Factory>(
 ///
 /// Each worker:
 /// 1. Receives `AccountWorkerJob` from crossbeam unbounded channel
-/// 2. Computes result using its dedicated long-lived transaction
-/// 3. Sends result directly to original caller via `std::mpsc`
-/// 4. Repeats until channel closes (graceful shutdown)
+/// 2. Decrements availability counter to mark itself as busy
+/// 3. Computes result using its dedicated long-lived transaction
+/// 4. Sends result directly to original caller via `std::mpsc`
+/// 5. Increments availability counter to mark itself as available
+/// 6. Repeats until channel closes (graceful shutdown)
 ///
 /// # Transaction Reuse
 ///
@@ -269,6 +291,7 @@ fn account_worker_loop<Factory>(
     work_rx: CrossbeamReceiver<AccountWorkerJob>,
     storage_work_tx: CrossbeamSender<StorageWorkerJob>,
     worker_id: usize,
+    available_workers: Arc<AtomicUsize>,
     #[cfg(feature = "metrics")] metrics: ProofTaskTrieMetrics,
 ) where
     Factory: DatabaseProviderFactory<Provider: BlockReader>,
@@ -298,6 +321,12 @@ fn account_worker_loop<Factory>(
     let mut account_nodes_processed = 0u64;
 
     while let Ok(job) = work_rx.recv() {
+        // Reserve a worker slot with saturating decrement to prevent underflow.
+        // Only decrement if counter is above zero.
+        let decremented = available_workers
+            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| current.checked_sub(1))
+            .is_ok();
+
         match job {
             AccountWorkerJob::AccountMultiproof { mut input, result_sender } => {
                 let span = tracing::debug_span!(
@@ -381,6 +410,12 @@ fn account_worker_loop<Factory>(
                     "Account multiproof completed"
                 );
                 drop(_span_guard);
+
+                // Release worker slot only if we successfully reserved it.
+                // Prevents counter from wrapping to usize::MAX on underflow.
+                if decremented {
+                    available_workers.fetch_add(1, Ordering::Relaxed);
+                }
             }
 
             AccountWorkerJob::BlindedAccountNode { path, result_sender } => {
@@ -420,6 +455,12 @@ fn account_worker_loop<Factory>(
                     "Blinded account node completed"
                 );
                 drop(_span_guard);
+
+                // Release worker slot only if we successfully reserved it.
+                // Prevents counter from wrapping to usize::MAX on underflow.
+                if decremented {
+                    available_workers.fetch_add(1, Ordering::Relaxed);
+                }
             }
         }
     }
@@ -866,6 +907,12 @@ pub struct ProofWorkerHandle {
     storage_work_tx: CrossbeamSender<StorageWorkerJob>,
     /// Direct sender to account worker pool
     account_work_tx: CrossbeamSender<AccountWorkerJob>,
+    /// Counter tracking available storage workers. Workers decrement when starting work,
+    /// increment when finishing. Used to determine whether to chunk multiproofs.
+    storage_available_workers: Arc<AtomicUsize>,
+    /// Counter tracking available account workers. Workers decrement when starting work,
+    /// increment when finishing. Used to determine whether to chunk multiproofs.
+    account_available_workers: Arc<AtomicUsize>,
 }
 
 impl ProofWorkerHandle {
@@ -893,6 +940,10 @@ impl ProofWorkerHandle {
         let (storage_work_tx, storage_work_rx) = unbounded::<StorageWorkerJob>();
         let (account_work_tx, account_work_rx) = unbounded::<AccountWorkerJob>();
 
+        // Initialize availability counters
+        let storage_available_workers = Arc::new(AtomicUsize::new(storage_worker_count));
+        let account_available_workers = Arc::new(AtomicUsize::new(account_worker_count));
+
         tracing::debug!(
             target: "trie::proof_task",
             storage_worker_count,
@@ -910,6 +961,7 @@ impl ProofWorkerHandle {
             let view_clone = view.clone();
             let task_ctx_clone = task_ctx.clone();
             let work_rx_clone = storage_work_rx.clone();
+            let storage_available_workers_clone = storage_available_workers.clone();
 
             executor.spawn_blocking(move || {
                 #[cfg(feature = "metrics")]
@@ -921,6 +973,7 @@ impl ProofWorkerHandle {
                     task_ctx_clone,
                     work_rx_clone,
                     worker_id,
+                    storage_available_workers_clone,
                     #[cfg(feature = "metrics")]
                     metrics,
                 )
@@ -946,6 +999,7 @@ impl ProofWorkerHandle {
             let task_ctx_clone = task_ctx.clone();
             let work_rx_clone = account_work_rx.clone();
             let storage_work_tx_clone = storage_work_tx.clone();
+            let account_available_workers_clone = account_available_workers.clone();
 
             executor.spawn_blocking(move || {
                 #[cfg(feature = "metrics")]
@@ -958,6 +1012,7 @@ impl ProofWorkerHandle {
                     work_rx_clone,
                     storage_work_tx_clone,
                     worker_id,
+                    account_available_workers_clone,
                     #[cfg(feature = "metrics")]
                     metrics,
                 )
@@ -972,7 +1027,12 @@ impl ProofWorkerHandle {
 
         drop(_guard);
 
-        Self::new_handle(storage_work_tx, account_work_tx)
+        Self::new_handle(
+            storage_work_tx,
+            account_work_tx,
+            storage_available_workers,
+            account_available_workers,
+        )
     }
 
     /// Creates a new [`ProofWorkerHandle`] with direct access to worker pools.
@@ -981,8 +1041,25 @@ impl ProofWorkerHandle {
     const fn new_handle(
         storage_work_tx: CrossbeamSender<StorageWorkerJob>,
         account_work_tx: CrossbeamSender<AccountWorkerJob>,
+        storage_available_workers: Arc<AtomicUsize>,
+        account_available_workers: Arc<AtomicUsize>,
     ) -> Self {
-        Self { storage_work_tx, account_work_tx }
+        Self {
+            storage_work_tx,
+            account_work_tx,
+            storage_available_workers,
+            account_available_workers,
+        }
+    }
+
+    /// Returns true if there are available storage workers to process tasks.
+    pub fn has_available_storage_workers(&self) -> bool {
+        self.storage_available_workers.load(Ordering::Relaxed) > 0
+    }
+
+    /// Returns true if there are available account workers to process tasks.
+    pub fn has_available_account_workers(&self) -> bool {
+        self.account_available_workers.load(Ordering::Relaxed) > 0
     }
 
     /// Dispatch a storage proof computation to storage worker pool