reth_trie_parallel/
value_encoder.rs

1use crate::proof_task::StorageProofResultMessage;
2use alloy_primitives::{map::B256Map, B256};
3use alloy_rlp::Encodable;
4use core::cell::RefCell;
5use crossbeam_channel::Receiver as CrossbeamReceiver;
6use reth_execution_errors::trie::StateProofError;
7use reth_primitives_traits::Account;
8use reth_storage_errors::db::DatabaseError;
9use reth_trie::{
10    hashed_cursor::HashedStorageCursor,
11    proof_v2::{DeferredValueEncoder, LeafValueEncoder, StorageProofCalculator},
12    trie_cursor::TrieStorageCursor,
13    ProofTrieNodeV2,
14};
15use std::{
16    rc::Rc,
17    time::{Duration, Instant},
18};
19
20/// Stats collected by [`AsyncAccountValueEncoder`] during proof computation.
21///
22/// Tracks time spent waiting for storage proofs and counts of each deferred encoder variant used.
23#[derive(Debug, Default, Clone, Copy)]
24pub(crate) struct ValueEncoderStats {
25    /// Accumulated time spent waiting for storage proof results from dispatched workers.
26    pub(crate) storage_wait_time: Duration,
27    /// Number of times the `Dispatched` variant was used (proof pre-dispatched to workers).
28    pub(crate) dispatched_count: u64,
29    /// Number of times the `Sync` variant was used (synchronous computation).
30    pub(crate) sync_count: u64,
31    /// Number of times a dispatched storage proof had no root node and fell back to sync
32    /// computation.
33    pub(crate) dispatched_missing_root_count: u64,
34}
35
36impl ValueEncoderStats {
37    /// Extends this metrics by adding the values from another.
38    pub(crate) fn extend(&mut self, other: &Self) {
39        self.storage_wait_time += other.storage_wait_time;
40        self.dispatched_count += other.dispatched_count;
41        self.sync_count += other.sync_count;
42        self.dispatched_missing_root_count += other.dispatched_missing_root_count;
43    }
44}
45
46/// Returned from [`AsyncAccountValueEncoder`], used to track an async storage root calculation.
47pub(crate) enum AsyncAccountDeferredValueEncoder<TC, HC> {
48    /// A storage proof job was dispatched to the worker pool.
49    Dispatched {
50        hashed_address: B256,
51        account: Account,
52        /// The receiver for the storage proof result. This is an `Option` so that `encode` can
53        /// take ownership of the receiver, preventing the `Drop` impl from trying to receive on
54        /// it again.
55        proof_result_rx:
56            Option<Result<CrossbeamReceiver<StorageProofResultMessage>, DatabaseError>>,
57        /// Shared storage proof results.
58        storage_proof_results: Rc<RefCell<B256Map<Vec<ProofTrieNodeV2>>>>,
59        /// Shared stats for tracking wait time and counts.
60        stats: Rc<RefCell<ValueEncoderStats>>,
61        /// Shared storage proof calculator for synchronous fallback when dispatched proof has no
62        /// root.
63        storage_calculator: Rc<RefCell<StorageProofCalculator<TC, HC>>>,
64    },
65    /// Synchronous storage root computation.
66    Sync {
67        /// Shared storage proof calculator for computing storage roots.
68        storage_calculator: Rc<RefCell<StorageProofCalculator<TC, HC>>>,
69        hashed_address: B256,
70        account: Account,
71    },
72}
73
74impl<TC, HC> Drop for AsyncAccountDeferredValueEncoder<TC, HC> {
75    fn drop(&mut self) {
76        // If this is a Dispatched encoder that was never consumed via encode(), we need to
77        // receive the storage proof result to avoid losing it.
78        let res = if let Self::Dispatched {
79            hashed_address,
80            proof_result_rx,
81            storage_proof_results,
82            stats,
83            ..
84        } = self
85        {
86            // Take the receiver out - if it's None (already consumed by encode), nothing to do
87            let Some(proof_result_rx) = proof_result_rx.take() else { return };
88
89            (|| -> Result<(), StateProofError> {
90                let rx = proof_result_rx?;
91
92                let wait_start = Instant::now();
93                let msg = rx.recv().map_err(|_| {
94                    StateProofError::Database(DatabaseError::Other(format!(
95                        "Storage proof channel closed for {hashed_address:?}",
96                    )))
97                })?;
98                let result = msg.result?;
99
100                stats.borrow_mut().storage_wait_time += wait_start.elapsed();
101
102                storage_proof_results.borrow_mut().insert(*hashed_address, result.proof);
103                Ok(())
104            })()
105        } else {
106            return;
107        };
108
109        if let Err(err) = res {
110            tracing::error!(target: "trie::parallel", %err, "Failed to collect storage proof in deferred encoder drop");
111        }
112    }
113}
114
115impl<TC, HC> DeferredValueEncoder for AsyncAccountDeferredValueEncoder<TC, HC>
116where
117    TC: TrieStorageCursor,
118    HC: HashedStorageCursor<Value = alloy_primitives::U256>,
119{
120    #[allow(clippy::clone_on_copy)]
121    fn encode(mut self, buf: &mut Vec<u8>) -> Result<(), StateProofError> {
122        let (account, root) = match &mut self {
123            Self::Dispatched {
124                hashed_address,
125                account,
126                proof_result_rx,
127                storage_proof_results,
128                stats,
129                storage_calculator,
130            } => {
131                let hashed_address = *hashed_address;
132                let account = account.clone();
133                // Take the receiver so Drop won't try to receive on it again
134                let proof_result_rx = proof_result_rx
135                    .take()
136                    .expect("encode called on already-consumed Dispatched encoder");
137                let wait_start = Instant::now();
138                let result = proof_result_rx?
139                    .recv()
140                    .map_err(|_| {
141                        StateProofError::Database(DatabaseError::Other(format!(
142                            "Storage proof channel closed for {hashed_address:?}",
143                        )))
144                    })?
145                    .result?;
146                stats.borrow_mut().storage_wait_time += wait_start.elapsed();
147
148                storage_proof_results.borrow_mut().insert(hashed_address, result.proof);
149
150                let root = match result.root {
151                    Some(root) => root,
152                    None => {
153                        // In `compute_v2_account_multiproof` we ensure that all dispatched storage
154                        // proofs computations for which there is also an account proof will return
155                        // a root node, but it could happen randomly that an account which is not in
156                        // the account proof targets, but _is_ in storage proof targets, will need
157                        // to be encoded as part of general trie traversal, so we need to handle
158                        // that case here.
159                        stats.borrow_mut().dispatched_missing_root_count += 1;
160
161                        let mut calculator = storage_calculator.borrow_mut();
162                        let root_node = calculator.storage_root_node(hashed_address)?;
163                        calculator
164                            .compute_root_hash(&[root_node])?
165                            .expect("storage_root_node returns a node at empty path")
166                    }
167                };
168
169                (account, root)
170            }
171            Self::Sync { storage_calculator, hashed_address, account } => {
172                let hashed_address = *hashed_address;
173                let account = account.clone();
174                let mut calculator = storage_calculator.borrow_mut();
175                let root_node = calculator.storage_root_node(hashed_address)?;
176                let storage_root = calculator
177                    .compute_root_hash(&[root_node])?
178                    .expect("storage_root_node returns a node at empty path");
179
180                (account, storage_root)
181            }
182        };
183
184        let account = account.into_trie_account(root);
185        account.encode(buf);
186        Ok(())
187    }
188}
189
190/// Implements the [`LeafValueEncoder`] trait for accounts.
191///
192/// Accepts a set of pre-dispatched storage proof receivers for accounts whose storage roots are
193/// being computed asynchronously by worker threads.
194///
195/// For accounts without pre-dispatched proofs, uses a shared
196/// [`StorageProofCalculator`] to compute storage roots synchronously, reusing cursors across
197/// multiple accounts.
198pub(crate) struct AsyncAccountValueEncoder<TC, HC> {
199    /// Storage proof jobs which were dispatched ahead of time.
200    dispatched: B256Map<CrossbeamReceiver<StorageProofResultMessage>>,
201    /// Tracks storage proof results received from the storage workers. [`Rc`] + [`RefCell`] is
202    /// required because [`DeferredValueEncoder`] cannot have a lifetime.
203    storage_proof_results: Rc<RefCell<B256Map<Vec<ProofTrieNodeV2>>>>,
204    /// Shared storage proof calculator for synchronous computation. Reuses cursors and internal
205    /// buffers across multiple storage root calculations.
206    storage_calculator: Rc<RefCell<StorageProofCalculator<TC, HC>>>,
207    /// Shared stats for tracking wait time and variant counts.
208    stats: Rc<RefCell<ValueEncoderStats>>,
209}
210
211impl<TC, HC> AsyncAccountValueEncoder<TC, HC> {
212    /// Initializes a [`Self`] using a storage proof calculator which will be reused to calculate
213    /// storage roots synchronously.
214    ///
215    /// # Parameters
216    /// - `dispatched`: Pre-dispatched storage proof receivers for target accounts
217    /// - `storage_calculator`: Shared storage proof calculator for synchronous computation
218    pub(crate) fn new(
219        dispatched: B256Map<CrossbeamReceiver<StorageProofResultMessage>>,
220        storage_calculator: Rc<RefCell<StorageProofCalculator<TC, HC>>>,
221    ) -> Self {
222        Self {
223            dispatched,
224            storage_proof_results: Default::default(),
225            storage_calculator,
226            stats: Default::default(),
227        }
228    }
229
230    /// Consume [`Self`] and return all collected storage proofs along with accumulated stats.
231    ///
232    /// This method collects any remaining dispatched proofs that weren't consumed during proof
233    /// calculation and includes their wait time in the returned stats.
234    ///
235    /// # Panics
236    ///
237    /// This method panics if any deferred encoders produced by [`Self::deferred_encoder`] have not
238    /// been dropped.
239    pub(crate) fn finalize(
240        self,
241    ) -> Result<(B256Map<Vec<ProofTrieNodeV2>>, ValueEncoderStats), StateProofError> {
242        let mut storage_proof_results = Rc::into_inner(self.storage_proof_results)
243            .expect("no deferred encoders are still allocated")
244            .into_inner();
245
246        let mut stats = Rc::into_inner(self.stats)
247            .expect("no deferred encoders are still allocated")
248            .into_inner();
249
250        // Any remaining dispatched proofs need to have their results collected.
251        // These are proofs that were pre-dispatched but not consumed during proof calculation.
252        for (hashed_address, rx) in &self.dispatched {
253            let wait_start = Instant::now();
254            let result = rx
255                .recv()
256                .map_err(|_| {
257                    StateProofError::Database(DatabaseError::Other(format!(
258                        "Storage proof channel closed for {hashed_address:?}",
259                    )))
260                })?
261                .result?;
262            stats.storage_wait_time += wait_start.elapsed();
263
264            storage_proof_results.insert(*hashed_address, result.proof);
265        }
266
267        Ok((storage_proof_results, stats))
268    }
269}
270
271impl<TC, HC> LeafValueEncoder for AsyncAccountValueEncoder<TC, HC>
272where
273    TC: TrieStorageCursor,
274    HC: HashedStorageCursor<Value = alloy_primitives::U256>,
275{
276    type Value = Account;
277    type DeferredEncoder = AsyncAccountDeferredValueEncoder<TC, HC>;
278
279    fn deferred_encoder(
280        &mut self,
281        hashed_address: B256,
282        account: Self::Value,
283    ) -> Self::DeferredEncoder {
284        // If the proof job has already been dispatched for this account then it's not necessary to
285        // dispatch another.
286        if let Some(rx) = self.dispatched.remove(&hashed_address) {
287            self.stats.borrow_mut().dispatched_count += 1;
288            return AsyncAccountDeferredValueEncoder::Dispatched {
289                hashed_address,
290                account,
291                proof_result_rx: Some(Ok(rx)),
292                storage_proof_results: self.storage_proof_results.clone(),
293                stats: self.stats.clone(),
294                storage_calculator: self.storage_calculator.clone(),
295            }
296        }
297
298        // If the address didn't have a job dispatched for it then we can assume it has no targets,
299        // and we only need its root.
300
301        // Compute storage root synchronously using the shared calculator
302        self.stats.borrow_mut().sync_count += 1;
303        AsyncAccountDeferredValueEncoder::Sync {
304            storage_calculator: self.storage_calculator.clone(),
305            hashed_address,
306            account,
307        }
308    }
309}
reth_trie_parallel/value_encoder.rs

reth_trie_parallel/
value_encoder.rs
