1use std::sync::Arc;
2
3use crate::storage::{codec::SolEventCodec, store::EventStore};
4use alloy::{json_abi::Event, rpc::types::Log};
5use alloy_primitives::B256;
6use anyhow::Result;
7use object_store::{ObjectStore, path::Path};
8use tracing::debug;
9
10pub struct EventStoreWriter {
14 store: EventStore,
15 batch_size: usize,
16 codec: SolEventCodec,
17 buffered_from_block: Option<u64>,
18 buffered_to_block: Option<u64>,
19 selector: B256,
20}
21
22impl EventStoreWriter {
23 pub fn new(
24 store: Arc<dyn ObjectStore>,
25 path: Path,
26 batch_size: usize,
27 event: Event,
28 ) -> Result<Self> {
29 let codec = SolEventCodec::new(&event)?;
30 let store = EventStore::new(store, path, &codec)?;
31
32 Ok(EventStoreWriter {
33 store,
34 batch_size,
35 codec,
36 buffered_from_block: None,
37 buffered_to_block: None,
38 selector: event.selector(),
39 })
40 }
41
42 pub fn selector(&self) -> B256 {
43 self.selector
44 }
45
46 pub fn event(&self) -> &Event {
47 &self.codec.event
48 }
49
50 fn name(&self) -> String {
51 self.codec.event.name.to_string()
52 }
53
54 pub fn append<'a>(
55 &mut self,
56 logs: impl Iterator<Item = &'a Log>,
57 from_block: u64,
58 to_block: u64,
59 ) -> Result<()> {
60 self.codec.encoder.append_iter(logs)?;
61 if self.buffered_from_block.is_none() {
62 self.buffered_from_block = Some(from_block);
63 }
64 self.buffered_to_block = Some(to_block);
65 Ok(())
66 }
67
68 pub async fn last_block(&self) -> Result<u64> {
69 self.store
70 .stored_range()
71 .await
72 .map(|(_, to_block)| to_block)
73 }
74
75 pub async fn flush(&mut self) -> Result<()> {
76 if self.buffered_from_block.is_none() || self.buffered_to_block.is_none() {
77 debug!("No buffered events to flush for {}", self.name());
78 return Ok(()); }
80
81 debug!(
82 "flushing {:08} records in block range ({:08}, {:08}) for {}",
83 self.codec.encoder.len(),
84 self.buffered_from_block.unwrap(),
85 self.buffered_to_block.unwrap(),
86 self.name(),
87 );
88
89 let records = self.codec.encoder.finish()?;
90 self.store
91 .write_records(
92 (
93 self.buffered_from_block.unwrap(),
94 self.buffered_to_block.unwrap(),
95 ),
96 records,
97 )
98 .await?;
99 self.buffered_from_block = None;
100 Ok(())
101 }
102
103 pub async fn append_and_maybe_flush<'a>(
104 &mut self,
105 logs: impl Iterator<Item = &'a Log>,
106 from_block: u64,
107 to_block: u64,
108 ) -> Result<bool> {
109 self.append(logs, from_block, to_block)?;
110 if self.codec.encoder.len() >= self.batch_size {
111 self.flush().await?;
112 Ok(true)
113 } else {
114 Ok(false)
115 }
116 }
117}