高性能TCP 服务器 （基于memcached源码修改）

背景

银行前置系统面临双十一支付压力，需对接入tcp socket进行性能升级改造。
改造前socket 处理方式：预启动N个子进程，父进程accept 连接后通过管道将socket 发送给子进程进行处理，性能较差。首先想到基于libevent来进行改造。之前项目使用memcached，其通信是基于libevent，就有了从memcached源码改造这个服务端程序。

程序

注意：程序运行编译需要先安装libevent

tcp.c
/* 
 *  Use and distribution licensed under the BSD license.  See
 *  the LICENSE file for full text.
 *  Authors:
 */
#include "tcp.h"
#include <sys/stat.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <signal.h>
#include <sys/param.h>
#include <sys/resource.h>
#include <sys/uio.h>
#include <ctype.h>
#include <stdarg.h>
#include <stdlib.h>
#include <stdio.h>
//注册用户消息处理函数范例， 通过user_process(event_handler) 设置处理函数
void event_handler(const int fd, const short which, void *arg) {
    myconn *c;
    int n;
    char buffer[128];
    char readbuffer[128];
    c = (myconn *)arg;
        
    assert(c != NULL);
        
    if (fd != c->sfd) {
        fprintf(stderr, "event_handler sanity\n");
        close(fd);
        return;
    }
    
    memset(readbuffer, 0, sizeof(readbuffer));
            
    n = read(fd, readbuffer, sizeof(readbuffer));
    if(n <= 0) {
        user_free(c);
        return;
    }
    
    memset(buffer, 0, sizeof(buffer));
        
    sprintf(buffer, "[sfd:%04d|thread_id:%08d %08d]", c->sfd , c->thread->thread_id,           pthread_self());
    printf( "read buffer [%s] write buffer %s", readbuffer, buffer);
        
    n = write(fd, buffer, strlen(buffer));
    if(n <= 0) {
        user_free(c);
        fprintf(stderr, "free fd[%d] gently\n", fd);
        return;
    }
    return;
}

static void sig_handler(const int sig) {
    printf("Signal handled: %s.\n", strsignal(sig));
    exit(EXIT_SUCCESS);
}
 
//主程序
int main (int argc, char **argv) {
    int sfd;
    int ret;
    signal(SIGINT, sig_handler);
    signal(SIGTERM, sig_handler);
    sfd = user_init(11211);//端口
    if(sfd < 0) {
        perror("failed to call user_init");
        exit(0);
    }
    user_process(event_handler);
    user_loop(sfd);
    return 0;
}

tcp.h
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <netinet/in.h>
#include <event.h>
#include <netdb.h>
#include <pthread.h>
#include <unistd.h>
#include <assert.h>
#include <grp.h>

struct thread_stats {
    pthread_mutex_t   mutex;
};

typedef struct {
    pthread_t thread_id;        /* unique ID of this thread */
    struct event_base *base;    /* libevent handle this thread uses */
    struct event notify_event;  /* listen event for notify pipe */
    int notify_receive_fd;      /* receiving end of notify pipe */
    int notify_send_fd;         /* sending end of notify pipe */
    struct thread_stats stats;  /* Stats generated by this thread */
    struct conn_queue *new_conn_queue; /* queue of new connections to handle */
    void *lru_bump_buf;         /* async LRU bump buffer */
} LIBEVENT_THREAD;

enum pause_thread_types {
    PAUSE_WORKER_THREADS = 0,
    PAUSE_ALL_THREADS,
    RESUME_ALL_THREADS,
    RESUME_WORKER_THREADS
};

typedef struct _myconn myconn;

struct _myconn {
    int    sfd;
    LIBEVENT_THREAD *thread; /* Pointer to the thread object serving this connection */
    struct event event;
} ;

typedef void (*handler)(const int fd, const short which, void *arg) ;

void user_loop(const int sfd);
int  user_init(short port);
void user_process(void (*func)(const int fd, const short which, void *arg));
void user_free(myconn *c);

thread.c :从memcached 修改而来
/* -*- Mode: C; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
/*
 */
#include "tcp.h"

#include <assert.h>
#include <stdio.h>
#include <errno.h>
#include <stdlib.h>
#include <string.h>
#include <pthread.h>


#define ITEMS_PER_ALLOC 64

static struct event_base *main_base;
static handler __g_fun;

