Skip to main content

jetstreamer_firehose/
node_reader.rs

1use crate::LOG_MODULE;
2use crate::SharedError;
3use crate::epochs::{epoch_to_slot_range, slot_to_epoch};
4use crate::firehose::FirehoseError;
5use crate::index::{SlotOffsetIndexError, slot_to_range};
6use crate::node::{Node, NodeWithCid, NodesWithCids, parse_any_from_cbordata};
7use crate::utils;
8use cid::Cid;
9use once_cell::sync::Lazy;
10use reqwest::RequestBuilder;
11use rseek::Seekable;
12use std::io;
13use std::io::SeekFrom;
14use std::sync::atomic::{AtomicU64, Ordering};
15use std::time::Instant;
16use std::vec::Vec;
17use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt};
18use tokio::time::sleep;
19
20const MAX_VARINT_LEN_64: usize = 10;
21const MIN_SEEK_SPACING_MS: u64 = 15;
22static SEEK_START_INSTANT: Lazy<Instant> = Lazy::new(Instant::now);
23static LAST_SEEK_HIT_TIME: AtomicU64 = AtomicU64::new(0);
24
25/// Total CAR section bytes read across all readers and threads since process start. Frontends
26/// (e.g. the TUI) sample this to derive an overall data rate.
27pub static TOTAL_BYTES_READ: AtomicU64 = AtomicU64::new(0);
28
29/// Reads an unsigned LEB128-encoded integer from the provided async reader.
30pub async fn read_uvarint<R: AsyncRead + Unpin>(reader: &mut R) -> io::Result<u64> {
31    let mut x = 0u64;
32    let mut s = 0u32;
33    let mut buffer = [0u8; 1];
34
35    for i in 0..MAX_VARINT_LEN_64 {
36        reader.read_exact(&mut buffer).await?;
37        let b = buffer[0];
38        if b < 0x80 {
39            if i == MAX_VARINT_LEN_64 - 1 && b > 1 {
40                return Err(io::Error::new(
41                    io::ErrorKind::InvalidData,
42                    "uvarint overflow",
43                ));
44            }
45            return Ok(x | ((b as u64) << s));
46        }
47        x |= ((b & 0x7f) as u64) << s;
48        s += 7;
49
50        if s > 63 {
51            return Err(io::Error::new(
52                io::ErrorKind::InvalidData,
53                "uvarint too long",
54            ));
55        }
56    }
57    Err(io::Error::new(
58        io::ErrorKind::InvalidData,
59        "uvarint overflow",
60    ))
61}
62
63/// Raw DAG-CBOR node paired with its [`Cid`].
64#[derive(Clone, PartialEq, Eq, Hash)]
65pub struct RawNode {
66    /// Content identifier for the node.
67    pub cid: Cid,
68    /// Raw CBOR-encoded bytes for the node.
69    pub data: Vec<u8>,
70}
71
72// Debug trait for RawNode
73impl core::fmt::Debug for RawNode {
74    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
75        f.debug_struct("RawNode")
76            .field("cid", &self.cid)
77            .field("data", &self.data)
78            .finish()
79    }
80}
81
82impl RawNode {
83    /// Creates a new [`RawNode`] from a CID and CBOR payload.
84    pub const fn new(cid: Cid, data: Vec<u8>) -> RawNode {
85        RawNode { cid, data }
86    }
87
88    /// Parses the CBOR payload into a typed [`Node`].
89    pub fn parse(&self) -> Result<Node, SharedError> {
90        match parse_any_from_cbordata(self.data.clone()) {
91            Ok(node) => Ok(node),
92            Err(err) => {
93                println!("Error: {:?}", err);
94                Err(Box::new(std::io::Error::other("Unknown type".to_owned())))
95            }
96        }
97    }
98
99    /// Reads a [`RawNode`] from a CAR section cursor.
100    pub async fn from_cursor(cursor: &mut io::Cursor<Vec<u8>>) -> Result<RawNode, SharedError> {
101        let cid_version = read_uvarint(cursor).await?;
102        // println!("CID version: {}", cid_version);
103
104        let multicodec = read_uvarint(cursor).await?;
105        // println!("Multicodec: {}", multicodec);
106
107        // Multihash hash function code.
108        let hash_function = read_uvarint(cursor).await?;
109        // println!("Hash function: {}", hash_function);
110
111        // Multihash digest length.
112        let digest_length = read_uvarint(cursor).await?;
113        // println!("Digest length: {}", digest_length);
114
115        if digest_length > 64 {
116            return Err(Box::new(std::io::Error::other(format!(
117                "Digest length too long, position={}",
118                cursor.position()
119            ))));
120        }
121
122        // reac actual digest
123        let mut digest = vec![0u8; digest_length as usize];
124        cursor.read_exact(&mut digest).await?;
125
126        // the rest is the data
127        let mut data = vec![];
128        cursor.read_to_end(&mut data).await?;
129
130        // println!("Data: {:?}", data);
131
132        let ha = multihash::Multihash::wrap(hash_function, digest.as_slice())?;
133
134        match cid_version {
135            0 => {
136                let cid = Cid::new_v0(ha)?;
137                let raw_node = RawNode::new(cid, data);
138                Ok(raw_node)
139            }
140            1 => {
141                let cid = Cid::new_v1(multicodec, ha);
142                let raw_node = RawNode::new(cid, data);
143                Ok(raw_node)
144            }
145            _ => Err(Box::new(std::io::Error::other(
146                "Unknown CID version".to_owned(),
147            ))),
148        }
149    }
150}
151
152/// Trait for readers that can report their total length.
153pub trait Len {
154    /// Returns the total number of bytes available.
155    fn len(&self) -> u64;
156    /// Returns `true` when the length is zero.
157    fn is_empty(&self) -> bool {
158        self.len() == 0
159    }
160}
161
162impl<F> Len for Seekable<F>
163where
164    F: Fn() -> RequestBuilder + Send + Sync + 'static,
165{
166    fn len(&self) -> u64 {
167        self.file_size.unwrap_or(0)
168    }
169}
170
171/// Incremental reader that produces typed nodes from an Old Faithful CAR stream.
172pub struct NodeReader<R: AsyncRead + AsyncSeek + Len> {
173    /// Underlying stream yielding Old Faithful CAR bytes.
174    pub reader: R,
175    /// Cached Old Faithful CAR header data.
176    pub header: Vec<u8>,
177    /// Number of Old Faithful items that have been read so far.
178    pub item_index: u64,
179    /// Whether the next seek already holds the global seek-spacing permit.
180    seek_permit_primed: bool,
181}
182
183impl<R: AsyncRead + Unpin + AsyncSeek + Len> NodeReader<R> {
184    /// Wraps an async reader and primes it for Old Faithful CAR decoding.
185    pub const fn new(reader: R) -> NodeReader<R> {
186        NodeReader {
187            reader,
188            header: vec![],
189            item_index: 0,
190            seek_permit_primed: false,
191        }
192    }
193
194    /// Returns the raw Old Faithful CAR header, fetching and caching it on first use.
195    pub async fn read_raw_header(&mut self) -> Result<Vec<u8>, SharedError> {
196        if !self.header.is_empty() {
197            return Ok(self.header.clone());
198        };
199        let header_length = read_uvarint(&mut self.reader).await?;
200        if header_length > 1024 {
201            return Err(Box::new(std::io::Error::other(
202                "Header length too long".to_owned(),
203            )));
204        }
205        let mut header = vec![0u8; header_length as usize];
206        self.reader.read_exact(&mut header).await?;
207
208        self.header.clone_from(&header);
209
210        let clone = header.clone();
211        Ok(clone.as_slice().to_owned())
212    }
213
214    /// Seeks the underlying reader to the first Old Faithful CAR section belonging to `slot`,
215    /// so that reading forward yields all of the slot's nodes (transactions, entries, rewards)
216    /// followed by its Block node. If `slot` is missing from the index (skipped on-chain), the
217    /// seek advances to the next present slot.
218    pub async fn seek_to_slot(&mut self, slot: u64) -> Result<(), FirehoseError> {
219        self.seek_to_slot_inner(slot).await
220    }
221
222    /// Acquires the global seek-spacing permit ahead of a [`Self::seek_to_slot`] call, so the
223    /// cross-thread pacing wait (which can queue for many seconds when hundreds of threads
224    /// seek at once) happens outside any timeout the caller wraps the seek in. The permit is
225    /// consumed by the next seek; an unprimed seek acquires it itself.
226    pub async fn prime_seek_permit(&mut self) {
227        wait_for_seek_hit_slot().await;
228        self.seek_permit_primed = true;
229    }
230
231    async fn seek_to_slot_inner(&mut self, slot: u64) -> Result<(), FirehoseError> {
232        if self.header.is_empty() {
233            self.read_raw_header()
234                .await
235                .map_err(FirehoseError::SeekToSlotError)?;
236        };
237
238        let epoch = slot_to_epoch(slot);
239        let (epoch_start, epoch_end_inclusive) = epoch_to_slot_range(epoch);
240        let mut current = slot;
241        loop {
242            match slot_to_range(current).await {
243                Ok((offset, _)) => {
244                    log::info!(
245                        target: LOG_MODULE,
246                        "Seeking to slot {} in epoch {} @ offset {}",
247                        current,
248                        epoch,
249                        offset
250                    );
251                    return self.seek_to_offset(offset).await;
252                }
253                Err(SlotOffsetIndexError::SlotNotFound(..)) => {
254                    if current >= epoch_end_inclusive {
255                        // No present slot remains between `slot` and the end of its epoch, and
256                        // offsets from other epochs are meaningless for this reader's CAR
257                        // stream. Position the reader just past the last present slot's data so
258                        // callers read the epoch's trailing non-block nodes and observe a clean
259                        // end-of-epoch.
260                        return self.seek_to_epoch_tail(slot, epoch, epoch_start).await;
261                    }
262                    log::warn!(
263                        target: LOG_MODULE,
264                        "Slot {} not found in index, seeking to next slot",
265                        current
266                    );
267                    current += 1;
268                }
269                // Surface index failures as SlotOffsetIndexError so callers can invalidate the
270                // cached epoch index and retry with fresh data.
271                Err(err) => return Err(FirehoseError::SlotOffsetIndexError(err)),
272            }
273        }
274    }
275
276    /// Seeks just past the end of the data of the last present slot preceding `slot`, for use
277    /// when no slot at or after `slot` exists in its epoch.
278    async fn seek_to_epoch_tail(
279        &mut self,
280        slot: u64,
281        epoch: u64,
282        epoch_start: u64,
283    ) -> Result<(), FirehoseError> {
284        let mut candidate = slot;
285        while candidate > epoch_start {
286            candidate -= 1;
287            match slot_to_range(candidate).await {
288                Ok((offset, length)) => {
289                    log::warn!(
290                        target: LOG_MODULE,
291                        "No slot at or after {} is present in epoch {}; seeking past the last \
292                         present slot {}",
293                        slot,
294                        epoch,
295                        candidate
296                    );
297                    return self.seek_to_offset(offset + length).await;
298                }
299                Err(SlotOffsetIndexError::SlotNotFound(..)) => continue,
300                Err(err) => return Err(FirehoseError::SlotOffsetIndexError(err)),
301            }
302        }
303        // An index that reports every slot of an epoch as skipped is corrupt; surface an index
304        // error so callers invalidate the cached epoch index and refetch instead of retrying
305        // against the same bad data forever.
306        Err(FirehoseError::SlotOffsetIndexError(
307            SlotOffsetIndexError::EpochHasNoIndexedSlots(epoch),
308        ))
309    }
310
311    async fn seek_to_offset(&mut self, offset: u64) -> Result<(), FirehoseError> {
312        if !std::mem::take(&mut self.seek_permit_primed) {
313            wait_for_seek_hit_slot().await;
314        }
315        self.reader
316            .seek(SeekFrom::Start(offset))
317            .await
318            .map_err(|e| FirehoseError::SeekToSlotError(Box::new(e)))?;
319        Ok(())
320    }
321
322    #[allow(clippy::should_implement_trait)]
323    /// Reads the next raw node from the Old Faithful stream without parsing it.
324    pub async fn next(&mut self) -> Result<RawNode, SharedError> {
325        if self.header.is_empty() {
326            self.read_raw_header().await?;
327        };
328
329        // println!("Item index: {}", item_index);
330        self.item_index += 1;
331
332        // Read and decode the uvarint prefix (length of CID + data)
333        let section_size = read_uvarint(&mut self.reader).await?;
334        // println!("Section size: {}", section_size);
335
336        if section_size > utils::MAX_ALLOWED_SECTION_SIZE as u64 {
337            return Err(Box::new(std::io::Error::other(
338                "Section size too long".to_owned(),
339            )));
340        }
341
342        // read whole item
343        let mut item = vec![0u8; section_size as usize];
344        self.reader.read_exact(&mut item).await?;
345        TOTAL_BYTES_READ.fetch_add(section_size, Ordering::Relaxed);
346
347        // dump item bytes as numbers
348        // println!("Item bytes: {:?}", item);
349
350        // now create a cursor over the item
351        let mut cursor = io::Cursor::new(item);
352
353        RawNode::from_cursor(&mut cursor).await
354    }
355
356    /// Reads and parses the next node, returning it paired with its [`Cid`].
357    pub async fn next_parsed(&mut self) -> Result<NodeWithCid, SharedError> {
358        let raw_node = self.next().await?;
359        let cid = raw_node.cid;
360        Ok(NodeWithCid::new(cid, raw_node.parse()?))
361    }
362
363    /// Continues reading nodes until the next block is encountered.
364    pub async fn read_until_block(&mut self) -> Result<NodesWithCids, SharedError> {
365        let mut nodes = NodesWithCids::new();
366        loop {
367            let node = match self.next_parsed().await {
368                Ok(node) => node,
369                Err(e)
370                    if e.downcast_ref::<io::Error>()
371                        .is_some_and(|io_err| io_err.kind() == io::ErrorKind::UnexpectedEof) =>
372                {
373                    break;
374                }
375                Err(e) => return Err(e),
376            };
377            if node.get_node().is_block() {
378                nodes.push(node);
379                break;
380            }
381            nodes.push(node);
382        }
383        Ok(nodes)
384    }
385
386    /// Returns the number of Old Faithful CAR items that have been yielded so far.
387    pub const fn get_item_index(&self) -> u64 {
388        self.item_index
389    }
390}
391
392fn seek_monotonic_millis() -> u64 {
393    let elapsed = SEEK_START_INSTANT.elapsed().as_millis();
394    if elapsed > u64::MAX as u128 {
395        u64::MAX
396    } else {
397        elapsed as u64
398    }
399}
400
401async fn wait_for_seek_hit_slot() {
402    loop {
403        let now_ms = seek_monotonic_millis();
404        let last_hit = LAST_SEEK_HIT_TIME.load(Ordering::Relaxed);
405        let since_last = now_ms.saturating_sub(last_hit);
406        if since_last < MIN_SEEK_SPACING_MS {
407            // Sleep out the remainder of the spacing window instead of spinning; with
408            // hundreds of threads queued this otherwise burns CPU for the whole drain.
409            sleep(std::time::Duration::from_millis(
410                MIN_SEEK_SPACING_MS - since_last,
411            ))
412            .await;
413            continue;
414        }
415        if LAST_SEEK_HIT_TIME
416            .compare_exchange(last_hit, now_ms, Ordering::SeqCst, Ordering::SeqCst)
417            .is_ok()
418        {
419            return;
420        }
421        // Lost the race for this window; the winner's spacing interval starts now.
422        sleep(std::time::Duration::from_millis(MIN_SEEK_SPACING_MS)).await;
423    }
424}
425
426/// Extracts a CID from a DAG-CBOR link value.
427pub fn cid_from_cbor_link(val: &serde_cbor::Value) -> Result<cid::Cid, SharedError> {
428    if let serde_cbor::Value::Bytes(b) = val
429        && b.first() == Some(&0)
430    {
431        return Ok(cid::Cid::try_from(b[1..].to_vec())?);
432    }
433    Err("invalid DAG‑CBOR link encoding".into())
434}
435
436#[cfg(test)]
437mod tests {
438    use super::*;
439    use serde_cbor::Value::{self, Array, Bytes, Integer, Map, Null, Text};
440    use sha2::{Digest, Sha256};
441
442    impl Len for io::Cursor<Vec<u8>> {
443        fn len(&self) -> u64 {
444            self.get_ref().len() as u64
445        }
446    }
447
448    fn append_varint(mut value: usize, bytes: &mut Vec<u8>) {
449        while value >= 128 {
450            bytes.push((value as u8 & 0x7f) | 0x80);
451            value >>= 7;
452        }
453        bytes.push(value as u8);
454    }
455
456    fn archive(nodes: &[Value]) -> (Vec<u8>, Vec<Cid>, Vec<u8>) {
457        let mut sections = Vec::new();
458        let mut cids = Vec::new();
459        for node in nodes {
460            let data = serde_cbor::to_vec(node).unwrap();
461            let digest = Sha256::digest(&data);
462            let cid = Cid::new_v1(0x71, multihash::Multihash::wrap(0x12, &digest).unwrap());
463            let cid_bytes = cid.to_bytes();
464            append_varint(cid_bytes.len() + data.len(), &mut sections);
465            sections.extend(cid_bytes);
466            sections.extend(data);
467            cids.push(cid);
468        }
469        let mut root = vec![0];
470        root.extend(cids.last().unwrap().to_bytes());
471        let header = serde_cbor::to_vec(&Map([
472            (Text("roots".into()), Array(vec![Bytes(root)])),
473            (Text("version".into()), Integer(1)),
474        ]
475        .into()))
476        .unwrap();
477        let mut car = Vec::new();
478        append_varint(header.len(), &mut car);
479        car.extend_from_slice(&header);
480        car.extend(sections);
481        (car, cids, header)
482    }
483
484    #[tokio::test]
485    async fn test_async_node_reader() {
486        // A payload larger than 127 bytes exercises a multi-byte section length.
487        let payload = vec![42; 256];
488        let block = |slot| {
489            Array(vec![
490                Integer(2),
491                Integer(slot),
492                Array(vec![]),
493                Array(vec![]),
494                Array(vec![Integer(slot - 1), Integer(1_700_000_000), Integer(7)]),
495                Null,
496            ])
497        };
498        let (car, cids, header) = archive(&[
499            Array(vec![Integer(6), Null, Null, Null, Bytes(payload.clone())]),
500            block(42),
501            block(43),
502        ]);
503        let mut reader = NodeReader::new(io::Cursor::new(car));
504
505        assert_eq!(reader.read_raw_header().await.unwrap(), header);
506        let section_start = reader.reader.position();
507        assert_eq!(reader.read_raw_header().await.unwrap(), header);
508        assert_eq!(reader.reader.position(), section_start);
509
510        let first = reader.read_until_block().await.unwrap();
511        assert_eq!(first.get_cids(), cids[..2]);
512        assert_eq!(
513            first
514                .get(0)
515                .get_node()
516                .get_dataframe()
517                .unwrap()
518                .data
519                .as_slice(),
520            payload,
521        );
522        let first_block = first.get_block().unwrap();
523        assert_eq!(first_block.slot, 42);
524        assert_eq!(first_block.meta.parent_slot, 41);
525        assert_eq!(first_block.meta.blocktime, 1_700_000_000);
526
527        let second = reader.read_until_block().await.unwrap();
528        assert_eq!(second.get_cids(), cids[2..]);
529        assert_eq!(second.get_block().unwrap().slot, 43);
530        assert_eq!(reader.get_item_index(), 3);
531        assert!(reader.read_until_block().await.unwrap().is_empty());
532    }
533
534    #[tokio::test]
535    async fn read_until_block_propagates_invalid_node() {
536        let (car, _, _) = archive(&[Array(vec![Integer(999)])]);
537        let mut reader = NodeReader::new(io::Cursor::new(car));
538        let error = reader.read_until_block().await.err().unwrap();
539        assert_eq!(
540            error.downcast_ref::<io::Error>().unwrap().kind(),
541            io::ErrorKind::Other,
542        );
543    }
544}