Skip to main content

perpl_sdk/testing/
indexer.rs

1use 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    /// Current state snapshot
97    pub fn snapshot(&'a self) -> RwLockReadGuard<'a, state::Exchange> {
98        self.snapshot.read().unwrap()
99    }
100
101    /// Next available batch of raw events
102    pub async fn next_raw_events(&mut self) -> Option<stream::RawBlockEvents> {
103        self.raw_events_rx.next().await
104    }
105
106    /// Next available batch of state events
107    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    /// Checks if particular request ID has been seen in consumed state events
124    pub fn request_id_seen(&self, request_id: u64) -> bool {
125        self.request_ids.contains(&request_id)
126    }
127
128    /// Waits for specific block or order request being applied to state
129    /// snapshot, skipping all previous state event batches
130    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}