syscall or die

sockets.zig

const std = @import("std");
const net = std.net;
const os = std.os;
const info = std.log.info;
const expect = std.testing.expect;
const system = os.system;
const linux = os.linux;
const c = std.c;

const MAX_CONNECTIONS = 32;
const N_CLIENTS = 3;
const BUFSIZE = 1024;

pub const mmsghdr = extern struct {
    msg_hdr: c.msghdr,
    msg_len: c_ulong,
};

// int sendmmsg(int sockfd, struct mmsghdr *msgvec, unsigned int vlen, int flags);
pub extern "c" fn sendmmsg(sockfd: c_int, msgvec: [*]mmsghdr, vlen: c_uint, flags: c_int) c_int;

// int recvmmsg(int sockfd, struct mmsghdr *msgvec, unsigned int vlen, int flags, struct timespec *timeout);
pub extern "c" fn recvmmsg(sockfd: c_int, msgvec: [*]mmsghdr, vlen: c_uint, flags: c_int, timeout: *c.timespec) c_int;

pub fn start_server(server_addr: net.Address) !void {
    // setup server socket
    var sockfd = try os.socket(os.AF.INET, os.SOCK.DGRAM, os.IPPROTO.UDP);
    errdefer os.close(sockfd);

    try os.bind(sockfd, &server_addr.any, server_addr.getOsSockLen());

    // init
    var connections: [MAX_CONNECTIONS]mmsghdr = undefined; 
    var bufs: [MAX_CONNECTIONS][BUFSIZE]u8 = undefined; 
    var iovec: [MAX_CONNECTIONS][1]os.iovec = undefined;

    for (0..MAX_CONNECTIONS) |i| { 
        iovec[i][0].iov_base = &bufs[i];
        iovec[i][0].iov_len = BUFSIZE;

        var addr: net.Address = undefined;
        connections[i] = mmsghdr { 
            .msg_hdr = c.msghdr { 
                .name = &addr.any, 
                .namelen = server_addr.getOsSockLen(), 
                .iov = &iovec[i], 
                .iovlen = 1, 
                .control = null, 
                .controllen = 0, 
                .flags = 0, 
            },
            .msg_len = 0,
        };
    }
    
    var timeout = c.timespec { 
        .tv_sec = 1, 
        .tv_nsec = 0, 
    };

    // listen
    while (true) {
        const n_events_c = recvmmsg(
            @intCast(c_int, sockfd), 
            &connections,
            @intCast(c_uint, MAX_CONNECTIONS), 
            @intCast(c_int, 0),
            &timeout
        );

        if (n_events_c == -1) { 
            const err_enum = c.getErrno(n_events_c);
            // std/os/linux/errno/generic.zig lookup
            const err_n = @intFromEnum(err_enum); 
            if (err_n == 11) { // AGAIN
                std.time.sleep(std.time.ns_per_ms);
                continue;
            }
            info("recvmmsg err number: {d}", .{err_n});
            return;
        } 

        const n_events = @intCast(usize, n_events_c);
        info("recvmmsg n_events: {}", .{ n_events });

        for (0..n_events) |i| { 
            var v = bufs[i][0];
            var len = connections[i].msg_len;
            info("=> recieved msg -- client_id: {} len: {} @ {*}", .{v, len, connections[i].msg_hdr.iov[0].iov_base});
        }

        // echo back 
        _ = sendmmsg(@intCast(c_int, sockfd), &connections, @intCast(c_uint, n_events), @intCast(c_int, 0));
    }
}

pub fn start_client(server_addr: net.Address, id: usize) !void { 
    // setup client socket
    const sockfd = try os.socket(os.AF.INET, os.SOCK.DGRAM, os.IPPROTO.UDP);
    errdefer os.closeSocket(sockfd);

    const client = try net.Address.resolveIp("127.0.0.1", @intCast(u16, 9999 + id));
    try os.bind(sockfd, &client.any, client.getOsSockLen());

    // talk to server
    info("client {}: sending...", .{id});
    const v = @intCast(u8, id);
    const req = [_]u8{ v, v + 1, v + 2 }; 

    var addr: net.Address = undefined; 
    var addr_len: system.socklen_t = client.getOsSockLen();
    var resp: [3]u8 = undefined;

    _ = try os.sendto(sockfd, &req, 0, &server_addr.any, server_addr.getOsSockLen());
    _ = try os.recvfrom(sockfd, &resp, 0, &addr.any, &addr_len);

    // valid echo?
    if (!std.mem.eql(u8, &resp, &req)) {
        info("client {}: failed...", .{id});
    }
    info("client {}: done", .{id});
}

pub fn main() !void {
    const server_addr = try net.Address.resolveIp("127.0.0.1", 1234);
    const server_handle = try std.Thread.spawn(.{}, start_server, .{server_addr});

    std.time.sleep(std.time.ns_per_ms * std.time.ms_per_s);
    var timer = try std.time.Timer.start();
    const start = timer.read();

    // should also use the `client.sh` script (msgs are dropped sometimes like this)
    var client_handles: [N_CLIENTS]std.Thread = undefined;
    for (0..N_CLIENTS) |i| { 
        const ct = try std.Thread.spawn(.{}, start_client, .{server_addr, i});
        client_handles[i] = ct; 
    }
    for (0..N_CLIENTS) |i| {
        client_handles[i].join();
    }

    const end = timer.read();
    const elapsed = end - start;

    // 2110123333 on poll (accept() syscall)
    info("elapsed: {d}", .{elapsed});

    server_handle.join();
}