Skip to main content

ethl/storage/
writer.rs

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
10/// This struct is used to archive events to an object store.
11/// It allows for appending logs and flushing them to the store in batches.
12/// It is designed to work with a specific event signature.
13pub 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(()); // Nothing to flush
79        }
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}