agave-xdp 4.2.1

#![allow(clippy::arithmetic_side_effects)]

use {
    crate::{
        device::{DeviceQueue, NetworkDevice, QueueId, RingSizes, TxCompletionRing},
        ecn_codepoint::EcnCodepoint,
        gre::{
            construct_gre_packet, gre_packet_size,
            packet::{GRE_HEADER_BASE_SIZE, INNER_PACKET_HEADER_SIZE},
        },
        netlink::MacAddress,
        packet::{
            IP_HEADER_SIZE, PACKET_HEADER_SIZE, UDP_HEADER_SIZE, VLAN_PACKET_HEADER_SIZE,
            construct_packet, construct_vlan_packet,
        },
        route::NextHop,
        socket::{Socket, Tx, TxRing},
        umem::{Frame, OwnedUmem, PageAlignedMemory, Umem},
    },
    agave_cpu_utils::set_cpu_affinity,
    crossbeam_channel::{Receiver, TryRecvError},
    libc::{_SC_PAGESIZE, sysconf},
    std::{
        io,
        net::{IpAddr, SocketAddr, SocketAddrV4},
        thread,
        time::Duration,
    },
};

pub struct TxLoopConfigBuilder {
    zero_copy: bool,
    maybe_src_mac: Option<MacAddress>,
}

impl TxLoopConfigBuilder {
    pub fn new() -> Self {
        Self {
            zero_copy: false,
            maybe_src_mac: None,
        }
    }

    pub fn zero_copy(&mut self, enable: bool) -> &mut Self {
        self.zero_copy = enable;
        self
    }

    pub fn override_src_mac(&mut self, mac: MacAddress) -> &mut Self {
        self.maybe_src_mac = Some(mac);
        self
    }

    pub fn build_with_src_device(self, src_device: &NetworkDevice) -> TxLoopConfig {
        let Self {
            zero_copy,
            maybe_src_mac,
        } = self;

        let src_mac = maybe_src_mac.unwrap_or_else(|| {
            // if no source MAC is provided, use the device's MAC address
            src_device
                .mac_addr()
                .expect("no src_mac provided, device must have a MAC address")
        });

        TxLoopConfig { zero_copy, src_mac }
    }
}

impl Default for TxLoopConfigBuilder {
    fn default() -> Self {
        Self::new()
    }
}

#[derive(Clone, Debug)]
pub struct TxLoopConfig {
    zero_copy: bool,
    src_mac: MacAddress,
}

pub struct TxLoopBuilder<U: Umem> {
    cpu_id: usize,
    zero_copy: bool,
    src_mac: MacAddress,
    queue: DeviceQueue,
    tx_size: usize,
    umem: U,
}

impl TxLoopBuilder<OwnedUmem<PageAlignedMemory>> {
    pub fn new(
        cpu_id: usize,
        queue_id: QueueId,
        config: TxLoopConfig,
        dev: &NetworkDevice,
    ) -> TxLoopBuilder<OwnedUmem<PageAlignedMemory>> {
        let TxLoopConfig { zero_copy, src_mac } = config;

        log::info!(
            "starting xdp loop on {} queue {queue_id:?} cpu {cpu_id}",
            dev.name()
        );

        // We don't support MTUs larger than page size due to AF_XDP limitations in single-buffer
        // mode and a possible workaround might be to use multi-buffer TX.

        // some drivers require frame_size=page_size
        let frame_size = unsafe { sysconf(_SC_PAGESIZE) } as usize;

        let queue = dev
            .open_queue(queue_id)
            .expect("failed to open queue for AF_XDP socket");
        let RingSizes {
            rx: rx_size,
            tx: tx_size,
        } = queue.ring_sizes().unwrap_or_else(|| {
            log::info!(
                "using default ring sizes for {} queue {queue_id:?}",
                dev.name()
            );
            RingSizes::default()
        });

        let frame_count = (rx_size + tx_size) * 2;

        // try to allocate huge pages first, then fall back to regular pages
        const HUGE_2MB: usize = 2 * 1024 * 1024;
        let memory =
            PageAlignedMemory::alloc_with_page_size(frame_size, frame_count, HUGE_2MB, true)
                .or_else(|_| {
                    log::warn!("huge page alloc failed, falling back to regular page size");
                    PageAlignedMemory::alloc(frame_size, frame_count)
                })
                .unwrap();
        let umem = OwnedUmem::new(memory, frame_size as u32).unwrap();

        TxLoopBuilder {
            cpu_id,
            zero_copy,
            src_mac,
            queue,
            tx_size,
            umem,
        }
    }

