agave-xdp 4.3.0-beta.0

#[cfg(target_os = "linux")]
pub use crate::{neighbors::NeighborIntervals, tx_loop::TrySendError};
use {
    crate::ecn_codepoint::EcnCodepoint,
    bytes::Bytes,
    std::{
        error::Error,
        io,
        net::{SocketAddr, SocketAddrV4},
        sync::{Arc, atomic::AtomicBool},
        thread,
    },
};
#[cfg(target_os = "linux")]
use {
    crate::{
        device::{NetworkDevice, QueueId},
        load_xdp_program,
        neighbors::NeighborsObserver,
        route::{RouteTable, Router, RoutingTables},
        route_monitor::RouteMonitor,
        tx_loop::{self, TxLoop, TxLoopBuilder, TxLoopConfigBuilder, TxPacket},
        umem::OwnedUmem,
    },
    agave_cpu_utils::{CpuId, cpu_affinity, set_cpu_affinity},
    arc_swap::ArcSwap,
    arrayvec::ArrayVec,
    aya::Ebpf,
    crossbeam_queue::ArrayQueue,
    log::info,
    std::{
        net::{IpAddr, Ipv4Addr},
        thread::Builder,
        time::Duration,
    },
};

#[cfg(target_os = "linux")]
const ROUTE_MONITOR_UPDATE_INTERVAL: Duration = Duration::from_millis(50);

/// Binding of a single NIC hardware TX queue to a CPU core.
///
/// Each binding becomes one TX worker thread, pinned to `cpu`, whose AF_XDP
/// socket is bound to hardware queue `queue` on the configured interface.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct QueueCpuBinding {
    /// NIC hardware TX queue id the AF_XDP socket binds to.
    pub queue: u32,
    /// Logical CPU core the worker thread is pinned to.
    pub cpu: usize,
}

#[derive(Clone, Debug)]
pub struct XdpConfig {
    pub interface: Option<String>,
    /// NIC-queue -> CPU-core bindings. One TX worker is created per entry, in
    /// order. The queue id is taken explicitly from the binding rather than
    /// inferred from position, so callers can target arbitrary hardware queues.
    pub queues: Vec<QueueCpuBinding>,
    pub zero_copy: bool,
    // The capacity of the channel that sits between senders and each XDP thread that enqueues
    // packets to the NIC.
    pub tx_channel_cap: usize,
}

impl XdpConfig {
    // A nice round number
    const DEFAULT_TX_CHANNEL_CAP: usize = 1_000_000;
}

impl Default for XdpConfig {
    fn default() -> Self {
        Self {
            interface: None,
            queues: vec![],
            zero_copy: false,
            tx_channel_cap: Self::DEFAULT_TX_CHANNEL_CAP,
        }
    }
}

impl XdpConfig {
    pub fn new(
        interface: Option<impl Into<String>>,
        queues: Vec<QueueCpuBinding>,
        zero_copy: bool,
    ) -> Self {
        Self {
            interface: interface.map(|s| s.into()),
            queues,
            zero_copy,
            tx_channel_cap: XdpConfig::DEFAULT_TX_CHANNEL_CAP,
        }
    }

    #[cfg(feature = "dev-context-only-utils")]
    pub fn with_tx_channel_cap(
        interface: Option<impl Into<String>>,
        queues: Vec<QueueCpuBinding>,
        zero_copy: bool,
        tx_channel_cap: usize,
    ) -> Self {
        Self {
            interface: interface.map(|s| s.into()),
            queues,
            zero_copy,
            tx_channel_cap,
        }
    }
}

/// [`BytesTxPacket`] encapsulates the information needed to transmit a packet via XDP. Besides
/// the payload and destination addresses, it includes the source address of the packet.
#[cfg(target_os = "linux")]
pub struct BytesTxPacket {
    src_addr: SocketAddrV4,
    dst_addrs: XdpAddrs,
    ecn: Option<EcnCodepoint>,
    allow_mtu_overflow: bool,
    payload: Bytes,
}

