use anyhow::Result;
use arancini_lib::sender::UpdateSender;
use arancini_lib::update::Update;
use async_nats::jetstream;
use crossfire::mpsc;
use crossfire::{AsyncRxTrait, AsyncTxTrait};
use futures::FutureExt;
use futures::stream::{FuturesUnordered, StreamExt};
use metrics::{counter, gauge, histogram};
use std::future::Future;
use std::net::IpAddr;
use std::pin::Pin;
use std::thread;
use std::time::Instant;
use tokio::sync::mpsc::Sender;
use tracing::{debug, error};
use crate::config::NatsConfig;
use crate::serializer::serialize_update;
type BridgeFlavor = mpsc::Array<Update>;
type BridgeTx = crossfire::MAsyncTx<BridgeFlavor>;
type BridgeRx = crossfire::AsyncRx<BridgeFlavor>;
type AckWaiter = Pin<Box<dyn Future<Output = (Instant, Result<()>)> + Send>>;
#[derive(Clone)]
pub(crate) struct BridgeSender {
tx: BridgeTx,
}
impl BridgeSender {
pub(crate) fn new(tx: BridgeTx) -> Self {
Self { tx }
}
pub(crate) fn queue_capacity(&self) -> Option<usize> {
self.tx.capacity()
}
}
impl UpdateSender for BridgeSender {
fn send<'a>(&'a self, update: Update) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
Box::pin(async move {
self.tx.send(update).await.map_err(|err| {
anyhow::anyhow!("failed to send update to arancini bridge: {}", err)
})?;
if let Some(capacity) = self.tx.capacity() {
gauge!("risotto_arancini_bridge_queue_fill_ratio")
.set(self.tx.len() as f64 / capacity as f64);
}
gauge!("risotto_arancini_bridge_queue_depth").set(self.tx.len() as f64);
counter!("risotto_arancini_bridge_tx_total").increment(1);
Ok(())
})
}
}
pub(crate) fn bounded_channel(size: usize) -> (BridgeSender, BridgeRx) {
let (tx, rx) = mpsc::bounded_async(size);
(BridgeSender::new(tx), rx)
}
pub(crate) async fn sidecar_forwarder(rx: BridgeRx, tx: Sender<Update>) -> Result<()> {
debug!("starting arancini tokio sidecar forwarder");
loop {
let update = match rx.recv().await {
Ok(update) => update,
Err(_) => {
debug!("arancini bridge receiver closed");
break;
}
};
tx.send(update)
.await
.map_err(|err| anyhow::anyhow!("failed forwarding bridge update: {}", err))?;
counter!("risotto_arancini_bridge_rx_total").increment(1);
if let Some(capacity) = rx.capacity() {
gauge!("risotto_arancini_bridge_queue_fill_ratio")
.set(rx.len() as f64 / capacity as f64);
}
gauge!("risotto_arancini_bridge_queue_depth").set(rx.len() as f64);
}
Ok(())
}
fn nats_subject_for_update(base_subject: &str, update: &Update) -> String {
let router_ip = router_ip_subject_token(update.router_addr);
let afi_safi = afi_safi_subject_token(update);
format!(
"{}.{}.{}.{}",
base_subject, router_ip, update.peer_asn, afi_safi
)
}
fn router_ip_subject_token(ip: IpAddr) -> String {
match ip {
IpAddr::V4(v4) => {
let [a, b, c, d] = v4.octets();
format!("v4_{}_{}_{}_{}", a, b, c, d)
}
IpAddr::V6(v6) => {
if let Some(v4_mapped) = v6.to_ipv4_mapped() {
return router_ip_subject_token(IpAddr::V4(v4_mapped));
}
let segments = v6.segments();
format!(
"v6_{:x}_{:x}_{:x}_{:x}_{:x}_{:x}_{:x}_{:x}",
segments[0],
segments[1],
segments[2],
segments[3],
segments[4],
segments[5],
segments[6],
segments[7]
)
}
}
}
fn afi_safi_subject_token(update: &Update) -> String {
let (afi, safi) = if update.announced {
(update.attrs.mp_reach_afi, update.attrs.mp_reach_safi)
} else {
(update.attrs.mp_unreach_afi, update.attrs.mp_unreach_safi)
};
let afi = afi.unwrap_or(match update.prefix_addr {
IpAddr::V4(_) => 1u16,
IpAddr::V6(_) => 2u16,
});
let safi = safi.unwrap_or(1u8);
format!("{}_{}", afi, safi)
}
async fn nats_sidecar_publisher(rx: BridgeRx, cfg: NatsConfig) -> Result<()> {
let client = crate::nats::connect(&cfg).await?;
let jetstream = jetstream::new(client);
let mut in_flight_acks: FuturesUnordered<AckWaiter> = FuturesUnordered::new();
debug!(
"starting arancini NATS sidecar publisher to {} subject {}",
cfg.server, cfg.subject
);
loop {
while let Some((enqueued_at, ack_result)) = in_flight_acks.next().now_or_never().flatten() {
match ack_result {
Ok(_) => {
counter!("risotto_arancini_nats_ack_success_total").increment(1);
histogram!("risotto_arancini_nats_ack_lag_seconds")
.record(enqueued_at.elapsed().as_secs_f64());
}
Err(err) => {
counter!("risotto_arancini_nats_ack_failure_total").increment(1);
error!("JetStream ACK failed: {}", err);
}
}
gauge!("risotto_arancini_nats_ack_inflight").set(in_flight_acks.len() as f64);
}
let update = match rx.recv().await {
Ok(update) => update,
Err(_) => {
debug!("arancini bridge receiver closed");
break;
}
};
counter!("risotto_arancini_bridge_rx_total").increment(1);
if let Some(capacity) = rx.capacity() {
gauge!("risotto_arancini_bridge_queue_fill_ratio")
.set(rx.len() as f64 / capacity as f64);
}
gauge!("risotto_arancini_bridge_queue_depth").set(rx.len() as f64);
let subject = nats_subject_for_update(&cfg.subject, &update);
let payload = serialize_update(&update);
let enqueued_at = Instant::now();
let publish_start = Instant::now();
match jetstream.publish(subject.clone(), payload.into()).await {
Ok(ack) => {
counter!("risotto_arancini_nats_publish_enqueued_total").increment(1);
histogram!("risotto_arancini_nats_publish_enqueue_latency_seconds")
.record(publish_start.elapsed().as_secs_f64());
in_flight_acks.push(Box::pin(async move {
let res = ack
.await
.map(|_| ())
.map_err(|err| anyhow::anyhow!("JetStream publish ACK error: {}", err));
(enqueued_at, res)
}));
gauge!("risotto_arancini_nats_ack_inflight").set(in_flight_acks.len() as f64);
}
Err(err) => {
counter!("risotto_arancini_nats_publish_errors_total").increment(1);
histogram!("risotto_arancini_nats_publish_enqueue_latency_seconds")
.record(publish_start.elapsed().as_secs_f64());
error!(
"JetStream publish enqueue failed for subject {}: {}",
subject, err
);
}
}
while in_flight_acks.len() >= cfg.max_in_flight_acks {
if let Some((enqueued_at, ack_result)) = in_flight_acks.next().await {
match ack_result {
Ok(_) => {
counter!("risotto_arancini_nats_ack_success_total").increment(1);
histogram!("risotto_arancini_nats_ack_lag_seconds")
.record(enqueued_at.elapsed().as_secs_f64());
}
Err(err) => {
counter!("risotto_arancini_nats_ack_failure_total").increment(1);
error!("JetStream ACK failed while draining inflight: {}", err);
}
}
gauge!("risotto_arancini_nats_ack_inflight").set(in_flight_acks.len() as f64);
}
}
}
while let Some((enqueued_at, ack_result)) = in_flight_acks.next().await {
match ack_result {
Ok(_) => {
counter!("risotto_arancini_nats_ack_success_total").increment(1);
histogram!("risotto_arancini_nats_ack_lag_seconds")
.record(enqueued_at.elapsed().as_secs_f64());
}
Err(err) => {
counter!("risotto_arancini_nats_ack_failure_total").increment(1);
error!("JetStream ACK failed during drain: {}", err);
}
}
}
gauge!("risotto_arancini_nats_ack_inflight").set(0.0);
Ok(())
}
pub(crate) fn spawn_sidecar_thread(rx: BridgeRx, tx: Sender<Update>) -> Result<()> {
thread::Builder::new()
.name("arancini-sidecar".to_string())
.spawn(move || {
let runtime = match tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
{
Ok(runtime) => runtime,
Err(err) => {
tracing::error!("failed to build arancini sidecar runtime: {}", err);
return;
}
};
runtime.block_on(async move {
if let Err(err) = sidecar_forwarder(rx, tx).await {
tracing::error!("arancini sidecar forwarder failed: {}", err);
}
});
})
.map(|_| ())
.map_err(|err| anyhow::anyhow!("failed to spawn arancini sidecar thread: {}", err))
}
pub(crate) fn spawn_nats_sidecar_thread(rx: BridgeRx, cfg: NatsConfig) -> Result<()> {
thread::Builder::new()
.name("arancini-sidecar-nats".to_string())
.spawn(move || {
let runtime = match tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
{
Ok(runtime) => runtime,
Err(err) => {
tracing::error!("failed to build arancini nats sidecar runtime: {}", err);
return;
}
};
runtime.block_on(async move {
if let Err(err) = nats_sidecar_publisher(rx, cfg).await {
tracing::error!("arancini NATS sidecar failed: {}", err);
}
});
})
.map(|_| ())
.map_err(|err| anyhow::anyhow!("failed to spawn arancini NATS sidecar thread: {}", err))
}
#[cfg(test)]
mod tests {
use super::*;
use arancini_lib::update::UpdateAttributes;
use chrono::Utc;
use futures::FutureExt;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio::time::{Duration, timeout};
fn test_update(router_addr: IpAddr, prefix_addr: IpAddr, announced: bool) -> Update {
Update {
time_received_ns: Utc::now(),
time_bmp_header_ns: Utc::now(),
router_addr,
router_port: 4000,
peer_addr: IpAddr::V4(Ipv4Addr::new(198, 51, 100, 1)),
peer_bgp_id: Ipv4Addr::new(198, 51, 100, 1),
peer_asn: 64512,
prefix_addr,
prefix_len: 24,
is_post_policy: false,
is_adj_rib_out: false,
announced,
synthetic: false,
attrs: Arc::new(UpdateAttributes::default()),
}
}
#[test]
fn nats_subject_uses_schema_fields() {
let mut update = test_update(
IpAddr::V4(Ipv4Addr::new(192, 0, 2, 10)),
IpAddr::V4(Ipv4Addr::new(203, 0, 113, 0)),
true,
);
let attrs = UpdateAttributes {
mp_reach_afi: Some(1),
mp_reach_safi: Some(1),
..Default::default()
};
update.attrs = Arc::new(attrs);
let subject = nats_subject_for_update("arancini.updates", &update);
assert_eq!(subject, "arancini.updates.v4_192_0_2_10.64512.1_1");
}
#[test]
fn nats_subject_normalizes_ipv4_mapped_ipv6_router() {
let mut update = test_update(
IpAddr::V6(Ipv4Addr::new(203, 0, 113, 9).to_ipv6_mapped()),
IpAddr::V4(Ipv4Addr::new(203, 0, 113, 0)),
true,
);
let attrs = UpdateAttributes {
mp_reach_afi: Some(1),
mp_reach_safi: Some(1),
..Default::default()
};
update.attrs = Arc::new(attrs);
let subject = nats_subject_for_update("arancini.updates", &update);
assert_eq!(subject, "arancini.updates.v4_203_0_113_9.64512.1_1");
}
#[test]
fn nats_subject_uses_unreach_for_withdraws() {
let mut update = test_update(
IpAddr::V6(Ipv6Addr::LOCALHOST),
IpAddr::V6(Ipv6Addr::LOCALHOST),
false,
);
let attrs = UpdateAttributes {
mp_unreach_afi: Some(2),
mp_unreach_safi: Some(128),
..Default::default()
};
update.attrs = Arc::new(attrs);
let subject = nats_subject_for_update("arancini.updates", &update);
assert_eq!(subject, "arancini.updates.v6_0_0_0_0_0_0_0_1.64512.2_128");
}
#[tokio::test(flavor = "current_thread")]
async fn bridge_sender_applies_backpressure_when_bounded_queue_is_full() {
let (sender, rx) = bounded_channel(1);
sender
.send(test_update(
IpAddr::V4(Ipv4Addr::new(192, 0, 2, 1)),
IpAddr::V4(Ipv4Addr::new(203, 0, 113, 1)),
true,
))
.await
.expect("first send should succeed");
let mut second_send = sender.send(test_update(
IpAddr::V4(Ipv4Addr::new(192, 0, 2, 2)),
IpAddr::V4(Ipv4Addr::new(203, 0, 113, 2)),
true,
));
assert!(
second_send.as_mut().now_or_never().is_none(),
"second send should be pending while queue is full"
);
let _ = timeout(Duration::from_secs(1), rx.recv())
.await
.expect("drain should not timeout")
.expect("bridge receiver should still be open");
timeout(Duration::from_secs(1), second_send)
.await
.expect("second send completion should not timeout")
.expect("second send should succeed once capacity is available");
}
#[tokio::test(flavor = "current_thread")]
async fn sidecar_forwarder_returns_error_if_downstream_channel_is_closed() {
let (bridge_tx, bridge_rx) = bounded_channel(4);
let (downstream_tx, downstream_rx) = mpsc::channel(1);
drop(downstream_rx);
let forwarder = tokio::spawn(sidecar_forwarder(bridge_rx, downstream_tx));
bridge_tx
.send(test_update(
IpAddr::V4(Ipv4Addr::new(198, 51, 100, 10)),
IpAddr::V4(Ipv4Addr::new(203, 0, 113, 10)),
true,
))
.await
.expect("bridge enqueue should succeed");
let result = timeout(Duration::from_secs(1), forwarder)
.await
.expect("forwarder completion should not timeout")
.expect("forwarder join should succeed");
assert!(
result.is_err(),
"forwarder should return error when downstream channel is closed"
);
}
}