/* Which thread we assigned a connection to most recently. */
static int last_thread = -1;

static int handler_flag = 0;

typedef struct conn_queue_item CQ_ITEM;
struct conn_queue_item {
    int               sfd;
    myconn *c;
    CQ_ITEM          *next;
};

/* A connection queue. */
typedef struct conn_queue CQ;
struct conn_queue {
    CQ_ITEM *head;
    CQ_ITEM *tail;
    pthread_mutex_t lock;
};

/* Free list of CQ_ITEM structs */
static CQ_ITEM *cqi_freelist;
static pthread_mutex_t cqi_freelist_lock;

/*
 * Each libevent instance has a wakeup pipe, which other threads
 * can use to signal that they've put a new connection on its queue.
 */
static LIBEVENT_THREAD *threads;

/*
 * Number of worker threads that have finished setting themselves up.
 */
static int init_count = 0;
static pthread_mutex_t init_lock;
static pthread_cond_t init_cond;

/* item_lock() must be held for an item before any modifications to either its
 * associated hash bucket, or the structure itself.
 * LRU modifications must hold the item lock, and the LRU lock.
 * LRU's accessing items must item_trylock() before modifying an item.
 * Items accessible from an LRU must not be freed or modified
 * without first locking and removing from the LRU.
 */

static void wait_for_thread_registration(int nthreads) {
    while (init_count < nthreads) {
        pthread_cond_wait(&init_cond, &init_lock);
    }
}

static void register_thread_initialized(void) {
    pthread_mutex_lock(&init_lock);
    init_count++;
    pthread_cond_signal(&init_cond);
    pthread_mutex_unlock(&init_lock);
}

/*
 * Initializes a connection queue.
 */
static void cq_init(CQ *cq) {
    pthread_mutex_init(&cq->lock, NULL);
    cq->head = NULL;
    cq->tail = NULL;
}

/*
 * Looks for an item on a connection queue, but doesn't block if there isn't
 * one.
 * Returns the item, or NULL if no item is available
 */
static CQ_ITEM *cq_pop(CQ *cq) {
    CQ_ITEM *item;

    pthread_mutex_lock(&cq->lock);
    item = cq->head;
    if (NULL != item) {
        cq->head = item->next;
        if (NULL == cq->head)
            cq->tail = NULL;
    }
    pthread_mutex_unlock(&cq->lock);

    return item;
}

/*
 * Adds an item to a connection queue.
 */
static void cq_push(CQ *cq, CQ_ITEM *item) {
    item->next = NULL;

    pthread_mutex_lock(&cq->lock);
    if (NULL == cq->tail)
        cq->head = item;
    else
        cq->tail->next = item;
    cq->tail = item;
    pthread_mutex_unlock(&cq->lock);
}

/*
 * Returns a fresh connection queue item.
 */
static CQ_ITEM *cqi_new(void) {
    CQ_ITEM *item = NULL;
    pthread_mutex_lock(&cqi_freelist_lock);
    if (cqi_freelist) {
        item = cqi_freelist;
        cqi_freelist = item->next;
    }
    pthread_mutex_unlock(&cqi_freelist_lock);

    if (NULL == item) {
        int i;

        /* Allocate a bunch of items at once to reduce fragmentation */
        item = malloc(sizeof(CQ_ITEM) * ITEMS_PER_ALLOC);
        if (NULL == item) {
            return NULL;
        }

        /*
         * Link together all the new items except the first one
         * (which we'll return to the caller) for placement on
         * the freelist.
         */
        for (i = 2; i < ITEMS_PER_ALLOC; i++){
            item[i - 1].next = &item[i];
        }

        pthread_mutex_lock(&cqi_freelist_lock);
        item[ITEMS_PER_ALLOC - 1].next = cqi_freelist;
        cqi_freelist = &item[1];
        pthread_mutex_unlock(&cqi_freelist_lock);
    }

    return item;
}


/*
 * Frees a connection queue item (adds it to the freelist.)
 */
static void cqi_free(CQ_ITEM *item) {
    pthread_mutex_lock(&cqi_freelist_lock);
    item->next = cqi_freelist;
    cqi_freelist = item;
    pthread_mutex_unlock(&cqi_freelist_lock);
}


/*
 * Creates a worker thread.
 */