#[cfg(not(target_os = "linux"))]
pub struct BytesTxPacket;

#[cfg(target_os = "linux")]
impl BytesTxPacket {
    pub fn new(
        src_addr: SocketAddrV4,
        dst_addrs: impl Into<XdpAddrs>,
        ecn: Option<EcnCodepoint>,
        payload: Bytes,
    ) -> Self {
        Self {
            src_addr,
            dst_addrs: dst_addrs.into(),
            ecn,
            allow_mtu_overflow: false,
            payload,
        }
    }

    /// Sets whether MTU overflow is possible for this packet.
    pub fn set_allow_mtu_overflow(&mut self, allow: bool) {
        self.allow_mtu_overflow = allow;
    }
}

#[cfg(not(target_os = "linux"))]
impl BytesTxPacket {
    pub fn new(
        _src_addr: SocketAddrV4,
        _dst_addrs: impl Into<XdpAddrs>,
        _ecn: Option<EcnCodepoint>,
        _payload: Bytes,
    ) -> Self {
        Self
    }

    pub fn set_allow_mtu_overflow(&mut self, _allow: bool) {}
}

#[cfg(not(target_os = "linux"))]
pub enum TrySendError<T> {
    Full(T),
    Disconnected(T),
}

#[cfg(not(target_os = "linux"))]
impl std::fmt::Debug for TrySendError<BytesTxPacket> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TrySendError::Full(_) => write!(f, "TrySendError::Full"),
            TrySendError::Disconnected(_) => write!(f, "TrySendError::Disconnected"),
        }
    }
}

#[cfg(target_os = "linux")]
impl TxPacket for BytesTxPacket {
    type Addrs = XdpAddrs;
    type Payload = Bytes;

    fn dst_addrs(&self) -> &Self::Addrs {
        &self.dst_addrs
    }

    fn payload(&self) -> &Self::Payload {
        &self.payload
    }

    fn src_addr(&self) -> SocketAddrV4 {
        self.src_addr
    }

    fn ecn(&self) -> Option<EcnCodepoint> {
        self.ecn
    }

    fn allow_mtu_overflow(&self) -> bool {
        self.allow_mtu_overflow
    }
}

#[derive(Clone)]
pub struct XdpSender {
    #[cfg(target_os = "linux")]
    senders: Vec<tx_loop::TxSender<BytesTxPacket>>,
}

pub enum XdpAddrs {
    Single(SocketAddr),
    Multi(Arc<[SocketAddr]>),
}

impl From<SocketAddr> for XdpAddrs {
    #[inline]
    fn from(addr: SocketAddr) -> Self {
        XdpAddrs::Single(addr)
    }
}

impl From<Vec<SocketAddr>> for XdpAddrs {
    #[inline]
    fn from(addrs: Vec<SocketAddr>) -> Self {
        XdpAddrs::Multi(addrs.into())
    }
}

impl From<Arc<[SocketAddr]>> for XdpAddrs {
    #[inline]
    fn from(addrs: Arc<[SocketAddr]>) -> Self {
        XdpAddrs::Multi(addrs)
    }
}

impl AsRef<[SocketAddr]> for XdpAddrs {
    #[inline]
    fn as_ref(&self) -> &[SocketAddr] {
        match self {
            XdpAddrs::Single(addr) => std::slice::from_ref(addr),
            XdpAddrs::Multi(addrs) => addrs,
        }
    }
}

impl XdpSender {
    /// Validate positions for a sender subset before an `XdpSender` is available.
    pub fn validate_subset_positions(
        positions: &[usize],
        sender_count: usize,
    ) -> Result<(), io::Error> {
        fn invalid_input(message: impl Into<String>) -> io::Error {
            io::Error::new(io::ErrorKind::InvalidInput, message.into())
        }

        if positions.is_empty() {
            return Err(invalid_input("XDP sender subset cannot be empty"));
        }
        if let Some(&position) = positions.iter().find(|&&position| position >= sender_count) {
            return Err(invalid_input(format!(
                "XDP sender subset position {position} is out of range for {sender_count} \
                 configured XDP sender(s)"
            )));
        }
        if let Some((_, &position)) = positions
            .iter()
            .enumerate()
            .find(|(i, position)| positions[..*i].contains(position))
        {
            return Err(invalid_input(format!(
                "XDP sender subset position {position} is repeated"
            )));
        }
        Ok(())
    }