    pub fn build(self) -> Result<TxLoop<OwnedUmem<PageAlignedMemory>>, io::Error> {
        let TxLoopBuilder {
            cpu_id,
            zero_copy,
            src_mac,
            queue,
            tx_size,
            umem,
        } = self;

        let queue_id = queue.id();
        let (socket, tx) =
            Socket::tx(queue, umem, zero_copy, tx_size * 2, tx_size).map_err(|err| {
                log::error!(
                    "failed to create AF_XDP TX socket for queue {queue_id:?} on CPU {cpu_id}: \
                     {err}"
                );
                err
            })?;

        let Tx {
            // this is where we'll queue frames
            ring,
            // this is where we'll get completion events once frames have been picked up by the NIC
            completion,
        } = tx;
        let ring = ring.unwrap();

        Ok(TxLoop {
            cpu_id,
            src_mac,
            socket,
            ring,
            completion,
        })
    }
}

pub struct TxLoop<U: Umem> {
    cpu_id: usize,
    src_mac: MacAddress,
    socket: Socket<U>,
    ring: TxRing<U::Frame>,
    completion: TxCompletionRing,
}

/// [`TxPacket`] represents a packet to transmit via XDP to a list of addresses with the provided
/// `payload` and source address.
pub trait TxPacket {
    type Addrs: AsRef<[SocketAddr]>;
    type Payload: AsRef<[u8]>;

    /// List of destination addresses to which the packet should be sent.
    fn dst_addrs(&self) -> &Self::Addrs;

    /// Payload of the packet to be sent.
    fn payload(&self) -> &Self::Payload;

    /// Source address used when sending the packet.
    fn src_addr(&self) -> SocketAddrV4;

    /// Explicit congestion notification bits to set on the packet.
    fn ecn(&self) -> Option<EcnCodepoint>;

    /// Returns true when this packet is expected to occasionally exceed the route MTU.
    fn allow_mtu_overflow(&self) -> bool;
}

