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> }, } /// 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, events: broadcast::Sender, } 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 { 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 { self.events.subscribe() } } async fn run_actor(mut rx: mpsc::Receiver, events: broadcast::Sender) { 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"); } } } } }