    /// Return a sender restricted to `positions` in this sender's queue list, in the given order.
    ///
    /// `positions` must be non-empty and free of duplicates, and each element must be less than
    /// `self.len()`. Position `i` of the returned sender maps to `positions[i]` of this one, so
    /// the order determines which queue each `try_send` index lands on.
    #[cfg(target_os = "linux")]
    pub fn subset(&self, positions: &[usize]) -> Result<XdpSender, io::Error> {
        Self::validate_subset_positions(positions, self.len())?;

        Ok(XdpSender {
            senders: positions.iter().map(|&i| self.senders[i].clone()).collect(),
        })
    }

    #[cfg(not(target_os = "linux"))]
    pub fn subset(&self, _positions: &[usize]) -> Result<XdpSender, io::Error> {
        Err(io::Error::new(
            io::ErrorKind::Unsupported,
            "XDP is only supported on Linux",
        ))
    }

    #[inline]
    pub fn try_send(
        &self,
        sender_index: usize,
        packet: BytesTxPacket,
    ) -> Result<(), TrySendError<BytesTxPacket>> {
        #[cfg(target_os = "linux")]
        {
            let idx = sender_index
                .checked_rem(self.senders.len())
                .expect("XdpSender::senders should not be empty");
            self.senders[idx].try_send(packet)
        }
        #[cfg(not(target_os = "linux"))]
        {
            let _ = sender_index;
            Err(TrySendError::Disconnected(packet))
        }
    }

    pub fn len(&self) -> usize {
        #[cfg(target_os = "linux")]
        return self.senders.len();

        #[cfg(not(target_os = "linux"))]
        0
    }

    pub fn is_empty(&self) -> bool {
        #[cfg(target_os = "linux")]
        return self.senders.is_empty();

        #[cfg(not(target_os = "linux"))]
        true
    }
}

pub struct Transmitter {
    threads: Vec<thread::JoinHandle<()>>,
}

#[cfg(not(target_os = "linux"))]
pub struct TransmitterBuilder {}

#[cfg(target_os = "linux")]
pub struct TransmitterBuilder {
    tx_loops: Vec<TxLoop<OwnedUmem>>,
    tx_channel_cap: usize,
    maybe_ebpf: Option<Ebpf>,
    atomic_router: Arc<ArcSwap<Router>>,
    neighbors: NeighborsObserver,
    neighbors_monitor_handle: thread::JoinHandle<()>,
    route_monitor_handle: thread::JoinHandle<()>,
}

impl TransmitterBuilder {
    #[cfg(not(target_os = "linux"))]
    pub fn new(_config: XdpConfig, _exit: Arc<AtomicBool>) -> Result<Self, Box<dyn Error>> {
        Err("XDP is only supported on Linux".into())
    }

    #[cfg(target_os = "linux")]
    pub fn new(config: XdpConfig, exit: Arc<AtomicBool>) -> Result<Self, Box<dyn Error>> {
        Self::new_with_intervals(
            config,
            exit,
            NeighborIntervals {
                use_interval: Duration::from_secs(30),
                miss_interval: Duration::from_secs(1),
            },
        )
    }

