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