多线程-非阻塞模型

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <pthread.h>
#include <errno.h>

#define MAX_SIZE     2048
#define MAX_THREADS  4
#define QUEUE_SIZE   1024

// 线程安全队列
int socket_queue[QUEUE_SIZE];
int queue_front = 0;
int queue_rear = 0;
pthread_mutex_t queue_mutex = PTHREAD_MUTEX_INITIALIZER;
pthread_cond_t queue_cond = PTHREAD_COND_INITIALIZER;

// 线程函数：处理客户端 I/O
void* worker_task(void* arg) {
    printf("Worker thread %ld started.\n", (long)pthread_self());

    while (1) {
        int client_fd = -1;

        // 从队列中安全取出一个 socket
        pthread_mutex_lock(&queue_mutex);
        while (queue_front == queue_rear) {
            pthread_cond_wait(&queue_cond, &queue_mutex);
        }
        client_fd = socket_queue[queue_front];
        queue_front = (queue_front + 1) % QUEUE_SIZE;
        pthread_mutex_unlock(&queue_mutex);

        if (client_fd == -1) continue;

        char buffer[MAX_SIZE];
        ssize_t recv_len;

        // 非阻塞 recv
        while ((recv_len = recv(client_fd, buffer, MAX_SIZE - 1, 0)) > 0) {
            buffer[recv_len] = '\0';
            printf("Thread %ld: recv from %d: %s\n", (long)pthread_self(), client_fd, buffer);

            // 回显
            if (send(client_fd, buffer, recv_len, 0) != recv_len) {
                perror("send");
                break;
            }
        }

        // recv 返回 0：客户端关闭连接
        if (recv_len == 0) {
            printf("Client %d disconnected.\n", client_fd);
        } 
        // recv 返回 -1：错误或 EAGAIN（理论上非阻塞不会阻塞，但可能有其他错误）
        else if (recv_len == -1 && errno != EAGAIN && errno != EWOULDBLOCK) {
            perror("recv error");
        }

        // 关闭连接
        close(client_fd);
        printf("Closed client %d\n", client_fd);
    }

    return NULL;
}

// 向队列添加 socket（线程安全）
int enqueue_socket(int fd) {
    pthread_mutex_lock(&queue_mutex);
    int next_rear = (queue_rear + 1) % QUEUE_SIZE;
    if (next_rear == queue_front) {
        pthread_mutex_unlock(&queue_mutex);
        return -1; // 队列满
    }
    socket_queue[queue_rear] = fd;
    queue_rear = next_rear;
    pthread_cond_signal(&queue_cond);
    pthread_mutex_unlock(&queue_mutex);
    return 0;
}

// 启动服务器
void server_run() {
    unsigned short port = 8000;
    printf("TCP Server started on port %d!\n", port);

    int server_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (server_fd < 0) {
        perror("socket");
        exit(EXIT_FAILURE);
    }

    // 端口复用
    int opt = 1;
    setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

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

    if (bind(server_fd, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
        perror("bind");
        close(server_fd);
        exit(EXIT_FAILURE);
    }

    if (listen(server_fd, 10) < 0) {
        perror("listen");
        close(server_fd);
        exit(EXIT_FAILURE);
    }

    printf("Server listening on port %d, waiting for clients...\n", port);

    // 创建工作线程池
    pthread_t threads[MAX_THREADS];
    for (int i = 0; i < MAX_THREADS; i++) {
        if (pthread_create(&threads[i], NULL, worker_task, NULL) != 0) {
            perror("pthread_create");
            exit(EXIT_FAILURE);
        }
    }

    // 主线程 accept 新连接
    while (1) {
        struct sockaddr_in client_addr;
        socklen_t cli_len = sizeof(client_addr);
        int client_fd = accept(server_fd, (struct sockaddr*)&client_addr, &cli_len);

        if (client_fd < 0) {
            perror("accept");
            continue;
        }

        // 设置为非阻塞
        int flags = fcntl(client_fd, F_GETFL, 0);
        fcntl(client_fd, F_SETFL, flags | O_NONBLOCK);

        char client_ip[INET_ADDRSTRLEN];
        inet_ntop(AF_INET, &client_addr.sin_addr, client_ip, INET_ADDRSTRLEN);
        printf("New connection from %s:%d (fd=%d)\n", client_ip, ntohs(client_addr.sin_port), client_fd);

        // 加入队列
        if (enqueue_socket(client_fd) == -1) {
            printf("Queue full, closing client %d\n", client_fd);
            close(client_fd);
        }
    }

    close(server_fd);
}

int main() {
    server_run();
    return 0;
}