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 #[allow(dead_code)]
123 pub(crate) fn with_ledger_strategy(
124 mut self,
125 ledger_strategy: BoxedStrategy<LedgerState>,
126 ) -> Self {
127 self.ledger_strategy = Some(ledger_strategy);
128 self
129 }
130}
131
132impl Strategy for PreparedChain {
133 type Tree = PreparedChainTree;
134 type Value = <PreparedChainTree as ValueTree>::Value;
135
136 #[allow(clippy::unwrap_in_result)]
137 fn new_tree(&self, runner: &mut TestRunner) -> NewTree<Self> {
138 let mut chain = self.chain.lock().unwrap();
139 if chain.is_none() {
140 let default_ledger_strategy =
142 LedgerState::genesis_strategy(None, NetworkUpgrade::Nu5, None, false);
143 let ledger_strategy = self
144 .ledger_strategy
145 .as_ref()
146 .unwrap_or(&default_ledger_strategy);
147
148 let (network, blocks) = ledger_strategy
149 .prop_flat_map(|ledger| {
150 (
151 Just(ledger.network.clone()),
152 Block::partial_chain_strategy(
153 ledger,
154 MAX_PARTIAL_CHAIN_BLOCKS,
155 check::utxo::transparent_coinbase_spend,
156 self.generate_valid_commitments,
157 ),
158 )
159 })
160 .prop_map(|(network, vec)| {
161 (
162 network,
163 vec.iter()
164 .map(|blk| blk.clone().prepare())
165 .collect::<Vec<_>>(),
166 )
167 })
168 .new_tree(runner)?
169 .current();
170 let history_tree = HistoryTree::from_block(
172 &network,
173 blocks[0].block.clone(),
174 &Default::default(),
176 &Default::default(),
177 )
178 .expect("history tree should be created");
179 *chain = Some((
180 network,
181 Arc::new(SummaryDebug(blocks)),
182 Arc::new(history_tree),
183 ));
184 }
185
186 let chain = chain.clone().expect("should be generated");
187 let count = (2..chain.1.len()).new_tree(runner)?;
191 Ok(PreparedChainTree {
192 chain: chain.1,
193 count,
194 network: chain.0,
195 history_tree: chain.2,
196 })
197 }
198}
199
200pub async fn populated_state(
206 blocks: impl IntoIterator<Item = Arc<Block>>,
207 network: &Network,
208) -> (
209 Buffer<BoxService<Request, Response, BoxError>, Request>,
210 ReadStateService,
211 LatestChainTip,
212 ChainTipChange,
213) {
214 let requests = blocks
215 .into_iter()
216 .map(|block| Request::CommitCheckpointVerifiedBlock(block.into()));
217
218 let (state, read_state, latest_chain_tip, mut chain_tip_change) =
221 StateService::new(Config::ephemeral(), network, Height::MAX, 0);
222 let mut state = Buffer::new(BoxService::new(state), 1);
223
224 let mut responses = FuturesUnordered::new();
225
226 for request in requests {
227 let rsp = state.ready().await.unwrap().call(request);
228 responses.push(rsp);
229 }
230
231 while let Some(rsp) = responses.next().await {
232 rsp.expect("unexpected block commit failure");
235
236 if let Err(timeout_error) = timeout(
238 CHAIN_TIP_UPDATE_WAIT_LIMIT,
239 chain_tip_change.wait_for_tip_change(),
240 )
241 .await
242 .map(|change_result| change_result.expect("unexpected chain tip update failure"))
243 {
244 debug!(
245 timeout = ?humantime_seconds(CHAIN_TIP_UPDATE_WAIT_LIMIT),
246 ?timeout_error,
247 "timeout waiting for chain tip change after committing block"
248 );
249 }
250 }
251
252 (state, read_state, latest_chain_tip, chain_tip_change)
253}