    #[cfg(target_os = "linux")]
    pub fn new_with_intervals(
        config: XdpConfig,
        exit: Arc<AtomicBool>,
        neighbor_intervals: NeighborIntervals,
    ) -> Result<Self, Box<dyn Error>> {
        use {
            crate::neighbors::NeighborsRefresher,
            caps::Capability::{CAP_BPF, CAP_NET_ADMIN, CAP_NET_RAW, CAP_PERFMON},
            log::debug,
            std::{collections::HashSet, io},
        };
        let XdpConfig {
            interface: maybe_interface,
            queues,
            zero_copy,
            tx_channel_cap,
        } = config;

        let dev = Arc::new(if let Some(interface) = maybe_interface {
            NetworkDevice::new(interface).unwrap()
        } else {
            NetworkDevice::new_from_default_route().unwrap()
        });

        let mut tx_loop_config_builder = TxLoopConfigBuilder::new();
        tx_loop_config_builder.zero_copy(zero_copy);
        let tx_loop_config = tx_loop_config_builder.build_with_src_device(&dev);

        let reserved_cores = queues
            .iter()
            .map(|binding| CpuId::new(binding.cpu))
            .collect::<io::Result<HashSet<_>>>()?;
        let unreserved_cores = cpu_affinity(None)?
            .into_iter()
            .filter(|core| !reserved_cores.contains(core))
            .collect::<Vec<_>>();

        if unreserved_cores.is_empty() {
            return Err("all CPUs are reserved; no CPU available for the main thread".into());
        }

        let mut tx_loop_builders = Vec::with_capacity(queues.len());
        for binding in queues {
            // since we aren't necessarily allocating from the thread that we intend to run on,
            // temporarily switch to the target cpu for each TxLoop to ensure that the Umem region
            // is allocated to the correct numa node
            let cpu = CpuId::new(binding.cpu)?;
            set_cpu_affinity(None, [cpu])?;
            let tx_loop_builder = TxLoopBuilder::new(
                binding.cpu,
                QueueId(binding.queue as u64),
                tx_loop_config.clone(),
                &dev,
            );
            // migrate main thread back off of the last xdp reserved cpu
            set_cpu_affinity(None, unreserved_cores.iter().copied())?;
            tx_loop_builders.push(tx_loop_builder);
        }

        // switch to higher caps while we setup XDP. We assume that an error in
        // this function is irrecoverable so we don't try to drop on errors.
        let _setup_caps =
            CapGuard::raise([CAP_NET_ADMIN, CAP_NET_RAW]).expect("raise net capabilities");

        let maybe_ebpf_result = if zero_copy {
            let _ebpf_caps =
                CapGuard::raise([CAP_BPF, CAP_PERFMON]).expect("raise ebpf capabilities");

            let load_result =
                load_xdp_program(&dev).map_err(|e| format!("failed to attach xdp program: {e}"));

            Some(load_result)
        } else {
            None
        };

        let tx_loops = tx_loop_builders
            .into_iter()
            .map(|tx_loop_builder| tx_loop_builder.build())
            .collect::<Result<Vec<_>, io::Error>>()?;

        let tables_result = RoutingTables::from_netlink(RouteTable::Main);

        let tables = tables_result?;
        let router = Router::from_tables(tables)?;
        debug!(
            "published router table {}:\n{}",
            RouteTable::Main,
            router.routing_table()
        );

        fn retain_cap_net_admin() {
            // we need to retain CAP_NET_ADMIN in case the netlink socket needs reinitialized
            let retained_caps = caps::CapsHashSet::from_iter([caps::Capability::CAP_NET_ADMIN]);
            caps::set(None, caps::CapSet::Effective, &retained_caps)
                .expect("linux allows effective capset to be set");
            caps::set(None, caps::CapSet::Permitted, &retained_caps)
                .expect("linux allows permitted capset to be set");
        }

        // Use ArcSwap for lock-free updates of the routing table
        let atomic_router = Arc::new(ArcSwap::from_pointee(router));
        let route_monitor_handle = RouteMonitor::start(
            Arc::clone(&atomic_router),
            RouteTable::Main,
            exit.clone(),
            ROUTE_MONITOR_UPDATE_INTERVAL,
            || {
                retain_cap_net_admin();
                info!("route monitor thread started");
            },
        );
        let (neighbors_monitor_handle, neighbors) =
            NeighborsRefresher::start(exit, neighbor_intervals, || {
                retain_cap_net_admin();
                info!("neighbors thread started");
            })?;

        let maybe_ebpf = maybe_ebpf_result.transpose()?;

        Ok(Self {
            tx_loops,
            tx_channel_cap,
            maybe_ebpf,
            atomic_router,
            neighbors,
            neighbors_monitor_handle,
            route_monitor_handle,
        })
    }

