Linux网络编程之poll

poll多路复用非阻塞模型

the poll is like select,it uses struct pollfd to replace the read_set,write_set,error_set of select,and it has no limitation of the fd number like select.
poll与select很类似，但是用struct pollfd代替了select的read_set,write_set,error_set,而且没有fd数量的限制

struct pollfd
{
int fd; /* File descriptor to poll. */
short int events; /* Types of events poller cares about. */
short int revents; /* Types of events that actually occurred. */
};
poll(<#(struct pollfd*)__fds#>, <#(nfds_t)__nfds#>, <#(int)__timeout#>);

//poll-nonblocking
#include <unistd.h>
#include <stdio.h>
#include <sys/types.h>
#include <stdlib.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <errno.h>
#include <assert.h>
#include <sys/socket.h>
#include <sys/poll.h>
#include <string.h>

#define MAX_READ_SIZE 20*1024*1024
#define MAX_WRITE_SIZE 20*1024*1024
#define MAX_CHUNCK_SIZE 2*1024*1024
#define KEEP_ALIVE 0
/*
 * a simple server using poll 
 *
 * feilengcui008@gmail.com 
 *
 */

void run();

int main(int argc, char *argv[])
{
    run();
    return 0;
}

//struct for socket reading and writing
struct poll_fd_state {
    char read_buf[MAX_READ_SIZE];
    char write_buf[MAX_WRITE_SIZE];
    int writing;
    ssize_t readlen;
    ssize_t written_upto;
    ssize_t write_pos;
};

//error handle
void error_exit()
{
    perror("errors happen");
    exit(EXIT_FAILURE);
}

//create listening socket
int create_server_socket()
{
    int fd;
    int reuse = 1;
    if ((fd= socket(AF_INET, SOCK_STREAM, 0))<0){
        perror("error create socket");
        return fd;
    }
    if (setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse))<0){
        perror("error setsockopt");
    }
    return fd;
}

//alloc poll_fd_state struct
struct poll_fd_state *alloc_poll_fd_state()
{
    //init fds_state
    struct poll_fd_state *p = (struct poll_fd_state *) malloc(sizeof(struct poll_fd_state));
    if(!p){
        return NULL;
    }
    p->readlen = p->written_upto = p->write_pos = p->writing = 0;
    // do not do the following,it will break the code
    //memset(p->read_buf, 0, sizeof(p->read_buf));
    //memset(p->write_buf, 0, sizeof(p->write_buf));
    //bzero(p->read_buf, sizeof(p->read_buf));
    //bzero(p->write_buf, sizeof(p->write_buf));

};

//free poll_fd_state struct
void free_poll_fd_state(struct poll_fd_state *p)
{
    if(p!=NULL){
        free(p);
    }
}

//set socket non-blocking
void made_nonblock(int fd)
{
    if (fcntl(fd, F_SETFL, O_NONBLOCK)<0)
        error_exit();
}

//handle read event
int do_read(struct pollfd *pfd, struct poll_fd_state *state)
{
    ssize_t result=10;
    char buf[MAX_CHUNCK_SIZE];
    while(1){
        result = recv(pfd->fd, buf, sizeof(buf), 0);
        //printf("readlen:%d\n",(int)result);
        //fflush(stdout);
        if (result<=0)
            break;
        for (int i = 0; i < result; ++i) {
            if (state->readlen < sizeof(state->read_buf)){
                state->read_buf[state->readlen++] = buf[i];
            }
            //read until '\n'
        /*
         * todo: handle the readbuffer for http
         *
         */
            /*if (buf[i]=='\n'){
                state->writing = 1;
                state->written_upto = state->readlen;
                pfd->events = POLLOUT;//register write event
                break;
            }*/
        }

    }
    //change state to write
    state->writing = 1;
    state->written_upto = state->readlen;
    pfd->events = POLLOUT;

    if (result==0)
        return 1;
    if (result<0){
        if (errno== EAGAIN)
            return 0;
        else
            return -1;
    }
    return 0;
}

//handle write event
int do_write(struct pollfd *pfd, struct poll_fd_state *state)
{
    ssize_t result;
    while(state->write_pos<state->written_upto){
        result = send(pfd->fd, state->read_buf+state->write_pos, MAX_CHUNCK_SIZE, 0);
        if (result<=0)
            break;
        state->write_pos+=result;
    }

    if (state->write_pos==state->written_upto)
        state->write_pos = state->written_upto = state->readlen = 0;
    state->writing = 0;

    if (result==0)
        return 1;
    if (result<0){
        if (errno== EAGAIN)
            return 0;
        else
            return -1;
    }
    return 0;
}

//poll loop
void do_poll(int serverfd)
{
    int MAX_SOCKET = 1024;
    struct pollfd fds[MAX_SOCKET];
    struct poll_fd_state *fds_state[MAX_SOCKET];

    //init pollfds
    for (int i = 0; i < MAX_SOCKET; ++i) {
        fds[i].fd = -1;
    }

    //add serverfd to pollfds
    fds[0].fd = serverfd;
    fds[0].events = POLLIN;
    fds[0].revents = 0;
    int max = 0;
    int ready = -1;
    //poll begin
    while(1){
        ready = poll(fds, max+1, 0);
        if (ready<0){
            perror("poll error");
            return;
        }
        //listening socket ready to be read?
        if(fds[0].revents & POLLIN==POLLIN){
            struct sockaddr_in client;
            socklen_t slen = sizeof(client);
            memset(&client, 0, sizeof(client));
            int clientfd = accept(fds[0].fd, (struct sockaddr *)&client, &slen);
            made_nonblock(clientfd);
            if (clientfd<0){
                error_exit();
            }
            //find a slot in fds
            int i;
            for (i = 0; i < MAX_SOCKET; ++i) {
                if (fds[i].fd==-1){
                    fds[i].fd=clientfd;
                    fds[i].events= POLLIN;
                    struct poll_fd_state *p = alloc_poll_fd_state();
                    if(p==NULL){
                    error_exit();
                            }
                    fds_state[i] = p;
                    break;
                }
            }
            //reset max
            if(i>max){
                max = i;
            }
        }
        //handle client socks
        for (int j = 1; j < max+1; ++j) {
            int flag=0;
            if (fds[j].fd==-1) continue;
            if ((fds[j].revents & POLLIN)== POLLIN) flag= do_read(&fds[j], fds_state[j]);
            if (flag==0&&((fds[j].revents& POLLOUT)== POLLOUT)){
            flag = do_write(&fds[j], fds_state[j]);
        //if not set keepalive,we close the sock nomatter what flag is
        if (!KEEP_ALIVE){
            close(fds[j].fd);
            fds[j].fd = -1;
            free_poll_fd_state(fds_state[j]);
            fds_state[j] = NULL;
        }
        }
        //handle read-write error
            if (flag){
                close(fds[j].fd);
                fds[j].fd = -1;
                free_poll_fd_state(fds_state[j]);
                fds_state[j] = NULL;
            }
        }
    } 

}

void run()
{

    struct sockaddr_in sin;
    memset(&sin, 0, sizeof(sin));
    sin.sin_family = AF_INET;
    sin.sin_addr.s_addr = 0;
    sin.sin_port = htons(8000);

    int serverfd = create_server_socket();
    if(bind(serverfd, (struct sockaddr *)&sin, sizeof(sin))<0){
        error_exit();
    }

    if (listen(serverfd, 20)<0){
        error_exit();
    }

    made_nonblock(serverfd);
    do_poll(serverfd);
}