基于liburing库多线程异步IO

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>
#include <arpa/inet.h>
#include <liburing.h>

#define PORT 4567
#define BACKLOG 1024
#define QUEUE_DEPTH 1024
#define MAX_THREADS 4
#define BUFFER_SIZE 4096

// 连接上下文
struct connection {
    int fd;
    char buffer[BUFFER_SIZE];
    size_t buffer_len;
    int thread_id;
};

// 线程上下文
struct thread_context {
    struct io_uring ring;
    int thread_id;
    int server_fd;  // 仅主线程使用
};

// 全局变量
static struct thread_context *threads;
static pthread_barrier_t start_barrier;

// 函数声明
static void *worker_thread(void *arg);
static void submit_accept(struct io_uring *ring, int server_fd);
static void submit_recv(struct io_uring *ring, struct connection *conn);
static void submit_send(struct io_uring *ring, struct connection *conn);

int main() {
    int server_fd, opt = 1;
    struct sockaddr_in addr;
    pthread_t tids[MAX_THREADS];

    // 创建监听 socket
    server_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (server_fd < 0) ERR_EXIT("socket");

    setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    setsockopt(server_fd, SOL_SOCKET, SO_REUSEPORT, &opt, sizeof(opt));

    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = INADDR_ANY;
    addr.sin_port = htons(PORT);

    if (bind(server_fd, (struct sockaddr*)&addr, sizeof(addr)) < 0)
        ERR_EXIT("bind");

    if (listen(server_fd, BACKLOG) < 0)
        ERR_EXIT("listen");

    printf("io_uring TCP Server listening on port %d\n", PORT);

    // 初始化线程上下文
    threads = calloc(MAX_THREADS, sizeof(struct thread_context));
    pthread_barrier_init(&start_barrier, NULL, MAX_THREADS + 1);

    // 创建工作线程
    for (int i = 0; i < MAX_THREADS; i++) {
        threads[i].thread_id = i;
        if (io_uring_queue_init(QUEUE_DEPTH, &threads[i].ring, 0)) {
            fprintf(stderr, "io_uring_queue_init failed for thread %d\n", i);
            exit(1);
        }
        pthread_create(&tids[i], NULL, worker_thread, &threads[i]);
    }

    // 等待所有线程初始化完成
    pthread_barrier_wait(&start_barrier);

    // 主线程：提交第一个 accept
    submit_accept(&threads[0].ring, server_fd);

    // 事件循环：轮询分发新连接
    int next_thread = 0;
    while (1) {
        struct io_uring_cqe *cqe;
        struct connection *conn;
        int ret;

        // 等待 accept 完成
        ret = io_uring_wait_cqe(&threads[0].ring, &cqe);
        if (ret < 0) {
            fprintf(stderr, "io_uring_wait_cqe failed\n");
            break;
        }

        if (cqe->res > 0) {
            int client_fd = cqe->res;
            struct sockaddr_in client_addr;
            socklen_t addr_len = sizeof(client_addr);
            getpeername(client_fd, (struct sockaddr*)&client_addr, &addr_len);

            printf("Accepted %s:%d (fd=%d) → Thread-%d\n",
                   inet_ntoa(client_addr.sin_addr),
                   ntohs(client_addr.sin_port),
                   client_fd, next_thread);

            // 分配连接上下文
            conn = malloc(sizeof(*conn));
            conn->fd = client_fd;
            conn->thread_id = next_thread;

            // 提交到目标工作线程
            submit_recv(&threads[next_thread].ring, conn);

            // 轮询分发
            next_thread = (next_thread + 1) % MAX_THREADS;
        } else {
            fprintf(stderr, "Accept failed: %s\n", strerror(-cqe->res));
        }

        io_uring_cqe_seen(&threads[0].ring, cqe);

        // 提交下一个 accept
        submit_accept(&threads[0].ring, server_fd);
    }

    // 清理
    close(server_fd);
    for (int i = 0; i < MAX_THREADS; i++) {
        io_uring_queue_exit(&threads[i].ring);
    }
    free(threads);
    pthread_barrier_destroy(&start_barrier);
    return 0;
}

// 工作线程：处理 recv/send
static void *worker_thread(void *arg) {
    struct thread_context *tctx = (struct thread_context*)arg;
    struct io_uring *ring = &tctx->ring;
    struct io_uring_cqe *cqe;
    struct connection *conn;

    // 通知主线程已就绪
    pthread_barrier_wait(&start_barrier);

    while (1) {
        if (io_uring_wait_cqe(ring, &cqe) < 0) continue;

        conn = io_uring_cqe_get_data(cqe);
        if (!conn) {
            io_uring_cqe_seen(ring, cqe);
            continue;
        }

        if (cqe->res < 0) {
            fprintf(stderr, "I/O error: %s (fd=%d)\n", strerror(-cqe->res), conn->fd);
            goto close_conn;
        }

        switch (cqe->user_data) {
            case 1: // recv 完成
                printf("Thread-%d: recv %d bytes from fd=%d\n", tctx->thread_id, cqe->res, conn->fd);
                if (cqe->res == 0) {
                    goto close_conn; // EOF
                }
                conn->buffer_len = cqe->res;
                // 回显
                submit_send(ring, conn);
                break;

            case 2: // send 完成
                printf("Thread-%d: sent %d bytes\n", tctx->thread_id, cqe->res);
                // 继续接收
                submit_recv(ring, conn);
                break;

            default:
                break;
        }

        io_uring_cqe_seen(ring, cqe);
        continue;

    close_conn:
        close(conn->fd);
        free(conn);
        io_uring_cqe_seen(ring, cqe);
    }

    return NULL;
}

// 提交 accept 请求
static void submit_accept(struct io_uring *ring, int server_fd) {
    struct io_uring_sqe *sqe = io_uring_get_sqe(ring);
    if (!sqe) {
        fprintf(stderr, "SQ is full!\n");
        return;
    }
    io_uring_prep_accept(sqe, server_fd, NULL, NULL, SOCK_NONBLOCK);
    io_uring_sqe_set_data(sqe, NULL); // accept 无上下文
    io_uring_submit(ring);
}

// 提交 recv 请求
static void submit_recv(struct io_uring *ring, struct connection *conn) {
    struct io_uring_sqe *sqe = io_uring_get_sqe(ring);
    if (!sqe) return;
    io_uring_prep_recv(sqe, conn->fd, conn->buffer, BUFFER_SIZE, 0);
    io_uring_sqe_set_data(sqe, conn);
    sqe->user_data = 1; // 标记为 recv
    io_uring_submit(ring);
}

// 提交 send 请求
static void submit_send(struct io_uring *ring, struct connection *conn) {
    struct io_uring_sqe *sqe = io_uring_get_sqe(ring);
    if (!sqe) return;
    io_uring_prep_send(sqe, conn->fd, conn->buffer, conn->buffer_len, 0);
    io_uring_sqe_set_data(sqe, conn);
    sqe->user_data = 2; // 标记为 send
    io_uring_submit(ring);
}

// 错误处理宏
#define ERR_EXIT(msg) \
    do { perror(msg); exit(EXIT_FAILURE); } while(0)