    pub fn sender_count(&self) -> usize {
        #[cfg(target_os = "linux")]
        return self.tx_loops.len();

        #[cfg(not(target_os = "linux"))]
        0
    }

    #[cfg(not(target_os = "linux"))]
    pub fn build(self) -> (Transmitter, XdpSender) {
        (Transmitter { threads: vec![] }, XdpSender {})
    }

    #[cfg(target_os = "linux")]
    pub fn build(self) -> (Transmitter, XdpSender) {
        const DROP_CHANNEL_CAP: usize = 1_000_000;

        let Self {
            tx_loops,
            tx_channel_cap,
            maybe_ebpf,
            atomic_router,
            neighbors,
            neighbors_monitor_handle,
            route_monitor_handle,
        } = self;

        let drop_queue = Arc::new(ArrayQueue::new(DROP_CHANNEL_CAP));
        let mut threads = vec![route_monitor_handle, neighbors_monitor_handle];

        threads.push(
            Builder::new()
                .name("solTransmDrop".to_owned())
                .spawn({
                    let drop_queue = Arc::clone(&drop_queue);
                    move || {
                        loop {
                            // drop shreds in a dedicated thread so that we never lock/madvise() from
                            // the xdp thread
                            match drop_queue.pop() {
                                Some(i) => {
                                    drop(i);
                                }
                                None if Arc::strong_count(&drop_queue) == 1 => break,
                                None => {
                                    thread::sleep(Duration::from_millis(1));
                                }
                            }
                        }
                        // move the ebpf program here so it stays attached until we exit
                        drop(maybe_ebpf);
                    }
                })
                .unwrap(),
        );

        let mut senders = vec![];
        for (i, tx_loop) in tx_loops.into_iter().enumerate() {
            let (sender, receiver) = tx_loop::channel(tx_channel_cap);
            let drop_queue = Arc::clone(&drop_queue);
            let atomic_router = Arc::clone(&atomic_router);
            let mut neighbors = neighbors.clone();
            threads.push(
                Builder::new()
                    .name(format!("solTransmIO{i:02}"))
                    .spawn(move || {
                        tx_loop.run(
                            receiver,
                            move |item| {
                                if let Err(item) = drop_queue.push(item) {
                                    drop(item);
                                }
                            },
                            move |ip| route(ip, &atomic_router.load(), &mut neighbors),
                        )
                    })
                    .unwrap(),
            );
            senders.push(sender);
        }

        (Transmitter { threads }, XdpSender { senders })
    }
}

#[cfg(target_os = "linux")]
fn route(
    ip: &IpAddr,
    router: &Router,
    neighbors: &mut NeighborsObserver,
) -> Option<crate::route::NextHop> {
    let IpAddr::V4(ip) = ip else {
        return None;
    };

    let next_hop = router.route_v4(*ip).ok()?;
    if next_hop.neigh_requires_refresh {
        if let Some(gre) = next_hop.gre.as_ref() {
            neighbors.observe(
                gre.underlay_if_index,
                gre.underlay_ip_addr,
                gre.underlay_mac_addr.is_some(),
            );
        } else {
            let IpAddr::V4(neighbor_ip) = next_hop.ip_addr else {
                return None;
            };
            neighbors.observe(next_hop.if_index, neighbor_ip, next_hop.mac_addr.is_some());
        }
    }
    Some(next_hop)
}

impl Transmitter {
    pub fn join(self) -> thread::Result<()> {
        for handle in self.threads {
            handle.join()?;
        }
        Ok(())
    }
}

