Skip to main content

yellowstone_block_machine/
stream.rs

1use {
2    crate::{
3        event::{GeyserEventAdapter, GeyserEventInfo},
4        state_machine::{
5            BlockStateMachineOutput, BlockstoreStats, DeadBlockDetected, DeadletterEvent,
6            ForkDetected, FrozenBlock, SlotCommitmentStatusUpdate,
7        },
8        wrapper::BlocksStateMachineWrapper,
9    },
10    derive_more::From,
11    futures_util::{Stream, TryStream, TryStreamExt},
12    rustc_hash::FxHashMap,
13    solana_clock::Slot,
14    solana_commitment_config::CommitmentLevel,
15    solana_hash::HASH_BYTES,
16    std::{cmp::Ordering, collections::VecDeque, marker::PhantomData},
17};
18
19///
20/// A fully reconstructed block, containing all events (accounts, transactions, entries) for a given slot.
21///
22#[derive(Debug, Clone)]
23pub struct Block<Storage> {
24    pub slot: Slot,
25    pub blockhash: [u8; HASH_BYTES],
26    pub events: Storage,
27}
28
29impl<Storage> AsRef<Storage> for Block<Storage> {
30    fn as_ref(&self) -> &Storage {
31        &self.events
32    }
33}
34
35///
36/// A trait for types that can store events for a block, and provide iterators over those events.
37///
38pub trait BlockEventStore {
39    type EventT;
40
41    type Iter<'a>: Iterator<Item = &'a Self::EventT>
42    where
43        Self: 'a,
44        Self::EventT: 'a;
45
46    type IntoIter: IntoIterator<Item = Self::EventT>;
47
48    fn len(&self) -> usize;
49
50    fn is_empty(&self) -> bool {
51        self.len() == 0
52    }
53
54    fn iter(&self) -> Self::Iter<'_>;
55
56    ///
57    /// Returns an iterator over the events in this block that are accounts.
58    fn account_iter(&self) -> Self::Iter<'_>;
59
60    ///
61    /// Returns the number of account events in this block.
62    fn account_len(&self) -> usize {
63        self.account_iter().count()
64    }
65
66    ///
67    /// Returns an iterator over the events in this block that are transactions.
68    fn transaction_iter(&self) -> Self::Iter<'_>;
69
70    ///
71    /// Returns the number of transaction events in this block.
72    fn transaction_len(&self) -> usize {
73        self.transaction_iter().count()
74    }
75
76    ///
77    /// Returns an iterator over the events in this block that are entries.
78    fn entry_iter(&self) -> Self::Iter<'_>;
79
80    ///
81    /// Returns the number of entry events in this block.
82    fn entry_len(&self) -> usize {
83        self.entry_iter().count()
84    }
85
86    ///
87    /// Returns an iterator over the events in this block that are neither accounts, transactions, nor entries.
88    fn other_iter(&self) -> Self::Iter<'_>;
89
90    ///
91    /// Returns the number of events in this block that are neither accounts, transactions, nor entries.
92    fn other_len(&self) -> usize {
93        self.other_iter().count()
94    }
95
96    fn into_iter(self) -> Self::IntoIter;
97}
98
99///
100/// A trait for types that can accumulate events into blocks.
101///
102pub trait BlockAccumulator {
103    ///
104    /// The type of events that this cumulator can handle. This is typically the same as the `EventT` associated type of the `GeyserEventAdapter` used by the `BlockStream`.
105    type EventT;
106
107    ///
108    /// The type of storage used to hold the events for a block. This is typically a `Vec<EventT>`, but can be any type that can hold the events for a block.
109    type EventStore: BlockEventStore;
110
111    ///
112    /// Inserts a new event into the block accumulator for the given slot, under the given
113    /// [`Bucket`].
114    fn add_event(&mut self, event: Self::EventT, slot: Slot, ev_info: &GeyserEventInfo);
115
116    ///
117    /// Marks a block as frozen, indicating that it has been fully reconstructed and is ready for processing.
118    ///
119    /// See [`BlockAccumulator::finish_block`] for how to retrieve the frozen block.
120    fn freeze_block(&mut self, frozen_block_info: FrozenBlock);
121
122    ///
123    /// Finishes a block and returns it, if it exists. This is typically called when the block has been fully processed and is ready to be consumed.
124    ///
125    /// # Note
126    ///
127    /// This function should only return Some if the block was previously `freeze_block`.
128    ///
129    /// # Idempotency
130    ///
131    /// This function is NOT idempotent. Calling it multiple times for the same slot will return None after the first call.
132    fn finish_block(&mut self, slot: Slot) -> Option<Block<Self::EventStore>>;
133
134    ///
135    /// Prunes a block from the accumulator, removing all associated events and data for the given slot.
136    /// This is typically called when a block is no longer needed, such as when it has been finalized or when a fork has been detected.
137    ///
138    fn prune_block(&mut self, slot: Slot);
139}
140
141impl<Storage> Block<Storage> {
142    // ///
143    // /// Returns the number of transactions in this block.
144    // ///
145    // pub fn txn_len(&self) -> usize {
146    //     self.transaction_idx_map.len()
147    // }
148
149    // ///
150    // /// Returns the number of accounts in this block.
151    // ///
152    // pub fn account_len(&self) -> usize {
153    //     self.account_idx_map.len()
154    // }
155
156    // ///
157    // /// Returns the number of entries in this block.
158    // ///
159    // pub fn entry_len(&self) -> usize {
160    //     self.entry_idx_map.len()
161    // }
162
163    // ///
164    // /// Checks if the block has no events.
165    // ///
166    // pub fn is_empty(&self) -> bool {
167    //     self.events.is_empty()
168    // }
169
170    // ///
171    // /// Returns the number of events in this block.
172    // ///
173    // pub fn len(&self) -> usize {
174    //     self.events.len()
175    // }
176}
177
178enum PendingEvent {
179    FrozenBlock(Slot),
180    SlotCommitmentUpdate(SlotCommitmentStatusUpdate),
181    ForkDetected(ForkDetected),
182    DeadBlockDetect(DeadBlockDetected),
183}
184
185///
186/// The different types of outputs produced by the Dragon's mouth block machine.
187///
188#[derive(Debug, From)]
189pub enum BlockMachineOutput<EventStore> {
190    ///
191    /// A fully reconstructed block, ready for processing.
192    ///
193    FrozenBlock(Block<EventStore>),
194    ///
195    /// An update on the commitment status of a slot.
196    /// Note: This is sent when the slot reaches or exceeds the minimum commitment level set during initialization.
197    /// It is guaranteed that the block for this slot has been sent before this update.
198    ///
199    SlotCommitmentUpdate(SlotCommitmentStatusUpdate),
200    ///
201    /// A notification that a fork has been detected.
202    ///
203    ForkDetected(ForkDetected),
204    ///
205    /// A notification that a dead block has been detected.
206    /// Note: All Dead blocks are Forks, but not all Forks are Dead blocks.
207    /// Dead blocks mostly come from corrupted entries early in the replay process of a slot.
208    ///
209    DeadBlockDetected(DeadBlockDetected),
210}
211
212///
213/// A stream that yields [`BlockMachineOutput`] items.
214///
215/// # Generic Parameters
216///
217/// - `Source`: The underlying source of raw Geyser events, typically a gRPC stream from the Geyser
218///   plugin. Its `Ok` item type must match `V::EventT`.
219/// - `V`: A [`GeyserEventAdapter`] that knows how to view the events yielded by `Source`. Use
220///   `yellowstone_grpc_proto::geyser::SubscribeUpdate` (behind the `dragonsmouth-thin` feature,
221///   which implements this trait on itself) or implement [`GeyserEventAdapter`] on your own type to
222///   avoid depending on a specific version of `yellowstone-grpc-proto`.
223///
224pub struct BlockStream<Source, Adaptor, Acc> {
225    min_commitment_level: CommitmentLevel,
226    source: Source,
227    machine: BlocksStateMachineWrapper,
228    storage: Acc,
229    pending: VecDeque<PendingEvent>,
230    _adapter: PhantomData<Adaptor>,
231}
232
233impl<Source, Adaptor, Acc> BlockStream<Source, Adaptor, Acc>
234where
235    Adaptor: GeyserEventAdapter,
236{
237    pub fn new(source: Source, block_acc: Acc, min_commitment_level: CommitmentLevel) -> Self {
238        Self {
239            min_commitment_level,
240            source,
241            machine: BlocksStateMachineWrapper::new_with_slot_gc_tracing(),
242            storage: block_acc,
243            pending: VecDeque::new(),
244            _adapter: PhantomData,
245        }
246    }
247}
248
249// Auto-derivation of `Unpin` doesn't see through the `Adaptor::EventT` associated-type projection
250// held (transitively) by `pending`, so it's implemented explicitly here instead.
251impl<Source, Adaptor, Acc> Unpin for BlockStream<Source, Adaptor, Acc>
252where
253    Source: Unpin,
254    Adaptor: GeyserEventAdapter,
255    Adaptor::EventT: Unpin,
256    Acc: Unpin,
257{
258}
259
260fn compare_commitment(cl1: CommitmentLevel, cl2: CommitmentLevel) -> Ordering {
261    match (cl1, cl2) {
262        (CommitmentLevel::Processed, CommitmentLevel::Processed) => Ordering::Equal,
263        (CommitmentLevel::Confirmed, CommitmentLevel::Confirmed) => Ordering::Equal,
264        (CommitmentLevel::Finalized, CommitmentLevel::Finalized) => Ordering::Equal,
265        (CommitmentLevel::Processed, _) => Ordering::Less,
266        (CommitmentLevel::Confirmed, CommitmentLevel::Processed) => Ordering::Greater,
267        (CommitmentLevel::Finalized, CommitmentLevel::Processed) => Ordering::Greater,
268        (CommitmentLevel::Finalized, CommitmentLevel::Confirmed) => Ordering::Greater,
269        (CommitmentLevel::Confirmed, CommitmentLevel::Finalized) => Ordering::Less,
270    }
271}
272
273impl<Source, Adaptor, Acc> BlockStream<Source, Adaptor, Acc>
274where
275    Adaptor: GeyserEventAdapter,
276    Acc: BlockAccumulator<EventT = Adaptor::EventT>,
277{
278    pub fn state_machine_stats(&self) -> BlockstoreStats {
279        self.machine.sm.stats()
280    }
281
282    fn insert_into_storage(&mut self, event: Adaptor::EventT, ev_info: &GeyserEventInfo) {
283        let slot = ev_info.slot();
284        self.storage.add_event(event, slot, ev_info);
285    }
286
287    fn on_new_frozen_block(&mut self) {
288        // Drain DLQ — clean up slots the state machine gave up on
289        while let Some(dlq_event) = self.machine.pop_next_dlq() {
290            match dlq_event {
291                DeadletterEvent::Incomplete(slot) => {
292                    self.storage.prune_block(slot);
293                }
294            }
295        }
296
297        while let Some(slot) = self.machine.pop_slot_gc_trace() {
298            self.storage.prune_block(slot);
299        }
300    }
301
302    fn process_state_machine_output(&mut self) {
303        while let Some(output) = self.machine.pop_next_state_machine_output() {
304            match output {
305                BlockStateMachineOutput::FrozenBlock(frozen_block) => {
306                    self.on_new_frozen_block();
307                    self.storage.freeze_block(frozen_block);
308                }
309                BlockStateMachineOutput::SlotStatus(slot_status) => {
310                    let slot = slot_status.slot;
311                    let cl = slot_status.commitment;
312                    match compare_commitment(cl, self.min_commitment_level) {
313                        Ordering::Less => continue,
314                        _ => {
315                            let commitment_level_update = SlotCommitmentStatusUpdate {
316                                parent_slot: slot_status.parent_slot,
317                                slot: slot_status.slot,
318                                commitment: cl,
319                            };
320
321                            self.pending.push_back(PendingEvent::FrozenBlock(slot));
322
323                            self.pending.push_back(PendingEvent::SlotCommitmentUpdate(
324                                commitment_level_update,
325                            ));
326                        }
327                    }
328                }
329                BlockStateMachineOutput::ForksDetected(fork_detected) => {
330                    self.storage.prune_block(fork_detected.slot);
331                    self.pending
332                        .push_back(PendingEvent::ForkDetected(fork_detected));
333                }
334                BlockStateMachineOutput::DeadSlotDetected(dead_block) => {
335                    self.storage.prune_block(dead_block.slot);
336                    self.pending
337                        .push_back(PendingEvent::DeadBlockDetect(dead_block));
338                }
339                BlockStateMachineOutput::BankCreated(_) => {}
340                BlockStateMachineOutput::BankReset(slot) => {
341                    self.storage.prune_block(slot);
342                }
343            }
344        }
345    }
346}
347
348impl<Source, Adaptor, Acc> Stream for BlockStream<Source, Adaptor, Acc>
349where
350    Source: TryStream<Ok = Adaptor::EventT> + Unpin,
351    Adaptor: GeyserEventAdapter,
352    Adaptor::EventT: Unpin,
353    Acc: BlockAccumulator<EventT = Adaptor::EventT> + Unpin,
354{
355    type Item = Result<BlockMachineOutput<Acc::EventStore>, Source::Error>;
356
357    fn poll_next(
358        mut self: std::pin::Pin<&mut Self>,
359        cx: &mut std::task::Context<'_>,
360    ) -> std::task::Poll<Option<Self::Item>> {
361        loop {
362            if let Some(pending_ev) = self.pending.pop_front() {
363                let output = match pending_ev {
364                    PendingEvent::FrozenBlock(slot) => {
365                        if let Some(block) = self.storage.finish_block(slot) {
366                            BlockMachineOutput::FrozenBlock(block)
367                        } else {
368                            continue;
369                        }
370                    }
371                    PendingEvent::SlotCommitmentUpdate(update) => {
372                        BlockMachineOutput::SlotCommitmentUpdate(update)
373                    }
374                    PendingEvent::ForkDetected(fork) => BlockMachineOutput::ForkDetected(fork),
375                    PendingEvent::DeadBlockDetect(dead) => {
376                        BlockMachineOutput::DeadBlockDetected(dead)
377                    }
378                };
379
380                return std::task::Poll::Ready(Some(Ok(output)));
381            }
382
383            match self.source.try_poll_next_unpin(cx) {
384                std::task::Poll::Ready(Some(Ok(ev))) => {
385                    let event_view = match Adaptor::extract_geyser_ev_info(&ev) {
386                        Some(ev) => ev,
387                        None => continue,
388                    };
389
390                    if self
391                        .machine
392                        .handle_new_geyser_event(event_view.clone())
393                        .is_ok()
394                    {
395                        self.insert_into_storage(ev, &event_view);
396                    }
397                }
398                std::task::Poll::Ready(Some(Err(e))) => {
399                    return std::task::Poll::Ready(Some(Err(e)));
400                }
401                std::task::Poll::Ready(None) => {
402                    return std::task::Poll::Ready(None);
403                }
404                std::task::Poll::Pending => {
405                    return std::task::Poll::Pending;
406                }
407            }
408            self.process_state_machine_output();
409        }
410    }
411}
412
413///
414/// Which per-block index map an event's position should be recorded in.
415///
416#[derive(Debug, Clone, Copy, PartialEq, Eq)]
417pub enum Bucket {
418    Account,
419    Transaction,
420    Entry,
421    Other,
422}
423
424#[derive(Debug)]
425struct BlockBuffer<E> {
426    blockhash: [u8; HASH_BYTES],
427    events: Vec<E>,
428    account_idx_map: Vec<usize>,
429    transaction_idx_map: Vec<usize>,
430    entry_idx_map: Vec<usize>,
431    other_idx_map: Vec<usize>,
432}
433
434impl<E> Default for BlockBuffer<E> {
435    fn default() -> Self {
436        Self {
437            blockhash: [0; HASH_BYTES],
438            events: Vec::new(),
439            account_idx_map: Vec::new(),
440            transaction_idx_map: Vec::new(),
441            entry_idx_map: Vec::new(),
442            other_idx_map: Vec::new(),
443        }
444    }
445}
446
447pub struct SimpleBlockStore<E> {
448    pub slot: Slot,
449    pub events: Vec<E>,
450    pub account_idx_map: Vec<usize>,
451    pub transaction_idx_map: Vec<usize>,
452    pub entry_idx_map: Vec<usize>,
453    pub other_idx_map: Vec<usize>,
454}
455
456pub struct SimpleBlockStoreIter<'a, E> {
457    events: &'a [E],
458    idx_map: Option<&'a [usize]>,
459    idx_pos: usize,
460}
461
462impl<'a, E> Iterator for SimpleBlockStoreIter<'a, E> {
463    type Item = &'a E;
464
465    fn next(&mut self) -> Option<Self::Item> {
466        if let Some(idx_map) = self.idx_map {
467            if self.idx_pos < idx_map.len() {
468                let idx = idx_map[self.idx_pos];
469                self.idx_pos += 1;
470                self.events.get(idx)
471            } else {
472                None
473            }
474        } else {
475            if self.idx_pos < self.events.len() {
476                let event = &self.events[self.idx_pos];
477                self.idx_pos += 1;
478                Some(event)
479            } else {
480                None
481            }
482        }
483    }
484}
485
486impl<E> BlockEventStore for SimpleBlockStore<E> {
487    type EventT = E;
488    type Iter<'a>
489        = SimpleBlockStoreIter<'a, E>
490    where
491        Self: 'a,
492        E: 'a;
493
494    type IntoIter = std::vec::IntoIter<E>;
495
496    fn len(&self) -> usize {
497        self.events.len()
498    }
499
500    fn iter(&self) -> Self::Iter<'_> {
501        SimpleBlockStoreIter {
502            events: &self.events,
503            idx_map: None,
504            idx_pos: 0,
505        }
506    }
507
508    fn account_iter(&self) -> Self::Iter<'_> {
509        SimpleBlockStoreIter {
510            events: &self.events,
511            idx_map: Some(&self.account_idx_map),
512            idx_pos: 0,
513        }
514    }
515
516    fn transaction_iter(&self) -> Self::Iter<'_> {
517        SimpleBlockStoreIter {
518            events: &self.events,
519            idx_map: Some(&self.transaction_idx_map),
520            idx_pos: 0,
521        }
522    }
523
524    fn entry_iter(&self) -> Self::Iter<'_> {
525        SimpleBlockStoreIter {
526            events: &self.events,
527            idx_map: Some(&self.entry_idx_map),
528            idx_pos: 0,
529        }
530    }
531
532    fn other_iter(&self) -> Self::Iter<'_> {
533        SimpleBlockStoreIter {
534            events: &self.events,
535            idx_map: Some(&self.other_idx_map),
536            idx_pos: 0,
537        }
538    }
539
540    fn account_len(&self) -> usize {
541        self.account_idx_map.len()
542    }
543
544    fn transaction_len(&self) -> usize {
545        self.transaction_idx_map.len()
546    }
547
548    fn entry_len(&self) -> usize {
549        self.entry_idx_map.len()
550    }
551
552    fn other_len(&self) -> usize {
553        self.other_idx_map.len()
554    }
555
556    fn into_iter(self) -> Self::IntoIter {
557        self.events.into_iter()
558    }
559}
560
561impl<E> BlockBuffer<E> {
562    fn finish(self, slot: Slot) -> Block<SimpleBlockStore<E>> {
563        Block {
564            slot,
565            blockhash: self.blockhash,
566            events: SimpleBlockStore {
567                slot,
568                events: self.events,
569                account_idx_map: self.account_idx_map,
570                transaction_idx_map: self.transaction_idx_map,
571                entry_idx_map: self.entry_idx_map,
572                other_idx_map: self.other_idx_map,
573            },
574        }
575    }
576}
577
578///
579/// An in-memory store for blocks being reconstructed.
580///
581/// It maintains active blocks (currently being reconstructed) and frozen blocks (fully reconstructed).
582pub struct SimpleBlockAccumulator<E> {
583    active_block_map: FxHashMap<Slot, BlockBuffer<E>>,
584    frozen_block_map: FxHashMap<Slot, BlockBuffer<E>>,
585}
586
587impl<E> Default for SimpleBlockAccumulator<E> {
588    fn default() -> Self {
589        Self {
590            active_block_map: FxHashMap::default(),
591            frozen_block_map: FxHashMap::default(),
592        }
593    }
594}
595
596impl<E> BlockAccumulator for SimpleBlockAccumulator<E> {
597    type EventT = E;
598    type EventStore = SimpleBlockStore<E>;
599
600    fn add_event(&mut self, event: E, slot: Slot, ev_info: &GeyserEventInfo) {
601        let block = self.active_block_map.entry(slot).or_default();
602        let idx = block.events.len();
603        match ev_info {
604            GeyserEventInfo::Account { .. } => block.account_idx_map.push(idx),
605            GeyserEventInfo::Transaction { .. } => block.transaction_idx_map.push(idx),
606            GeyserEventInfo::Entry(_) => block.entry_idx_map.push(idx),
607            GeyserEventInfo::Other { .. } => block.other_idx_map.push(idx),
608            _ => {
609                //block meta and slot are ignored
610                return;
611            }
612        }
613        block.events.push(event);
614    }
615
616    fn freeze_block(&mut self, frozen_block_info: FrozenBlock) {
617        let Some(mut block) = self.active_block_map.remove(&frozen_block_info.slot) else {
618            return;
619        };
620        block.blockhash = frozen_block_info.blockhash.to_bytes();
621        self.frozen_block_map.insert(frozen_block_info.slot, block);
622    }
623
624    fn finish_block(&mut self, slot: Slot) -> Option<Block<SimpleBlockStore<E>>> {
625        let acc = self.frozen_block_map.remove(&slot)?;
626        Some(acc.finish(slot))
627    }
628
629    fn prune_block(&mut self, slot: Slot) {
630        self.active_block_map.remove(&slot);
631        self.frozen_block_map.remove(&slot);
632    }
633}
634
635#[cfg(all(test, feature = "dragonsmouth-thin"))]
636mod tests {
637    use {
638        super::{
639            BlockEventStore, BlockMachineOutput, BlockStream, PendingEvent, SimpleBlockAccumulator,
640            SimpleBlockStore,
641        },
642        crate::{event::GeyserEventAdapter, state_machine::SlotCommitmentStatusUpdate},
643        futures_util::{Stream, stream},
644        solana_commitment_config::CommitmentLevel,
645        solana_hash::Hash,
646        std::{
647            io,
648            pin::Pin,
649            task::{Context, Poll},
650        },
651        yellowstone_grpc_proto::geyser::{
652            SlotStatus, SubscribeUpdate, SubscribeUpdateAccount, SubscribeUpdateBlockMeta,
653            SubscribeUpdateEntry, SubscribeUpdateSlot, SubscribeUpdateTransaction,
654            subscribe_update::UpdateOneof,
655        },
656    };
657
658    fn update(oneof: UpdateOneof, filters: Vec<String>) -> SubscribeUpdate {
659        SubscribeUpdate {
660            filters,
661            created_at: None,
662            update_oneof: Some(oneof),
663        }
664    }
665
666    fn slot_update(slot: u64, parent: Option<u64>, status: SlotStatus) -> SubscribeUpdate {
667        update(
668            UpdateOneof::Slot(SubscribeUpdateSlot {
669                slot,
670                parent,
671                status: status as i32,
672                dead_error: None,
673            }),
674            vec!["test".to_string()],
675        )
676    }
677
678    fn entry_update(slot: u64, index: u64) -> SubscribeUpdate {
679        update(
680            UpdateOneof::Entry(SubscribeUpdateEntry {
681                slot,
682                index,
683                num_hashes: 0,
684                hash: Hash::new_unique().to_bytes().to_vec(),
685                executed_transaction_count: 1,
686                starting_transaction_index: index,
687            }),
688            vec!["client-filter".to_string()],
689        )
690    }
691
692    fn tx_update(slot: u64) -> SubscribeUpdate {
693        update(
694            UpdateOneof::Transaction(SubscribeUpdateTransaction {
695                slot,
696                ..Default::default()
697            }),
698            vec!["client-filter".to_string()],
699        )
700    }
701
702    fn account_update(slot: u64) -> SubscribeUpdate {
703        update(
704            UpdateOneof::Account(SubscribeUpdateAccount {
705                slot,
706                ..Default::default()
707            }),
708            vec!["client-filter".to_string()],
709        )
710    }
711
712    fn block_meta_update(slot: u64, parent_slot: u64, entries_count: u64) -> SubscribeUpdate {
713        let blockhash = bs58::encode(Hash::new_unique().to_bytes()).into_string();
714        update(
715            UpdateOneof::BlockMeta(SubscribeUpdateBlockMeta {
716                slot,
717                parent_slot,
718                blockhash,
719                executed_transaction_count: entries_count,
720                entries_count,
721                ..Default::default()
722            }),
723            vec!["test".to_string()],
724        )
725    }
726
727    fn feed(
728        stream: &mut BlockStream<
729            stream::Iter<std::vec::IntoIter<Result<SubscribeUpdate, io::Error>>>,
730            SubscribeUpdate,
731            SimpleBlockAccumulator<SubscribeUpdate>,
732        >,
733        ev: SubscribeUpdate,
734    ) {
735        let ev_info = SubscribeUpdate::extract_geyser_ev_info(&ev).unwrap();
736        if stream
737            .machine
738            .handle_new_geyser_event(ev_info.clone())
739            .is_ok()
740        {
741            stream.insert_into_storage(ev, &ev_info);
742        }
743        stream.process_state_machine_output();
744    }
745
746    fn empty_source_stream(
747        min_commitment_level: CommitmentLevel,
748    ) -> BlockStream<
749        stream::Iter<std::vec::IntoIter<Result<SubscribeUpdate, io::Error>>>,
750        SubscribeUpdate,
751        SimpleBlockAccumulator<SubscribeUpdate>,
752    > {
753        BlockStream::new(
754            stream::iter(Vec::<Result<SubscribeUpdate, io::Error>>::new()),
755            SimpleBlockAccumulator::default(),
756            min_commitment_level,
757        )
758    }
759
760    #[test]
761    fn emits_frozen_block_before_slot_commitment_update() {
762        let mut bs = empty_source_stream(CommitmentLevel::Processed);
763
764        feed(
765            &mut bs,
766            slot_update(10, Some(9), SlotStatus::SlotFirstShredReceived),
767        );
768        feed(&mut bs, slot_update(10, Some(9), SlotStatus::SlotCompleted));
769        feed(&mut bs, entry_update(10, 0));
770        feed(&mut bs, tx_update(10));
771        feed(&mut bs, account_update(10));
772        feed(&mut bs, block_meta_update(10, 9, 1));
773        feed(&mut bs, slot_update(10, Some(9), SlotStatus::SlotProcessed));
774
775        let waker = futures_util::task::noop_waker();
776        let mut cx = Context::from_waker(&waker);
777
778        let first = Pin::new(&mut bs).poll_next(&mut cx);
779        let second = Pin::new(&mut bs).poll_next(&mut cx);
780
781        let Poll::Ready(Some(Ok(BlockMachineOutput::FrozenBlock(block)))) = first else {
782            panic!("expected FrozenBlock first");
783        };
784        assert_eq!(block.slot, 10);
785        assert_eq!(block.events.entry_idx_map.len(), 1);
786        assert_eq!(block.events.transaction_idx_map.len(), 1);
787        assert_eq!(block.events.account_idx_map.len(), 1);
788        assert_eq!(block.events.events.len(), 3);
789
790        let Poll::Ready(Some(Ok(BlockMachineOutput::SlotCommitmentUpdate(update)))) = second else {
791            panic!("expected SlotCommitmentUpdate second");
792        };
793        assert_eq!(update.slot, 10);
794        assert_eq!(update.commitment, CommitmentLevel::Processed);
795    }
796
797    #[test]
798    fn respects_minimum_commitment_filter() {
799        let mut bs = empty_source_stream(CommitmentLevel::Confirmed);
800
801        feed(
802            &mut bs,
803            slot_update(42, Some(41), SlotStatus::SlotFirstShredReceived),
804        );
805        feed(
806            &mut bs,
807            slot_update(42, Some(41), SlotStatus::SlotCompleted),
808        );
809        feed(&mut bs, entry_update(42, 0));
810        feed(&mut bs, block_meta_update(42, 41, 1));
811
812        // Processed is below minimum commitment and should produce no output.
813        feed(
814            &mut bs,
815            slot_update(42, Some(41), SlotStatus::SlotProcessed),
816        );
817
818        let waker = futures_util::task::noop_waker();
819        let mut cx = Context::from_waker(&waker);
820        let none_after_processed = Pin::new(&mut bs).poll_next(&mut cx);
821        assert!(matches!(none_after_processed, Poll::Ready(None)));
822
823        // Confirmed reaches minimum commitment and should emit both block and commitment update.
824        feed(
825            &mut bs,
826            slot_update(42, Some(41), SlotStatus::SlotConfirmed),
827        );
828
829        let first = Pin::new(&mut bs).poll_next(&mut cx);
830        assert!(matches!(
831            first,
832            Poll::Ready(Some(Ok(BlockMachineOutput::FrozenBlock(_))))
833        ));
834
835        let second = Pin::new(&mut bs).poll_next(&mut cx);
836        assert!(matches!(
837            second,
838            Poll::Ready(Some(Ok(BlockMachineOutput::SlotCommitmentUpdate(_))))
839        ));
840
841        let third = Pin::new(&mut bs).poll_next(&mut cx);
842        assert!(matches!(third, Poll::Ready(None)));
843    }
844
845    #[test]
846    fn stream_forwards_source_error_and_end_of_stream() {
847        let source = stream::iter(vec![Err::<SubscribeUpdate, _>(io::Error::other("boom"))]);
848        let mut bs =
849            BlockStream::<_, SubscribeUpdate, SimpleBlockAccumulator<SubscribeUpdate>>::new(
850                source,
851                SimpleBlockAccumulator::default(),
852                CommitmentLevel::Processed,
853            );
854        let waker = futures_util::task::noop_waker();
855        let mut cx = Context::from_waker(&waker);
856
857        let first = Pin::new(&mut bs).poll_next(&mut cx);
858        assert!(matches!(first, Poll::Ready(Some(Err(_)))));
859
860        let source = stream::iter(Vec::<Result<SubscribeUpdate, io::Error>>::new());
861        let mut bs =
862            BlockStream::<_, SubscribeUpdate, SimpleBlockAccumulator<SubscribeUpdate>>::new(
863                source,
864                SimpleBlockAccumulator::default(),
865                CommitmentLevel::Processed,
866            );
867        let second = Pin::new(&mut bs).poll_next(&mut cx);
868        assert!(matches!(second, Poll::Ready(None)));
869    }
870
871    #[test]
872    fn simple_block_store_empty_iterators_are_empty() {
873        let store = SimpleBlockStore::<u64> {
874            slot: 99,
875            events: Vec::new(),
876            account_idx_map: Vec::new(),
877            transaction_idx_map: Vec::new(),
878            entry_idx_map: Vec::new(),
879            other_idx_map: Vec::new(),
880        };
881
882        assert!(store.is_empty());
883        assert_eq!(store.len(), 0);
884        assert_eq!(store.account_len(), 0);
885        assert_eq!(store.transaction_len(), 0);
886        assert_eq!(store.entry_len(), 0);
887        assert_eq!(store.other_len(), 0);
888        assert_eq!(store.iter().count(), 0);
889        assert_eq!(store.account_iter().count(), 0);
890        assert_eq!(store.transaction_iter().count(), 0);
891        assert_eq!(store.entry_iter().count(), 0);
892        assert_eq!(store.other_iter().count(), 0);
893    }
894
895    #[test]
896    fn simple_block_store_partition_iterators_return_expected_events() {
897        let store = SimpleBlockStore {
898            slot: 7,
899            events: vec![10_u64, 11, 12, 13, 14],
900            account_idx_map: vec![1, 4],
901            transaction_idx_map: vec![0, 3],
902            entry_idx_map: vec![2],
903            other_idx_map: vec![4],
904        };
905
906        let all: Vec<u64> = store.iter().copied().collect();
907        let accounts: Vec<u64> = store.account_iter().copied().collect();
908        let txs: Vec<u64> = store.transaction_iter().copied().collect();
909        let entries: Vec<u64> = store.entry_iter().copied().collect();
910        let others: Vec<u64> = store.other_iter().copied().collect();
911
912        assert_eq!(all, vec![10, 11, 12, 13, 14]);
913        assert_eq!(accounts, vec![11, 14]);
914        assert_eq!(txs, vec![10, 13]);
915        assert_eq!(entries, vec![12]);
916        assert_eq!(others, vec![14]);
917    }
918
919    #[test]
920    fn skips_missing_frozen_block_and_emits_following_commitment_update() {
921        let mut bs = empty_source_stream(CommitmentLevel::Processed);
922
923        // Simulate a pending FrozenBlock for a slot that no longer exists in storage,
924        // followed by a valid commitment update for the same slot.
925        bs.pending.push_back(PendingEvent::FrozenBlock(77));
926        bs.pending.push_back(PendingEvent::SlotCommitmentUpdate(
927            SlotCommitmentStatusUpdate {
928                parent_slot: Some(76),
929                slot: 77,
930                commitment: CommitmentLevel::Processed,
931            },
932        ));
933
934        let waker = futures_util::task::noop_waker();
935        let mut cx = Context::from_waker(&waker);
936
937        let first = Pin::new(&mut bs).poll_next(&mut cx);
938        assert!(matches!(
939            first,
940            Poll::Ready(Some(Ok(BlockMachineOutput::SlotCommitmentUpdate(_))))
941        ));
942
943        let second = Pin::new(&mut bs).poll_next(&mut cx);
944        assert!(matches!(second, Poll::Ready(None)));
945    }
946}