Game Servers on Linux: Part 1

// Create the socket and the channel for passing
// received packets back to the main thread
let (socket, rx) = setup_socket(port);

let time_per_tick = Duration::from_millis(20); // 50 ticks per second
let mut last_tick = Instant::now();
let mut addr_set: HashSet<SocketAddr> = HashSet::new();

loop {
    let now = Instant::now();
    // Check if it's time for the next tick
    if now.duration_since(last_tick) >= time_per_tick {
        // Run our "per-tick" game logic (e.g. broadcasting packets)
        game_tick(&rx, &socket, &mut addr_set);
        last_tick = now;
    } else {
        // Sleep until the next tick
        let time_to_sleep = time_per_tick - now.duration_since(last_tick);
        if time_to_sleep > Duration::from_millis(0) {
            thread::sleep(time_to_sleep);
        }
    }
}

fn setup_socket(port: u16) -> (UdpSocket, Receiver<(SocketAddr, Vec<u8>)>) {
    let socket = UdpSocket::bind(("0.0.0.0", port)).expect("Failed to bind socket");
    let (tx, rx) = mpsc::channel();
    let cloned_socket = socket.try_clone().expect("Failed to clone socket");
    thread::spawn(move || {
        // Use a buffer of 1500 bytes, the maximum size of a typical Ethernet frame
        let mut buf = [0u8; 1500];
        loop {
            match cloned_socket.recv_from(&mut buf) {
                Ok((size, src)) => {
                    // Create a copy of the packet that we can pass to the main thread
                    let packet = buf[..size].to_vec();
                    tx.send((src, packet)).expect("Failed to send packet to main thread");
                }
                Err(e) => {
                    // Do nothing on timeout errors
                    if e.kind() != std::io::ErrorKind::WouldBlock
                        && e.kind() != std::io::ErrorKind::TimedOut
                    {
                        eprintln!("Failed to receive packet: {}", e);
                    }
                }
            }
        }
    });
    return (socket, rx);
}

fn game_tick(
    rx: &Receiver<(SocketAddr, Vec<u8>)>,
    socket: &UdpSocket,
    addr_set: &mut HashSet<SocketAddr>,
) {
    while let Ok((src, packet)) = rx.try_recv() {
        addr_set.insert(src);
        for addr in addr_set.iter() {
            socket
                .send_to(&packet, addr)
                .expect("Failed to send packet");
        }
    }
}

% time     seconds  usecs/call     calls    errors syscall
------ ----------- ----------- --------- --------- ------------------
 57.72   16.807322         168    100009         9 recvfrom
 38.14   11.103956          11    999955           sendto
  4.13    1.203268         218      5504         1 clock_nanosleep
  0.00    0.000461          19        24           write
  ...
------ ----------- ----------- --------- --------- ------------------
100.00   29.116518          26   1105625        11 total

// This is zero in this case since we're only using one socket
const SOCKET_IDENTIFIER: Fixed = Fixed(0);

// This is an arbitrary identifier for our buffer group
const BUFFER_GROUP_ID: u16 = 0xbead;

enum TokenType {
    Recv,
    Send(Vec<u8>),
}

fn setup_iouring_socket(port: u16) -> (UdpSocket, IoUring, FixedSizeBufRing, Slab<TokenType>) {
    // Build our io_uring instance
    let mut io_uring = IoUring::builder()
        // Signal to the kernel that we will only be submitting entries from a single thread
        .setup_single_issuer()
        // Setup a large completion queue to allow for many in-flight operations
        .setup_cqsize(8192)
        // Don't stop submitting if we hit an error on one entry
        .setup_submit_all()
        // Build with a submission queue size of 8192 entries
        .build(8192)
        .expect("Failed to create IoUring");

    // There is some complexity here that I'm glossing over for the sake of brevity:
    // You can think of this as a large buffer that will be broken into smaller,
    // uniform chunks for use with recv operations.
    // See: https://man7.org/linux/man-pages/man3/io_uring_register_buf_ring.3.html
    let buf_ring = register_buf_ring::Builder::new(BUFFER_GROUP_ID)
        .buf_cnt(1024)
        .buf_len(MAX_PACKET_SIZE)
        .build()
        .expect("Failed to create buffer ring");

    // Register the buffer ring with io_uring. This allows us to use the buffer
    // ring for recv operations (necessary for multishot recv).
    buf_ring.rc.register(&mut io_uring).expect("Failed to register buffer ring");

    let socket = UdpSocket::bind(("0.0.0.0", port)).expect("Failed to bind socket");

    // This is used to track when it's safe to reuse/drop buffers. At a high level,
    // this gives us a usize that we can pass with the io_uring entry to identify
    // the operation later.
    let mut slab = Slab::with_capacity(256);
    let recv_token_index = slab.insert(TokenType::Recv);

    // Register the socket's file descriptor with io_uring.
    io_uring.submitter()
        .register_files(&[socket.as_raw_fd()])
        .expect("Failed to register socket fd");

    // Prepare the multishot recv entry
    let recv = prep_multi_recv(recv_token_index, &MSG_HDR as *const _);

    unsafe {
        // Push the multishot recv entry to the submission queue.
        io_uring
            .submission()
            .push(&recv)
            .expect("Failed to push multi-recv entry");
    }

    // Submit the io_uring entries to the kernel
    // This should generate a single syscall to io_uring_enter
    io_uring
        .submit()
        .expect("Failed to submit io_uring entries");

    return (socket, io_uring, buf_ring, slab);
}

