epoll实现压测工具

代码：

 /*
  * g++ -o stress_test ./stress_test.cpp
  */

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

 #include <unistd.h>
 #include <sys/types.h>
 #include <sys/epoll.h>
 #include <fcntl.h>
 #include <sys/socket.h>
 #include <netinet/in.h>
 #include <arpa/inet.h>

 static const char *request = "GET / HTTP/1.1\r\nHost: 192.168.1.5\r\nConnection: keep-alive\r\n\r\n";

 int setnonblocking(int fd)
 {
     int old_option = fcntl(fd, F_GETFL);
     int new_option = old_option | O_NONBLOCK;
     fcntl(fd, F_SETFL, new_option);

     return old_option;
 }

 int add_fd(int epoll_fd, int fd)
 {
     struct epoll_event event;
     event.events = EPOLLOUT | EPOLLET | EPOLLERR;
     event.data.fd = fd;
     epoll_ctl(epoll_fd, EPOLL_CTL_ADD, fd, &event);
     setnonblocking(fd);
 }

 bool write_nbytes(int sockfd, const char *buffer, int len)
 {
     int bytes_write = ;
     printf("write out %d bytes to socket %d\n", len, sockfd);
     while()
     {
         bytes_write = send(sockfd, buffer, len, );
         if (bytes_write == -)
         {
             printf("send failed, errno[%d], error[%s]\n", errno, strerror(errno));
             return false;
         }
         else if (bytes_write == )
         {
             printf("send 0 bytes\n");
             return false;
         }

         len -= bytes_write;
         buffer = buffer + bytes_write;
         if (len <= )
         {
             return true;
         }
     }
 }

 bool read_once(int sockfd, char *buffer, int len)
 {
     int bytes_read = ;
     memset(buffer, '\0', len);
     bytes_read = recv(sockfd, buffer, len, );
     if (bytes_read == -)
     {
         printf("recv failed, errno[%d], error[%s]\n", errno, strerror(errno));
         return false;
     }
     else if (bytes_read == )
     {
         printf("recv 0 bytes\n");
         return false;
     }
     printf("read in %d bytes from socket %d\n", bytes_read, sockfd);

     return true;
 }

 int start_conn(int epoll_fd, const char *ip, int port, int num)
 {
     struct sockaddr_in addr;
     bzero(&addr, sizeof(addr));
     addr.sin_family = AF_INET;
     addr.sin_addr.s_addr = inet_addr(ip);
     addr.sin_port = htons(port);

     socklen_t len = sizeof(addr);
     for (int i = ; i < num; ++i)
     {
         sleep();
         int fd = socket(AF_INET, SOCK_STREAM, );
         if (fd == -)
         {
             printf("socket failed, errno[%d], error[%s]\n", errno, strerror(errno));
             continue;
         }

         if (connect(fd, (struct sockaddr*)&addr, len) == )
         {
             printf("build connection %d\n", i);
             add_fd(epoll_fd, fd);
         }
         else
         {
             printf("connect failed, errno[%d], error[%s]\n", errno, strerror(errno));
             continue;
         }
     }
 }

 int close_conn(int epoll_fd, int fd)
 {
     epoll_ctl(epoll_fd, EPOLL_CTL_DEL, fd, NULL);

     printf("close socket[%d]\n", fd);
     close(fd);
 }

 int main(int argc, char **argv)
 {
     if (argc != )
     {
         printf("usage:stress_test ip port num\n");
         exit(-);
     }

     int epoll_fd = epoll_create();
     start_conn(epoll_fd, argv[], atoi(argv[]), atoi(argv[]));
     epoll_event events[];
     char buffer[];
     while ()
     {
         int fds = epoll_wait(epoll_fd, events, , );
         for (int i = ; i < fds; ++i)
         {
             int sockfd = events[i].data.fd;
             if (events[i].events & EPOLLIN)
             {
                 if (!read_once(sockfd, buffer, ))
                 {
                     close_conn(epoll_fd, sockfd);
                 }

                 struct epoll_event event;
                 event.events = EPOLLOUT | EPOLLET | EPOLLERR;
                 event.data.fd = sockfd;
                 epoll_ctl(epoll_fd, EPOLL_CTL_MOD, sockfd, &event);
             }
             else if (events[i].events & EPOLLOUT)
             {
                 if (!write_nbytes(sockfd, request, strlen(request)))
                 {
                     close_conn(epoll_fd, sockfd);
                 }

                 struct epoll_event event;
                 event.events = EPOLLIN | EPOLLET | EPOLLERR;
                 event.data.fd = sockfd;
                 epoll_ctl(epoll_fd, EPOLL_CTL_MOD, sockfd, &event);
             }
             else if (events[i].events & EPOLLERR)
             {
                 close_conn(epoll_fd, sockfd);
             }
         }
     }

     return ;
 }