Skip to main content

zebra_state/service/
arbitrary.rs

1//! Arbitrary data generation and test setup for Zebra's state.
2
3use std::{sync::Arc, time::Duration};
4
5use futures::{stream::FuturesUnordered, StreamExt};
6use proptest::{
7    num::usize::BinarySearch,
8    prelude::*,
9    strategy::{NewTree, ValueTree},
10    test_runner::TestRunner,
11};
12use tokio::time::timeout;
13use tower::{buffer::Buffer, util::BoxService, Service, ServiceExt};
14
15use zebra_chain::{
16    block::{Block, Height},
17    chain_tip::ChainTip,
18    fmt::{humantime_seconds, SummaryDebug},
19    history_tree::HistoryTree,
20    parameters::{Network, NetworkUpgrade},
21    primitives::zcash_history::BlockCommitmentTreeRoots,
22    LedgerState,
23};
24
25use crate::{
26    arbitrary::Prepare,
27    service::{check, ReadStateService, StateService},
28    BoxError, ChainTipChange, CheckpointVerifiedBlock, Config, LatestChainTip, Request, Response,
29    SemanticallyVerifiedBlock,
30};
31
32pub use zebra_chain::block::arbitrary::MAX_PARTIAL_CHAIN_BLOCKS;
33
34/// How long we wait for chain tip updates before skipping them.
35pub const CHAIN_TIP_UPDATE_WAIT_LIMIT: Duration = Duration::from_secs(2);
36
37#[derive(Debug)]
38pub struct PreparedChainTree {
39    chain: Arc<SummaryDebug<Vec<SemanticallyVerifiedBlock>>>,
40    count: BinarySearch,
41    network: Network,
42    history_tree: Arc<HistoryTree>,
43}
44
45impl ValueTree for PreparedChainTree {
46    type Value = (
47        Arc<SummaryDebug<Vec<SemanticallyVerifiedBlock>>>,
48        <BinarySearch as ValueTree>::Value,
49        Network,
50        Arc<HistoryTree>,
51    );
52
53    fn current(&self) -> Self::Value {
54        (
55            self.chain.clone(),
56            self.count.current(),
57            self.network.clone(),
58            self.history_tree.clone(),
59        )
60    }
61
62    fn simplify(&mut self) -> bool {
63        self.count.simplify()
64    }
65
66    fn complicate(&mut self) -> bool {
67        self.count.complicate()
68    }
69}
70
71#[derive(Debug, Default)]
72pub struct PreparedChain {
73    // the proptests are threaded (not async), so we want to use a threaded mutex here
74    chain: std::sync::Mutex<
75        Option<(
76            Network,
77            Arc<SummaryDebug<Vec<SemanticallyVerifiedBlock>>>,
78            Arc<HistoryTree>,
79        )>,
80    >,
81    // the strategy for generating LedgerStates. If None, it calls [`LedgerState::genesis_strategy`].
82    ledger_strategy: Option<BoxedStrategy<LedgerState>>,
83    generate_valid_commitments: bool,
84}
85
86impl PreparedChain {
87    /// Create a PreparedChain strategy with Heartwood-onward blocks.
88    // dead_code is allowed because the function is called only by tests,
89    // but the code is also compiled when proptest-impl is activated.
90    #[allow(dead_code)]
91    pub(crate) fn new_heartwood() -> Self {
92        // The history tree only works with Heartwood onward.
93        // Since the network will be chosen later, we pick the larger
94        // between the mainnet and testnet Heartwood activation heights.
95        let height = Network::iter()
96            .map(|network| {
97                NetworkUpgrade::Heartwood
98                    .activation_height(&network)
99                    .expect("must have height")
100            })
101            .max()
102            .expect("Network::iter() must return non-empty iterator");
103
104        PreparedChain {
105            ledger_strategy: Some(LedgerState::height_strategy(
106                height,
107                NetworkUpgrade::Nu5,
108                None,
109                false,
110            )),
111            ..Default::default()
112        }
113    }
114
115    /// Transform the strategy to use valid commitments in the block.
116    ///
117    /// This is slower so it should be used only when needed.
118    #[allow(dead_code)]
119    pub(crate) fn with_valid_commitments(mut self) -> Self {
120        self.generate_valid_commitments = true;
121        self
122    }
123
124    /// Set the ledger strategy for the prepared chain.
125    #[allow(dead_code)]
126    pub(crate) fn with_ledger_strategy(
127        mut self,
128        ledger_strategy: BoxedStrategy<LedgerState>,
129    ) -> Self {
130        self.ledger_strategy = Some(ledger_strategy);
131        self
132    }
133}
134
135impl Strategy for PreparedChain {
136    type Tree = PreparedChainTree;
137    type Value = <PreparedChainTree as ValueTree>::Value;
138
139    #[allow(clippy::unwrap_in_result)]
140    fn new_tree(&self, runner: &mut TestRunner) -> NewTree<Self> {
141        let mut chain = self.chain.lock().unwrap();
142        if chain.is_none() {
143            // TODO: use the latest network upgrade (#1974)
144            let default_ledger_strategy =
145                LedgerState::genesis_strategy(None, NetworkUpgrade::Nu5, None, false);
146            let ledger_strategy = self
147                .ledger_strategy
148                .as_ref()
149                .unwrap_or(&default_ledger_strategy);
150
151            let (network, blocks) = ledger_strategy
152                .prop_flat_map(|ledger| {
153                    (
154                        Just(ledger.network.clone()),
155                        Block::partial_chain_strategy(
156                            ledger,
157                            MAX_PARTIAL_CHAIN_BLOCKS,
158                            check::utxo::transparent_coinbase_spend,
159                            self.generate_valid_commitments,
160                        ),
161                    )
162                })
163                .prop_map(|(network, vec)| {
164                    (
165                        network,
166                        vec.iter()
167                            .map(|blk| blk.clone().prepare())
168                            .collect::<Vec<_>>(),
169                    )
170                })
171                .new_tree(runner)?
172                .current();
173            // Generate a history tree from the first block
174            let history_tree = HistoryTree::from_block(
175                &network,
176                blocks[0].block.clone(),
177                // Dummy roots since this is only used for tests
178                BlockCommitmentTreeRoots {
179                    sapling: &Default::default(),
180                    orchard: &Default::default(),
181                    ironwood: &Default::default(),
182                },
183            )
184            .expect("history tree should be created");
185            *chain = Some((
186                network,
187                Arc::new(SummaryDebug(blocks)),
188                Arc::new(history_tree),
189            ));
190        }
191
192        let chain = chain.clone().expect("should be generated");
193        // The generated chain should contain at least two blocks:
194        // 1. the zeroth genesis block, and
195        // 2. a first block.
196        let count = (2..chain.1.len()).new_tree(runner)?;
197        Ok(PreparedChainTree {
198            chain: chain.1,
199            count,
200            network: chain.0,
201            history_tree: chain.2,
202        })
203    }
204}
205
206/// Initialize a state service with blocks, and return:
207/// - a read-write [`StateService`]
208/// - a read-only [`ReadStateService`]
209/// - a [`LatestChainTip`]
210/// - a [`ChainTipChange`] tracker
211pub async fn populated_state(
212    blocks: impl IntoIterator<Item = Arc<Block>>,
213    network: &Network,
214) -> (
215    Buffer<BoxService<Request, Response, BoxError>, Request>,
216    ReadStateService,
217    LatestChainTip,
218    ChainTipChange,
219) {
220    let blocks: Vec<CheckpointVerifiedBlock> = blocks.into_iter().map(Into::into).collect();
221    let tip_height = blocks.iter().map(|block| block.height).max();
222    let requests = blocks
223        .into_iter()
224        .map(Request::CommitCheckpointVerifiedBlock);
225
226    // TODO: write a test that checks the finalized to non-finalized transition with UTXOs,
227    //       and set max_checkpoint_height and checkpoint_verify_concurrency_limit correctly.
228    let (state, read_state, latest_chain_tip, mut chain_tip_change) =
229        StateService::new(Config::ephemeral(), network, Height::MAX, 0).await;
230    let mut state = Buffer::new(BoxService::new(state), 1);
231
232    let mut responses = FuturesUnordered::new();
233
234    for request in requests {
235        let rsp = state.ready().await.unwrap().call(request);
236        responses.push(rsp);
237    }
238
239    while let Some(rsp) = responses.next().await {
240        rsp.expect("unexpected block commit failure");
241    }
242
243    // The block write thread updates the chain tip just after it sends each commit response.
244    // The tip channel only keeps the latest tip, so blocks committed close together produce a
245    // single tip change. Waiting for one tip change per block would time out once for each of
246    // those blocks, so wait for the tip to reach the highest block instead.
247    if let Err(timeout_error) = timeout(CHAIN_TIP_UPDATE_WAIT_LIMIT, async {
248        while latest_chain_tip.best_tip_height() < tip_height {
249            chain_tip_change
250                .wait_for_tip_change()
251                .await
252                .expect("unexpected chain tip update failure");
253        }
254    })
255    .await
256    {
257        debug!(
258            timeout = ?humantime_seconds(CHAIN_TIP_UPDATE_WAIT_LIMIT),
259            ?timeout_error,
260            "timeout waiting for chain tip change after committing blocks"
261        );
262    }
263
264    // Mark the current tip as seen, so callers only see tip changes after these blocks.
265    let _ = chain_tip_change.last_tip_change();
266
267    (state, read_state, latest_chain_tip, chain_tip_change)
268}