/// Returns the IPv4 address of the master interface if the given interface is part of a bond.
#[cfg(target_os = "linux")]
pub(crate) fn master_ip_if_bonded(interface: &str) -> Option<Ipv4Addr> {
    let master_ifindex_path = format!("/sys/class/net/{interface}/master/ifindex");
    if let Ok(contents) = std::fs::read_to_string(&master_ifindex_path) {
        let idx = contents.trim().parse().unwrap();
        return Some(
            NetworkDevice::new_from_index(idx)
                .and_then(|dev| dev.ipv4_addr())
                .unwrap_or_else(|e| {
                    panic!(
                        "failed to open bond master interface for {interface}: master index \
                         {idx}: {e}"
                    )
                }),
        );
    }
    None
}

#[cfg(target_os = "linux")]
const CAP_GUARD_CAPACITY: usize = 2;

#[cfg(target_os = "linux")]
#[must_use = "capabilities are dropped when the guard goes out of scope"]
struct CapGuard {
    capabilities: ArrayVec<caps::Capability, CAP_GUARD_CAPACITY>,
}

#[cfg(target_os = "linux")]
impl CapGuard {
    fn raise(
        raised_capabilities: impl IntoIterator<Item = caps::Capability>,
    ) -> Result<Self, caps::errors::CapsError> {
        let mut capabilities = ArrayVec::new();
        for capability in raised_capabilities {
            capabilities.try_push(capability).unwrap_or_else(|_| {
                panic!("CapGuard supports at most {CAP_GUARD_CAPACITY} capabilities")
            });
            caps::raise(None, caps::CapSet::Effective, capability)?;
        }
        Ok(Self { capabilities })
    }
}

#[cfg(target_os = "linux")]
impl Drop for CapGuard {
    fn drop(&mut self) {
        for capability in self.capabilities.iter().rev() {
            caps::drop(None, caps::CapSet::Effective, *capability)
                .unwrap_or_else(|err| panic!("drop {capability:?} capability: {err}"));
        }
    }
}

#[cfg(all(test, target_os = "linux"))]
mod tests {
    use {
        super::*,
        crate::tx_loop::{Receiver, TryRecvError, TxReceiver},
    };

    /// The receivers are returned so they stay connected for the lifetime of the test.
    fn sender_with_receivers(sender_count: usize) -> (XdpSender, Vec<TxReceiver<BytesTxPacket>>) {
        let (senders, receivers) = (0..sender_count).map(|_| tx_loop::channel(1)).unzip();
        (XdpSender { senders }, receivers)
    }

    fn packet() -> BytesTxPacket {
        BytesTxPacket::new(
            SocketAddrV4::new(Ipv4Addr::LOCALHOST, 1),
            SocketAddr::from((Ipv4Addr::LOCALHOST, 2)),
            None,
            Bytes::new(),
        )
    }

    #[test]
    fn subset_rejects_invalid_positions() {
        let (sender, _receivers) = sender_with_receivers(2);
        for (positions, expected) in [
            (&[][..], "cannot be empty"),
            (&[0, 2][..], "out of range"),
            (&[1, 0, 1][..], "is repeated"),
        ] {
            let Err(error) = sender.subset(positions) else {
                panic!("invalid subset {positions:?} must fail");
            };
            assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
            assert!(
                error.to_string().contains(expected),
                "unexpected error for {positions:?}: {error}"
            );
        }
    }

    #[test]
    fn subset_maps_positions_in_order() {
        let (sender, receivers) = sender_with_receivers(3);
        let subset = sender.subset(&[2, 0]).unwrap();
        assert_eq!(subset.len(), 2);

        subset.try_send(0, packet()).unwrap();
        subset.try_send(1, packet()).unwrap();

        assert!(receivers[2].try_recv().is_ok());
        assert!(receivers[0].try_recv().is_ok());
        assert!(matches!(receivers[1].try_recv(), Err(TryRecvError::Empty)));
    }
}