1use 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 LedgerState,
21};
22
23use crate::{
24 arbitrary::Prepare,
25 service::{check, ReadStateService, StateService},
26 BoxError, ChainTipChange, Config, LatestChainTip, Request, Response, SemanticallyVerifiedBlock,
27};
28
29pub use zebra_chain::block::arbitrary::MAX_PARTIAL_CHAIN_BLOCKS;
30
31pub const CHAIN_TIP_UPDATE_WAIT_LIMIT: Duration = Duration::from_secs(2);
33
34#[derive(Debug)]
35pub struct PreparedChainTree {
36 chain: Arc<SummaryDebug<Vec<SemanticallyVerifiedBlock>>>,
37 count: BinarySearch,
38 network: Network,
39 history_tree: Arc<HistoryTree>,
40}
41
42impl ValueTree for PreparedChainTree {
43 type Value = (
44 Arc<SummaryDebug<Vec<SemanticallyVerifiedBlock>>>,
45 <BinarySearch as ValueTree>::Value,
46 Network,
47 Arc<HistoryTree>,
48 );
49
50 fn current(&self) -> Self::Value {
51 (
52 self.chain.clone(),
53 self.count.current(),
54 self.network.clone(),
55 self.history_tree.clone(),
56 )
57 }
58
59 fn simplify(&mut self) -> bool {
60 self.count.simplify()
61 }
62
63 fn complicate(&mut self) -> bool {
64 self.count.complicate()
65 }
66}
67
68#[derive(Debug, Default)]
69pub struct PreparedChain {
70 chain: std::sync::Mutex<
72 Option<(
73 Network,
74 Arc<SummaryDebug<Vec<SemanticallyVerifiedBlock>>>,
75 Arc<HistoryTree>,
76 )>,
77 >,
78 ledger_strategy: Option<BoxedStrategy<LedgerState>>,
80 generate_valid_commitments: bool,
81}
82
83impl PreparedChain {
84 #[allow(dead_code)]
88 pub(crate) fn new_heartwood() -> Self {
89 let height = Network::iter()
93 .map(|network| {
94 NetworkUpgrade::Heartwood
95 .activation_height(&network)
96 .expect("must have height")
97 })
98 .max()
99 .expect("Network::iter() must return non-empty iterator");
100
101 PreparedChain {
102 ledger_strategy: Some(LedgerState::height_strategy(
103 height,
104 NetworkUpgrade::Nu5,
105 None,
106 false,
107 )),
108 ..Default::default()
109 }
110 }
111
112 #[allow(dead_code)]
116 pub(crate) fn with_valid_commitments(mut self) -> Self {
117 self.generate_valid_commitments = true;
118 self
119 }
120}
121
122impl Strategy for PreparedChain {
123 type Tree = PreparedChainTree;
124 type Value = <PreparedChainTree as ValueTree>::Value;
125
126 #[allow(clippy::unwrap_in_result)]
127 fn new_tree(&self, runner: &mut TestRunner) -> NewTree<Self> {
128 let mut chain = self.chain.lock().unwrap();
129 if chain.is_none() {
130 let default_ledger_strategy =
132 LedgerState::genesis_strategy(NetworkUpgrade::Nu5, None, false);
133 let ledger_strategy = self
134 .ledger_strategy
135 .as_ref()
136 .unwrap_or(&default_ledger_strategy);
137
138 let (network, blocks) = ledger_strategy
139 .prop_flat_map(|ledger| {
140 (
141 Just(ledger.network.clone()),
142 Block::partial_chain_strategy(
143 ledger,
144 MAX_PARTIAL_CHAIN_BLOCKS,
145 check::utxo::transparent_coinbase_spend,
146 self.generate_valid_commitments,
147 ),
148 )
149 })
150 .prop_map(|(network, vec)| {
151 (
152 network,
153 vec.iter()
154 .map(|blk| blk.clone().prepare())
155 .collect::<Vec<_>>(),
156 )
157 })
158 .new_tree(runner)?
159 .current();
160 let history_tree = HistoryTree::from_block(
162 &network,
163 blocks[0].block.clone(),
164 &Default::default(),
166 &Default::default(),
167 )
168 .expect("history tree should be created");
169 *chain = Some((
170 network,
171 Arc::new(SummaryDebug(blocks)),
172 Arc::new(history_tree),
173 ));
174 }
175
176 let chain = chain.clone().expect("should be generated");
177 let count = (2..chain.1.len()).new_tree(runner)?;
181 Ok(PreparedChainTree {
182 chain: chain.1,
183 count,
184 network: chain.0,
185 history_tree: chain.2,
186 })
187 }
188}
189
190pub async fn populated_state(
196 blocks: impl IntoIterator<Item = Arc<Block>>,
197 network: &Network,
198) -> (
199 Buffer<BoxService<Request, Response, BoxError>, Request>,
200 ReadStateService,
201 LatestChainTip,
202 ChainTipChange,
203) {
204 let requests = blocks
205 .into_iter()
206 .map(|block| Request::CommitCheckpointVerifiedBlock(block.into()));
207
208 let (state, read_state, latest_chain_tip, mut chain_tip_change) =
211 StateService::new(Config::ephemeral(), network, Height::MAX, 0);
212 let mut state = Buffer::new(BoxService::new(state), 1);
213
214 let mut responses = FuturesUnordered::new();
215
216 for request in requests {
217 let rsp = state.ready().await.unwrap().call(request);
218 responses.push(rsp);
219 }
220
221 while let Some(rsp) = responses.next().await {
222 rsp.expect("unexpected block commit failure");
225
226 if let Err(timeout_error) = timeout(
228 CHAIN_TIP_UPDATE_WAIT_LIMIT,
229 chain_tip_change.wait_for_tip_change(),
230 )
231 .await
232 .map(|change_result| change_result.expect("unexpected chain tip update failure"))
233 {
234 debug!(
235 timeout = ?humantime_seconds(CHAIN_TIP_UPDATE_WAIT_LIMIT),
236 ?timeout_error,
237 "timeout waiting for chain tip change after committing block"
238 );
239 }
240 }
241
242 (state, read_state, latest_chain_tip, chain_tip_change)
243}