perpl_sdk/testing/
indexer.rs1use std::{
2 collections::HashSet,
3 pin::pin,
4 sync::{Arc, RwLock, RwLockReadGuard},
5 time::Duration,
6};
7
8use alloy::providers::DynProvider;
9use futures::{
10 SinkExt, StreamExt,
11 channel::mpsc::{self, UnboundedReceiver, UnboundedSender},
12};
13
14use super::TestExchange;
15use crate::{Chain, state, stream, types};
16pub struct Indexer {
17 chain: Chain,
18 provider: DynProvider,
19 snapshot: Arc<RwLock<state::Exchange>>,
20 raw_events_tx: UnboundedSender<stream::RawBlockEvents>,
21 state_events_tx: UnboundedSender<state::StateBlockEvents>,
22}
23
24pub struct IndexedState {
25 snapshot: Arc<RwLock<state::Exchange>>,
26 raw_events_rx: UnboundedReceiver<stream::RawBlockEvents>,
27 state_events_rx: UnboundedReceiver<state::StateBlockEvents>,
28 request_ids: HashSet<u64>,
29}
30
31impl Indexer {
32 pub async fn new(exchange: &TestExchange) -> (Self, IndexedState) {
33 let snapshot = Arc::new(RwLock::new(
34 state::SnapshotBuilder::new(&exchange.chain(), exchange.provider.clone())
35 .with_accounts(
36 exchange
37 .account_address
38 .iter()
39 .map(|e| types::AccountAddressOrID::ID(*e.key()))
40 .collect(),
41 )
42 .build()
43 .await
44 .unwrap(),
45 ));
46
47 let (raw_events_tx, raw_events_rx) = mpsc::unbounded();
48 let (state_events_tx, state_events_rx) = mpsc::unbounded();
49
50 (
51 Self {
52 chain: exchange.chain().clone(),
53 provider: exchange.provider.clone(),
54 snapshot: snapshot.clone(),
55 raw_events_tx,
56 state_events_tx,
57 },
58 IndexedState { snapshot, raw_events_rx, state_events_rx, request_ids: HashSet::new() },
59 )
60 }
61
62 pub async fn run<S, SFut>(mut self, sleep: S)
63 where
64 S: Fn(Duration) -> SFut + Copy,
65 SFut: Future<Output = ()>,
66 {
67 let mut stream = pin!(stream::raw(
68 &self.chain,
69 self.provider,
70 self.snapshot.read().unwrap().instant(),
71 sleep,
72 ));
73 while let Some(batch) = stream.next().await {
74 let batch = batch.unwrap();
75 let res = self.snapshot.write().unwrap().apply_events(&batch);
76 if self.raw_events_tx.send(batch).await.is_err() {
77 break;
78 };
79 match res {
80 Ok(Some(result)) => {
81 if self.state_events_tx.send(result).await.is_err() {
82 break;
83 }
84 },
85 Ok(None) => (),
86 Err(err) => {
87 println!("failed to apply_events: {:#?}", err);
88 break;
89 },
90 }
91 }
92 }
93}
94
95impl<'a> IndexedState {
96 pub fn snapshot(&'a self) -> RwLockReadGuard<'a, state::Exchange> {
98 self.snapshot.read().unwrap()
99 }
100
101 pub async fn next_raw_events(&mut self) -> Option<stream::RawBlockEvents> {
103 self.raw_events_rx.next().await
104 }
105
106 pub async fn next_state_events(&mut self) -> Option<state::StateBlockEvents> {
108 let batch = self.state_events_rx.next().await;
109 if let Some(be) = &batch {
110 be.events().iter().for_each(|ec| {
111 ec.event().iter().for_each(|e| {
112 if let Some(oe) = e.as_order_event()
113 && let Some(rid) = oe.request_id
114 {
115 self.request_ids.insert(rid);
116 }
117 });
118 });
119 }
120 batch
121 }
122
123 pub fn request_id_seen(&self, request_id: u64) -> bool {
125 self.request_ids.contains(&request_id)
126 }
127
128 pub async fn wait_for(&mut self, block_num: Option<u64>, request_id: Option<u64>) -> bool {
131 while let Some(be) = self.state_events_rx.next().await {
132 if block_num.is_some_and(|bn| be.instant().block_number() == bn)
133 || request_id.is_some_and(|rid| {
134 be.events().iter().any(|ec| {
135 ec.event().iter().any(|e| {
136 e.as_order_event()
137 .is_some_and(|oe| oe.request_id.unwrap_or_default() == rid)
138 })
139 })
140 })
141 {
142 return true;
143 }
144 }
145 false
146 }
147}