sol-shred-sdk 4.0.0

Solana raw shred and ShredStream decoding SDK with multi-DEX event parsing.
use std::net::UdpSocket as StdUdpSocket;
use std::time::{Duration, Instant};

use socket2::{Domain, Protocol, Socket, Type};
use tokio::net::UdpSocket;
use tokio::time::{interval_at, MissedTickBehavior};

use super::config::RawShredConfig;
use super::decoder::{entries_to_tx_batch, ShredEntryBatch, ShredTxBatch};
use super::error::ShredResult;
use super::reassembler::{RawShredDecoder, ShredDecoderStats};

const MAX_DATAGRAMS_PER_POLL: usize = 64;

/// Async UDP client for raw Solana shreds.
pub struct RawShredClient {
    socket: UdpSocket,
    decoder: RawShredDecoder,
    config: RawShredConfig,
}

impl RawShredClient {
    pub async fn bind(config: RawShredConfig) -> ShredResult<Self> {
        let socket = bind_udp_socket(&config)?;
        let decoder = RawShredDecoder::new(config.clone());

        Ok(Self {
            socket,
            decoder,
            config,
        })
    }

    #[inline]
    pub fn stats(&self) -> ShredDecoderStats {
        self.decoder.stats()
    }

    #[inline]
    pub fn decoder_mut(&mut self) -> &mut RawShredDecoder {
        &mut self.decoder
    }

    /// Run the receive loop and callback on each completed Entry batch.
    pub async fn run_entries<F>(&mut self, mut callback: F) -> ShredResult<()>
    where
        F: FnMut(ShredEntryBatch) + Send,
    {
        let mut buf = vec![0u8; self.config.max_datagram_size.max(1280)];
        let eviction_period = self
            .config
            .reassembly_gap_timeout
            .max(Duration::from_millis(1));
        let mut eviction_interval = interval_at(
            tokio::time::Instant::now() + eviction_period,
            eviction_period,
        );
        eviction_interval.set_missed_tick_behavior(MissedTickBehavior::Skip);

        loop {
            tokio::select! {
                ready = self.socket.readable() => {
                    ready?;
                    let batch_now = Instant::now();
                    for _ in 0..MAX_DATAGRAMS_PER_POLL {
                        match self.socket.try_recv(&mut buf) {
                            Ok(n) => {
                                for batch in self.decoder.push_packet(&buf[..n], batch_now) {
                                    callback(batch);
                                }
                            }
                            Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => break,
                            Err(error) => return Err(error.into()),
                        }
                    }
                    tokio::task::yield_now().await;
                }
                _ = eviction_interval.tick() => {
                    self.decoder.evict_stale_slots(Instant::now());
                }
            }
        }
    }

    /// Run the receive loop and callback on each completed transaction batch.
    pub async fn run_transactions<F>(&mut self, mut callback: F) -> ShredResult<()>
    where
        F: FnMut(ShredTxBatch) + Send,
    {
        self.run_entries(|batch| callback(entries_to_tx_batch(batch)))
            .await
    }
}

fn bind_udp_socket(config: &RawShredConfig) -> ShredResult<UdpSocket> {
    let socket = Socket::new(
        Domain::for_address(config.udp_bind),
        Type::DGRAM,
        Some(Protocol::UDP),
    )?;
    socket.set_reuse_address(true)?;
    if config.udp_recv_buffer_bytes > 0 {
        socket.set_recv_buffer_size(config.udp_recv_buffer_bytes)?;
    }
    socket.set_nonblocking(true)?;
    socket.bind(&config.udp_bind.into())?;

    let std_socket: StdUdpSocket = socket.into();
    Ok(UdpSocket::from_std(std_socket)?)
}