ferrix/crates/rib/src/adj_rib_in.rs

use std::collections::HashMap;

use bgp_primitives::ip_prefix::IpPrefix;
use tokio::sync::{broadcast, mpsc, oneshot};
use tracing::warn;

use crate::entry::{PeerId, RibEntry, SnapshotEntry};
use crate::event::RibEvent;

const EVENT_CHANNEL_CAPACITY: usize = 1024;
const CMD_CHANNEL_CAPACITY: usize = 256;

enum RibCommand {
    Announce { peer: PeerId, prefix: IpPrefix, entry: RibEntry },
    Withdraw { peer: PeerId, prefix: IpPrefix },
    Snapshot { reply: oneshot::Sender<Vec<SnapshotEntry>> },
}

/// Handle to the Adj-RIB-In actor. Cheap to clone — all clones share the same
/// backing task. The actor is stopped when all handles are dropped.
#[derive(Clone)]
pub struct AdjRibIn {
    tx: mpsc::Sender<RibCommand>,
    events: broadcast::Sender<RibEvent>,
}

impl AdjRibIn {
    /// Spawns the backing actor and returns a handle.
    pub fn spawn() -> Self {
        let (tx, rx) = mpsc::channel(CMD_CHANNEL_CAPACITY);
        let (events, _) = broadcast::channel(EVENT_CHANNEL_CAPACITY);
        let events_clone = events.clone();
        tokio::spawn(run_actor(rx, events_clone));
        Self { tx, events }
    }

    pub async fn announce(&self, peer: PeerId, prefix: IpPrefix, entry: RibEntry) {
        let _ = self.tx.send(RibCommand::Announce { peer, prefix, entry }).await;
    }

    pub async fn withdraw(&self, peer: PeerId, prefix: IpPrefix) {
        let _ = self.tx.send(RibCommand::Withdraw { peer, prefix }).await;
    }

    /// Returns a point-in-time snapshot of all routes in the RIB.
    pub async fn snapshot(&self) -> Vec<SnapshotEntry> {
        let (reply_tx, reply_rx) = oneshot::channel();
        let _ = self.tx.send(RibCommand::Snapshot { reply: reply_tx }).await;
        reply_rx.await.unwrap_or_default()
    }

    /// Subscribe to a stream of RIB mutation events.
    ///
    /// The receiver will see every [`RibEvent`] from the point of subscription.
    /// If the consumer falls behind by more than [`EVENT_CHANNEL_CAPACITY`]
    /// events it will receive a [`broadcast::error::RecvError::Lagged`] error.
    pub fn subscribe(&self) -> broadcast::Receiver<RibEvent> {
        self.events.subscribe()
    }
}

async fn run_actor(mut rx: mpsc::Receiver<RibCommand>, events: broadcast::Sender<RibEvent>) {
    let mut table: HashMap<(PeerId, IpPrefix), RibEntry> = HashMap::new();

    while let Some(cmd) = rx.recv().await {
        match cmd {
            RibCommand::Announce { peer, prefix, entry } => {
                table.insert((peer, prefix.clone()), entry.clone());
                if let Err(e) = events.send(RibEvent::Announced { peer, prefix, entry }) {
                    // No active subscribers — not an error.
                    let _ = e;
                }
            }
            RibCommand::Withdraw { peer, prefix } => {
                if table.remove(&(peer, prefix.clone())).is_some() {
                    if let Err(e) = events.send(RibEvent::Withdrawn { peer, prefix }) {
                        let _ = e;
                    }
                }
            }
            RibCommand::Snapshot { reply } => {
                let snap = table
                    .iter()
                    .map(|((peer, prefix), entry)| SnapshotEntry {
                        peer: *peer,
                        prefix: prefix.clone(),
                        entry: entry.clone(),
                    })
                    .collect();
                if reply.send(snap).is_err() {
                    warn!("Snapshot reply channel closed before send");
                }
            }
        }
    }
}