impl<U: Umem> TxLoop<U> {
    pub fn run<T, D, R>(self, receiver: Receiver<T>, mut drop_item: D, route_fn: R)
    where
        T: TxPacket,
        D: FnMut(T),
        R: Fn(&IpAddr) -> Option<NextHop>,
    {
        // How long we sleep waiting to receive packets from the channel.
        const RECV_TIMEOUT: Duration = Duration::from_nanos(1000);

        const MAX_TIMEOUTS: usize = 1;

        // We try to collect _at least_ BATCH_SIZE packets before queueing into the NIC. This is to
        // avoid introducing too much per-packet overhead and giving the NIC time to complete work
        // before we queue the next chunk of packets.
        const BATCH_SIZE: usize = 64;

        let TxLoop {
            cpu_id,
            src_mac,
            mut socket,
            mut ring,
            mut completion,
        } = self;

        // each queue is bound to its own CPU core
        set_cpu_affinity(None, [agave_cpu_utils::CpuId::new(cpu_id).unwrap()]).unwrap();

        let umem = socket.umem();
        let umem_tx_capacity = umem.available();
        let umem_frame_size = umem.frame_size();

        // Local buffer where we store packets before sending them.
        let mut batched_items = Vec::with_capacity(BATCH_SIZE);

        // How many packets we've batched. This is _not_ batched_items.len(), but item * peers. For
        // example if we have 3 packets to transmit to 2 destination addresses each, we have 6 batched
        // packets.
        let mut batched_packets = 0;
        // How many descriptors are written into the TX ring but not yet committed.
        let mut written_uncommitted = 0;

        let mut timeouts = 0;
        loop {
            match receiver.try_recv() {
                Ok(item) => {
                    batched_packets += item.dst_addrs().as_ref().len();
                    batched_items.push(item);
                    timeouts = 0;
                    if batched_packets < BATCH_SIZE {
                        continue;
                    }
                }
                Err(TryRecvError::Empty) => {
                    if timeouts < MAX_TIMEOUTS {
                        timeouts += 1;
                        thread::sleep(RECV_TIMEOUT);
                    } else {
                        timeouts = 0;
                        commit_pending(&mut ring, &mut written_uncommitted);
                        // we haven't received anything in a while, kick the driver
                        kick(&ring);
                    }
                }
                Err(TryRecvError::Disconnected) => {
                    // keep looping until we've flushed all the packets
                    if batched_packets == 0 {
                        break;
                    }
                }
            };

            for item in batched_items.drain(..) {
                let src_addr = item.src_addr();
                let src_ip = src_addr.ip();
                let src_port = src_addr.port();
                let ecn = item.ecn();
                let can_overflow_mtu = item.allow_mtu_overflow();
                for addr in item.dst_addrs().as_ref() {
                    if ring.available() == 0 || umem.available() == 0 {
                        commit_pending(&mut ring, &mut written_uncommitted);
                        kick(&ring);

                        // loop until we have space for the next packet
                        loop {
                            completion.sync(true);
                            // we haven't written any frames so we only need to sync the consumer position
                            ring.sync(false);

                            // check if any frames were completed
                            while let Some(frame_offset) = completion.read() {
                                umem.release(frame_offset);
                            }

                            if ring.available() > 0 && umem.available() > 0 {
                                // we have space for the next packet, break out of the loop
                                break;
                            }

                            // queues are full, if NEEDS_WAKEUP is set kick the driver so hopefully it'll
                            // complete some work
                            kick(&ring);
                        }
                    }

                    // at this point we're guaranteed to have a frame to write the next packet into and
                    // a slot in the ring to submit it
                    let mut frame = umem.reserve().unwrap();
                    let IpAddr::V4(dst_ip) = addr.ip() else {
                        panic!("IPv6 not supported");
                    };

                    let payload = item.payload().as_ref();
                    let len = payload.len();

                    let dst = addr.ip();
                    let Some(next_hop) = route_fn(&dst) else {
                        log::warn!("dropping packet: no route for peer {addr}");
                        batched_packets -= 1;
                        umem.release(frame.offset());
                        continue;
                    };

                    if let Some(gre) = &next_hop.gre {
                        let l3_inner_packet_len = INNER_PACKET_HEADER_SIZE + len;
                        let l3_outer_gre_packet_len =
                            IP_HEADER_SIZE + GRE_HEADER_BASE_SIZE + l3_inner_packet_len;

                        if l3_inner_packet_len > gre.mtu as usize
                            || l3_outer_gre_packet_len > next_hop.mtu as usize
                        {
                            if !can_overflow_mtu {
                                log::warn!(
                                    "dropping packet: GRE payload exceeds MTU for {addr}: L3 \
                                     inner packet length {l3_inner_packet_len}, L3 outer GRE \
                                     packet length {l3_outer_gre_packet_len}, MTU: {mtu}, \
                                     underlay_mtu: {underlay_mtu}.",
                                    mtu = gre.mtu,
                                    underlay_mtu = next_hop.mtu
                                );
                            }
                            batched_packets -= 1;
                            umem.release(frame.offset());
                            continue;
                        }

                        let packet_len = gre_packet_size(len);
                        if packet_len > umem_frame_size {
                            log::warn!(
                                "dropping packet: GRE packet size {packet_len} exceeds frame size \
                                 {umem_frame_size} for {addr}"
                            );
                            batched_packets -= 1;
                            umem.release(frame.offset());
                            continue;
                        }

                        frame.set_len(packet_len);
                        let packet = umem.map_frame_mut(&frame);
                        let inner_src_ip = next_hop.preferred_src_ip.as_ref().unwrap_or(src_ip);
                        if let Err(err) = construct_gre_packet(
                            packet,
                            &src_mac,
                            &gre.mac_addr,
                            inner_src_ip,
                            &dst_ip,
                            src_port,
                            addr.port(),
                            payload,
                            ecn,
                            &gre.tunnel_info,
                        ) {
                            log::warn!("dropping packet: {err}");
                            batched_packets -= 1;
                            umem.release(frame.offset());
                            continue;
                        }
                    } else if let Some(vlan) = &next_hop.vlan {
                        // we need the MAC address to send the packet
                        let Some(dest_mac) = next_hop.mac_addr else {
                            log::warn!(
                                "dropping packet: peer {addr} must be routed through {} which has \
                                 no known MAC address",
                                next_hop.ip_addr
                            );
                            batched_packets -= 1;
                            umem.release(frame.offset());
                            continue;
                        };

                        // The 802.1Q tag is added at L2, so the L3 size compared against the MTU
                        // is the same as the untagged path.
                        let l3_packet_len = IP_HEADER_SIZE + UDP_HEADER_SIZE + len;
                        if l3_packet_len > next_hop.mtu as usize {
                            if !can_overflow_mtu {
                                log::warn!(
                                    "dropping packet: packet size {l3_packet_len} exceeds MTU \
                                     {mtu} for {addr}",
                                    mtu = next_hop.mtu
                                );
                            }
                            batched_packets -= 1;
                            umem.release(frame.offset());
                            continue;
                        }

                        let packet_len = VLAN_PACKET_HEADER_SIZE + len;
                        if packet_len > umem_frame_size {
                            log::warn!(
                                "dropping packet: VLAN packet size {packet_len} exceeds frame \
                                 size {umem_frame_size} for {addr}"
                            );
                            batched_packets -= 1;
                            umem.release(frame.offset());
                            continue;
                        }

                        frame.set_len(packet_len);
                        let packet = umem.map_frame_mut(&frame);

                        // The route's preferred src is the IP assigned to the VLAN sub-interface,
                        // which is the right inner src for traffic egressing this VLAN. Fall back
                        // to the device's src IP if the route did not carry one.
                        let inner_src_ip = next_hop.preferred_src_ip.as_ref().unwrap_or(src_ip);

                        if !construct_vlan_packet(
                            packet,
                            &src_mac.0,
                            &dest_mac.0,
                            inner_src_ip,
                            &dst_ip,
                            src_port,
                            addr.port(),
                            vlan.vid,
                            vlan.pcp,
                            payload,
                            ecn,
                        ) {
                            log::warn!("dropping packet: VLAN frame did not fit in UMEM slot");
                            batched_packets -= 1;
                            umem.release(frame.offset());
                            continue;
                        }
                    } else {
                        // we need the MAC address to send the packet
                        let Some(dest_mac) = next_hop.mac_addr else {
                            log::warn!(
                                "dropping packet: peer {addr} must be routed through {} which has \
                                 no known MAC address",
                                next_hop.ip_addr
                            );
                            batched_packets -= 1;
                            umem.release(frame.offset());
                            continue;
                        };

                        let l3_packet_len = IP_HEADER_SIZE + UDP_HEADER_SIZE + len;
                        if l3_packet_len > next_hop.mtu as usize {
                            if !can_overflow_mtu {
                                log::warn!(
                                    "dropping packet: packet size {l3_packet_len} exceeds MTU \
                                     {mtu} for {addr}",
                                    mtu = next_hop.mtu
                                );
                            }
                            batched_packets -= 1;
                            umem.release(frame.offset());
                            continue;
                        }

                        let packet_len = PACKET_HEADER_SIZE + len;
                        if packet_len > umem_frame_size {
                            log::warn!(
                                "dropping packet: packet size {packet_len} exceeds frame size \
                                 {umem_frame_size} for {addr}"
                            );
                            batched_packets -= 1;
                            umem.release(frame.offset());
                            continue;
                        }

                        frame.set_len(packet_len);
                        let packet = umem.map_frame_mut(&frame);

                        if !construct_packet(
                            packet,
                            &src_mac.0,
                            &dest_mac.0,
                            src_ip,
                            &dst_ip,
                            src_port,
                            addr.port(),
                            payload,
                            ecn,
                        ) {
                            log::warn!("dropping packet: frame did not fit in UMEM slot");
                            batched_packets -= 1;
                            umem.release(frame.offset());
                            continue;
                        }
                    }

                    ring.write(frame, 0)
                        .map_err(|_| "ring full")
                        // this should never happen as we check for available slots above
                        .expect("failed to write to ring");

                    batched_packets -= 1;
                    written_uncommitted += 1;

                    // check if it's time to publish descriptors and kick the driver
                    if written_uncommitted >= BATCH_SIZE {
                        commit_pending(&mut ring, &mut written_uncommitted);
                        kick(&ring);
                    }
                }
                drop_item(item);
            }
            debug_assert_eq!(batched_packets, 0);
        }
        assert_eq!(batched_packets, 0);
        commit_pending(&mut ring, &mut written_uncommitted);
        kick(&ring);

        // drain the ring
        while umem.available() < umem_tx_capacity || ring.available() < ring.capacity() {
            log::debug!(
                "draining xdp ring umem {}/{} ring {}/{}",
                umem.available(),
                umem_tx_capacity,
                ring.available(),
                ring.capacity()
            );

            completion.sync(true);
            while let Some(frame_offset) = completion.read() {
                umem.release(frame_offset);
            }

            ring.sync(false);
            kick(&ring);
        }
    }
}

#[inline(always)]
fn commit_pending<F: Frame>(ring: &mut TxRing<F>, pending_uncommitted: &mut usize) {
    if *pending_uncommitted == 0 {
        return;
    }
    ring.commit();
    *pending_uncommitted = 0;
}

// With some drivers, or always when we work in SKB mode, we need to explicitly kick the driver once
// we want the NIC to do something.
#[inline(always)]
fn kick<F: Frame>(ring: &TxRing<F>) {
    if !ring.needs_wakeup() {
        return;
    }

    if let Err(e) = ring.wake() {
        kick_error(e);
    }
}

#[inline(never)]
fn kick_error(e: std::io::Error) {
    match e.raw_os_error() {
        // these are non-fatal errors
        Some(libc::EBUSY | libc::ENOBUFS | libc::EAGAIN) => {}
        // this can temporarily happen with some drivers when changing
        // settings (eg with ethtool)
        Some(libc::ENETDOWN) => {
            log::warn!("network interface is down")
        }
        // we should never get here, hopefully the driver recovers?
        _ => {
            log::error!("network interface driver error: {e:?}");
        }
    }
}