io_uring でエコーサーバー

README.md

io_uring でエコーサーバー

sudo apt update
sudo apt install -y build-essential liburing-dev

gcc -O2 -Wall -Wextra -pedantic -std=c11 -o uring_echo uring_echo.c -luring
./uring_echo 12345

nc 127.0.0.1 12345

何か文字を入力してみる。

server.c

#define _GNU_SOURCE
#include <arpa/inet.h>
#include <errno.h>
#include <liburing.h>
#include <netinet/in.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

#ifndef QUEUE_DEPTH
#define QUEUE_DEPTH  256
#endif

#ifndef BUF_SIZE
#define BUF_SIZE     4096
#endif

static volatile sig_atomic_t g_stop = 0;
static void on_sigint(int signo) { (void)signo; g_stop = 1; }

typedef enum {
    OP_ACCEPT = 1,
    OP_RECV   = 2,
    OP_SEND   = 3,
    OP_CLOSE  = 4,
} op_kind_t;

typedef struct conn {
    int fd;
    char buf[BUF_SIZE];
    size_t len;   // echo する総バイト数
    size_t off;   // すでに送ったオフセット
} conn_t;

typedef struct op {
    op_kind_t kind;
    int listen_fd;
    conn_t *c;
    struct sockaddr_in addr;
    socklen_t addrlen;
} op_t;

static void die(const char *msg) {
    perror(msg);
    exit(1);
}

static int make_listen_socket(uint16_t port) {
    int fd = socket(AF_INET, SOCK_STREAM | SOCK_CLOEXEC, 0);
    if (fd < 0) die("socket");

    int yes = 1;
    if (setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes)) < 0)
        die("setsockopt(SO_REUSEADDR)");

#ifdef SO_REUSEPORT
    if (setsockopt(fd, SOL_SOCKET, SO_REUSEPORT, &yes, sizeof(yes)) < 0)
        die("setsockopt(SO_REUSEPORT)");
#endif

    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = htonl(INADDR_ANY);
    addr.sin_port = htons(port);

    if (bind(fd, (struct sockaddr *)&addr, sizeof(addr)) < 0) die("bind");
    if (listen(fd, 512) < 0) die("listen");
    return fd;
}

static void submit_accept(struct io_uring *ring, int listen_fd) {
    op_t *op = (op_t *)calloc(1, sizeof(op_t));
    if (!op) die("calloc(op accept)");

    op->kind = OP_ACCEPT;
    op->listen_fd = listen_fd;
    op->addrlen = sizeof(op->addr);

    struct io_uring_sqe *sqe = io_uring_get_sqe(ring);
    if (!sqe) die("io_uring_get_sqe(accept)");
    io_uring_prep_accept(sqe, listen_fd, (struct sockaddr *)&op->addr, &op->addrlen, SOCK_CLOEXEC);
    io_uring_sqe_set_data(sqe, op);
}

static void submit_recv(struct io_uring *ring, conn_t *c) {
    op_t *op = (op_t *)calloc(1, sizeof(op_t));
    if (!op) die("calloc(op recv)");
    op->kind = OP_RECV;
    op->c = c;

    struct io_uring_sqe *sqe = io_uring_get_sqe(ring);
    if (!sqe) die("io_uring_get_sqe(recv)");
    io_uring_prep_recv(sqe, c->fd, c->buf, sizeof(c->buf), 0);
    io_uring_sqe_set_data(sqe, op);
}

static void submit_send_remaining(struct io_uring *ring, conn_t *c) {
    op_t *op = (op_t *)calloc(1, sizeof(op_t));
    if (!op) die("calloc(op send)");
    op->kind = OP_SEND;
    op->c = c;

    size_t remain = c->len - c->off;

    struct io_uring_sqe *sqe = io_uring_get_sqe(ring);
    if (!sqe) die("io_uring_get_sqe(send)");
    io_uring_prep_send(sqe, c->fd, c->buf + c->off, remain, 0);
    io_uring_sqe_set_data(sqe, op);
}