// This prepares a "multishot" recv operation. These allow us to continuously receive
// packets until the operation returns an entry without the IORING_CQE_F_MORE flag set.
fn prep_multi_recv(
    token_index: usize,
    msghdr: *const libc::msghdr,
) -> io_uring::squeue::Entry {
    return opcode::RecvMsgMulti::new(
        SOCKET_IDENTIFIER,
        msghdr,
        BUFFER_GROUP_ID,
    )
    .build()
    .user_data(token_index as _);
}

fn iouring_recv_game_tick(
    socket: &UdpSocket,
    io_uring: &mut IoUring,
    buf_ring: &register_buf_ring::FixedSizeBufRing,
    slab: &mut Slab<TokenType>,
    addr_set: &mut HashSet<SocketAddr>,
    addr_map: &mut HashMap<SocketAddr, Box<libc::sockaddr_in>>,
) {
    // Process the completion queue
    let (submitter, mut sq, cq) = io_uring.split();
    let mut count = sq.len();

    // Iterate through all the completion queue entries
    for cqe in cq {
        // Get the token index we stored in the user_data field
        let token_index = cqe.user_data() as usize;
        match slab.get(token_index) {
            Some(TokenType::Recv) => {
                let result = cqe.result();
                if result > 1500 || result < 0 {
                    eprintln!("Multi-recv returned unexpected result: {}", result);
                    continue;
                }

                // Once again, I'm going to gloss over the details of buffer rings here.
                // For the purposes of this example, just know that we can use the flags
                // on the cqe to get the buffer we used for this recv operation.
                let bufs = buf_ring.rc.get_bufs(buf_ring, result as u32, cqe.flags());
                if bufs.len() != 1 {
                    continue;
                }
                let buf_entry = &bufs[0];
                let data = buf_entry.as_slice();

                // Create empty msghdr to parse the `from` address
                let mut msghdr: libc::msghdr = unsafe { mem::zeroed() };
                msghdr.msg_namelen = 32;
                msghdr.msg_controllen = 0;
                let msg = RecvMsgOut::parse(data, &msghdr).expect("Failed to parse RecvMsgOut");
                let name_data = unsafe { *(msg.name_data().as_ptr() as *const libc::sockaddr_in) };

                let addr = SocketAddr::V4(SocketAddrV4::new(
                    u32::from_be(name_data.sin_addr.s_addr).into(),
                    u16::from_be(name_data.sin_port),
                ));

                addr_set.insert(addr);
                for addr in addr_set.iter() {
                    // Get or create a cached sockaddr_in for this address
                    let sockaddr = addr_map.entry(*addr).or_insert_with(|| Box::new(libc::sockaddr_in {
                        sin_family: libc::AF_INET as u16,
                        sin_port: u16::to_be(addr.port()),
                        sin_addr: libc::in_addr {
                            s_addr: u32::to_be(match addr {
                                SocketAddr::V4(a) => a.ip().to_owned().into(),
                                _ => 0,
                            }),
                        },
                        sin_zero: [0; 8],
                    }));

                    // Get a pointer to the sockaddr_in
                    let addr_ptr: *const libc::sockaddr_in = &**sockaddr;

                    // Create an owned copy of the payload to send
                    let payload = msg.payload_data().to_vec();
                    let payload_ptr = payload.as_ptr();
                    let payload_len = payload.len();

                    let index = slab.insert(TokenType::Send(payload));

                    // Use a zero-copy send operation to send the packet
                    let send = opcode::SendZc::new(
                                    SOCKET_IDENTIFIER_V4,
                                    payload_ptr,
                                    payload_len as u32,
                                )
                                .dest_addr(addr_ptr as *const _)
                                .dest_addr_len(size_of::<libc::sockaddr_in>() as u32)
                                .build()
                                .user_data(index as _);
                    unsafe {
                        sq.push(&send).expect("Failed to push send entry");
                    }

                    // If we have enough entries queued up, submit them to the kernel,
                    // otherwise we may run out of space in the submission queue.
                    if sq.len() - count >= 1024 {
                        submitter.submit().expect("Failed to submit io_uring entries");
                        count = sq.len();
                    }
                }

                // If the MORE flag is not set, we need to re-submit a new multishot recv
                if cqe.flags() & IORING_CQE_F_MORE == 0 {
                    let recv = prep_multi_recv(token_index, &MSG_HDR as *const _);
                    unsafe {
                        sq.push(&recv).expect("Failed to push multi-recv entry");
                    }
                }
            }
            Some(TokenType::Send(_)) => {
                // SendZc can produce multiple completions for a single send operation.
                // We only care about the notification completion (has IORING_CQE_F_NOTIF
                // flag set).
                if !cqueue::notif(cqe.flags()) {
                    continue;
                }
                slab.remove(token_index);
            },
            None => {
                eprintln!("Received unknown token index: {}", token_index);
            }
        }
    }
    // Sync the submission queue before submitting
    sq.sync();

    submitter.submit().expect("Failed to submit io_uring entries");
}

% time     seconds  usecs/call     calls    errors syscall
------ ----------- ----------- --------- --------- ------------------
 82.87    4.685087         824      5680       635 io_uring_enter
 13.65    0.771613          22     34957     29938 restart_syscall
  3.46    0.195726          38      5027      5019 clock_nanosleep
  0.01    0.000474          18        25         1 write
...
------ ----------- ----------- --------- --------- ------------------
100.00    5.653862         123     45824     35594 total