static void create_worker(void *(*func)(void *), void *arg) {
    pthread_attr_t  attr;
    int             ret;

    pthread_attr_init(&attr);

    if ((ret = pthread_create(&((LIBEVENT_THREAD*)arg)->thread_id, &attr, func, arg)) != 0) {
        fprintf(stderr, "Can't create thread: %s\n",
                strerror(ret));
        exit(1);
    }
}


/*
 * Processes an incoming "handle a new connection" item. This is called when
 * input arrives on the libevent wakeup pipe.
 */
static void thread_libevent_process(int fd, short which, void *arg) {
    LIBEVENT_THREAD *me = arg;
    CQ_ITEM *item;
    char buf[1];
    myconn *c;
    unsigned int timeout_fd;
    char buffer[128];
    
    if (read(fd, buf, 1) != 1) {
        fprintf(stderr, "Can't read from libevent pipe\n");
        return;
    }

    if (buf[0] == 'c') {

        item = cq_pop(me->new_conn_queue);

        if (NULL == item) {
            return;
        }

        c = item->c;
        
        c->thread = me;
        
        event_set(&c->event, item->sfd, EV_READ|EV_PERSIST, __g_fun, (void *)c);
        
        event_base_set(me->base, &c->event);

        if (event_add(&c->event, 0) == -1) {
            perror("event_add");
            return ;
        }
        
        cqi_free(item);
    }
    
}


/*
 * Set up a thread's information.
 */
static void setup_thread(LIBEVENT_THREAD *me) {

    me->base = event_init();

    if (! me->base) {
        fprintf(stderr, "Can't allocate event base\n");
        exit(1);
    }

    /* Listen for notifications from other threads */
    event_set(&me->notify_event, me->notify_receive_fd,
              EV_READ | EV_PERSIST, thread_libevent_process, me);
    
    event_base_set(me->base, &me->notify_event);

    if (event_add(&me->notify_event, 0) == -1) {
        fprintf(stderr, "Can't monitor libevent notify pipe\n");
        exit(1);
    }

    me->new_conn_queue = malloc(sizeof(struct conn_queue));
    if (me->new_conn_queue == NULL) {
        perror("Failed to allocate memory for connection queue");
        exit(EXIT_FAILURE);
    }
    
    cq_init(me->new_conn_queue);

}

/*
 * Worker thread: main event loop
 */
static void *worker_libevent(void *arg) {
    LIBEVENT_THREAD *me = arg;

    register_thread_initialized();

    event_base_loop(me->base, 0);

    event_base_free(me->base);
    
    return NULL;
}

/*
 * Dispatches a new connection to another thread. This is only ever called
 * from the main thread, either during initialization (for UDP) or because
 * of an incoming connection.
 */
static void dispatch_conn_new(int sfd,  myconn *c) {
    CQ_ITEM *item = cqi_new();
    
    char buf[1];
    
    if (item == NULL) {
        close(sfd);
        /* given that malloc failed this may also fail, but let's try */
        fprintf(stderr, "Failed to allocate memory for connection object\n");
        free(c);
        return ;
    }
 
    int tid = (last_thread + 1) % 4;

    LIBEVENT_THREAD *thread = threads + tid;

    last_thread = tid;

    item->sfd = sfd;
    item->c = c;
    
    cq_push(thread->new_conn_queue, item);
    
    buf[0] = 'c';
    
    if (write(thread->notify_send_fd, buf, 1) != 1) {
        perror("Writing to thread notify pipe");
    }
    
}

/*
 * Initializes the thread subsystem, creating various worker threads.
 *
 * nthreads  Number of worker event handler threads to spawn
 */
static void event_thread_init(int nthreads) {
    int  i;

    pthread_mutex_init(&init_lock, NULL);
    pthread_cond_init(&init_cond, NULL);

    threads = calloc(nthreads, sizeof(LIBEVENT_THREAD));
    if (! threads) {
        perror("Can't allocate thread descriptors");
        exit(1);
    }

    for (i = 0; i < nthreads; i++) {
        int fds[2];
        if (pipe(fds)) {
            perror("Can't create notify pipe");
            exit(1);
        }

        threads[i].notify_receive_fd = fds[0];
        threads[i].notify_send_fd = fds[1];

        setup_thread(&threads[i]);
    }

    /* Create threads after we've done all the libevent setup. */
    for (i = 0; i < nthreads; i++) {
        create_worker(worker_libevent, &threads[i]);
    }

    /* Wait for all the threads to set themselves up before returning. */
    pthread_mutex_lock(&init_lock);
    wait_for_thread_registration(nthreads);
    pthread_mutex_unlock(&init_lock);
}