static void submit_close(struct io_uring *ring, conn_t *c) {
    op_t *op = (op_t *)calloc(1, sizeof(op_t));
    if (!op) die("calloc(op close)");
    op->kind = OP_CLOSE;
    op->c = c;

    struct io_uring_sqe *sqe = io_uring_get_sqe(ring);
    if (!sqe) die("io_uring_get_sqe(close)");
    io_uring_prep_close(sqe, c->fd);
    io_uring_sqe_set_data(sqe, op);
}

int main(int argc, char **argv) {
    signal(SIGINT, on_sigint);
    signal(SIGPIPE, SIG_IGN); // 重要: 相手切断で落ちないようにする

    uint16_t port = 12345;
    if (argc >= 2) {
        long p = strtol(argv[1], NULL, 10);
        if (p <= 0 || p > 65535) {
            fprintf(stderr, "Invalid port: %s\n", argv[1]);
            return 2;
        }
        port = (uint16_t)p;
    }

    int listen_fd = make_listen_socket(port);
    printf("io_uring echo server listening on 0.0.0.0:%u (Ctrl+C to stop)\n", port);

    struct io_uring ring;
    int ret = io_uring_queue_init(QUEUE_DEPTH, &ring, 0);
    if (ret < 0) {
        fprintf(stderr, "io_uring_queue_init: %s\n", strerror(-ret));
        return 1;
    }

    for (int i = 0; i < 32; i++) submit_accept(&ring, listen_fd);
    ret = io_uring_submit(&ring);
    if (ret < 0) {
        fprintf(stderr, "io_uring_submit: %s\n", strerror(-ret));
        return 1;
    }

    while (!g_stop) {
        struct io_uring_cqe *cqe = NULL;
        ret = io_uring_wait_cqe(&ring, &cqe);
        if (ret < 0) {
            if (ret == -EINTR) continue;
            fprintf(stderr, "io_uring_wait_cqe: %s\n", strerror(-ret));
            break;
        }

        op_t *op = (op_t *)io_uring_cqe_get_data(cqe);
        int res = cqe->res;
        io_uring_cqe_seen(&ring, cqe);

        if (!op) continue;

        switch (op->kind) {
        case OP_ACCEPT: {
            submit_accept(&ring, listen_fd); // accept は常に再投入

            if (res < 0) {
                free(op);
                break;
            }

            conn_t *c = (conn_t *)calloc(1, sizeof(conn_t));
            if (!c) die("calloc(conn)");
            c->fd = res;
            c->len = 0;
            c->off = 0;

            submit_recv(&ring, c);
            free(op);
            break;
        }

        case OP_RECV: {
            conn_t *c = op->c;
            free(op);

            if (res <= 0) {
                submit_close(&ring, c);
                break;
            }

            c->len = (size_t)res;
            c->off = 0;
            submit_send_remaining(&ring, c);
            break;
        }

        case OP_SEND: {
            conn_t *c = op->c;
            free(op);

            if (res < 0) {
                submit_close(&ring, c);
                break;
            }

            c->off += (size_t)res;

            if (c->off < c->len) {
                // まだ送れてない分がある
                submit_send_remaining(&ring, c);
            } else {
                // 送信完了 -> 次の受信へ
                c->len = 0;
                c->off = 0;
                submit_recv(&ring, c);
            }
            break;
        }

        case OP_CLOSE: {
            conn_t *c = op->c;
            free(op);
            free(c);
            break;
        }

        default:
            free(op);
            break;
        }

        ret = io_uring_submit(&ring);
        if (ret < 0) {
            fprintf(stderr, "io_uring_submit: %s\n", strerror(-ret));
            break;
        }
    }

    printf("\nshutting down...\n");
    close(listen_fd);
    io_uring_queue_exit(&ring);
    return 0;
}