static void __event_handler(const int fd, const short which, void *arg) {
    myconn *c;
    socklen_t addrlen;
    struct sockaddr_storage addr;
    int sfd;

    assert(fd >= 0);
    
    addrlen = sizeof(addr);
    sfd = accept(fd, (struct sockaddr *)&addr, &addrlen);

    if (sfd == -1) {
        fprintf(stderr, "accept error: [%s]\n", strerror(errno));
        return;
    }
    
    if (!(c = (myconn *)calloc(1, sizeof(myconn)))) {
            fprintf(stderr, "Failed to allocate connection object\n");
            return ;
    }

    c->sfd = sfd;
    
    dispatch_conn_new(sfd, c);

    return;
}

#ifndef HAVE_SIGIGNORE
static int sigignore(int sig) {
    struct sigaction sa = { .sa_handler = SIG_IGN, .sa_flags = 0 };

    if (sigemptyset(&sa.sa_mask) == -1 || sigaction(sig, &sa, 0) == -1) {
        return -1;
    }
    return 0;
}
#endif

/**************************
 * for user call     
 **************************/
void user_loop(const int sfd) {
    struct event* ev;

    if(handler_flag == 0)
    {
        perror("set_handler had not called! \n");
        return ;
    }
    
    assert(sfd >= 0);

    ev = event_new(main_base, sfd, EV_READ | EV_PERSIST, __event_handler, NULL);
    
    event_base_set(main_base, ev);

    if (event_add(ev, 0) == -1) {
        perror("event_add");
        return ;
    }

    fprintf(stdout,"user_loop sucess...\n");
    
    /* enter the event loop */
    event_base_loop(main_base, 0) ;

    event_free(ev);

    /* cleanup base */
    event_base_free(main_base);
    
    return ;
}

int user_init(short port) {
    int sfd;
    struct linger ling = {0, 0};
    struct sockaddr_in server_addr;
    int flags = 1;

    if (sigignore(SIGPIPE) == -1) {
        perror("failed to ignore SIGPIPE; sigaction");
        exit(0);
    }
    
    if((sfd = socket(AF_INET, SOCK_STREAM, 0))==-1)
    {
        fprintf(stderr,"Socket error:%s\n\a",strerror(errno));
        return -1;
    }
    
    bzero(&server_addr,sizeof(struct sockaddr_in));
    
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = htonl(INADDR_ANY);
    server_addr.sin_port = htons(11211);

    setsockopt(sfd, SOL_SOCKET, SO_REUSEADDR, (void *)&flags, sizeof(flags));
    setsockopt(sfd, SOL_SOCKET, SO_KEEPALIVE, (void *)&flags, sizeof(flags));
    setsockopt(sfd, SOL_SOCKET, SO_LINGER, (void *)&ling, sizeof(ling));
    setsockopt(sfd, IPPROTO_TCP, TCP_NODELAY, (void *)&flags, sizeof(flags));

    if (bind(sfd, (struct sockaddr *)(&server_addr), sizeof(struct sockaddr)) == -1) {
        close(sfd);
        return -1;
    }
    if (listen(sfd, 10) == -1) {
        perror("listen()");
        close(sfd);
        return -1;
    }

    event_thread_init(4);

    main_base = event_init();
    if(NULL == main_base){
        perror("failed to call event_init");
        return -1;
    }

    fprintf(stdout,"user_init sucess...\n");
    
    return sfd;
}

void user_process(void (*func)(const int fd, const short which, void *arg))
{
    __g_fun = func;
    handler_flag = 1;
    fprintf(stdout,"user_process sucess...\n");
}

void user_free(myconn *c)
{
    fprintf(stderr, "user_free fd[%d] gently\n", c->sfd);
    event_del(&c->event);
    close(c->sfd);
    free(c);
}

测试

从网上下载的TCP_UDP_PerformanceTest作为tcp发包工具。

1. 服务端IBMT450 window 10 基于cygwin测试情况:2048 长连接，运行OK。
   

2. 服务端AIX 10C 512G 10000 长连接，运行OK。