1 /*
  2  * g++ -o stress_test ./stress_test.cpp
  3  */
  4
  5 #include <stdlib.h>
  6 #include <stdio.h>
  7 #include <string.h>
  8 #include <errno.h>
  9
 10 #include <unistd.h>
 11 #include <sys/types.h>
 12 #include <sys/epoll.h>
 13 #include <fcntl.h>
 14 #include <sys/socket.h>
 15 #include <netinet/in.h>
 16 #include <arpa/inet.h>
 17
 18 static const char *request = "GET / HTTP/1.1\r\nHost: 192.168.1.5\r\nConnection: keep-alive\r\n\r\n";
 19
 20 int setnonblocking(int fd)
 21 {
 22     int old_option = fcntl(fd, F_GETFL);
 23     int new_option = old_option | O_NONBLOCK;
 24     fcntl(fd, F_SETFL, new_option);
 25
 26     return old_option;
 27 }
 28
 29 int add_fd(int epoll_fd, int fd)
 30 {
 31     struct epoll_event event;
 32     event.events = EPOLLOUT | EPOLLET | EPOLLERR;
 33     event.data.fd = fd;
 34     epoll_ctl(epoll_fd, EPOLL_CTL_ADD, fd, &event);
 35     setnonblocking(fd);
 36 }
 37
 38 bool write_nbytes(int sockfd, const char *buffer, int len)
 39 {
 40     int bytes_write = 0;
 41     printf("write out %d bytes to socket %d\n", len, sockfd);
 42     while(1)
 43     {
 44         bytes_write = send(sockfd, buffer, len, 0);
 45         if (bytes_write == -1)
 46         {
 47             printf("send failed, errno[%d], error[%s]\n", errno, strerror(errno));
 48             return false;
 49         }
 50         else if (bytes_write == 0)
 51         {
 52             printf("send 0 bytes\n");
 53             return false;
 54         }
 55
 56         len -= bytes_write;
 57         buffer = buffer + bytes_write;
 58         if (len <= 0)
 59         {
 60             return true;
 61         }
 62     }
 63 }
 64
 65 bool read_once(int sockfd, char *buffer, int len)
 66 {
 67     int bytes_read = 0;
 68     memset(buffer, '\0', len);
 69     bytes_read = recv(sockfd, buffer, len, 0);
 70     if (bytes_read == -1)
 71     {
 72         printf("recv failed, errno[%d], error[%s]\n", errno, strerror(errno));
 73         return false;
 74     }
 75     else if (bytes_read == 0)
 76     {
 77         printf("recv 0 bytes\n");
 78         return false;
 79     }
 80     printf("read in %d bytes from socket %d\n", bytes_read, sockfd);
 81
 82     return true;
 83 }
 84
 85 int start_conn(int epoll_fd, const char *ip, int port, int num)
 86 {
 87     struct sockaddr_in addr;
 88     bzero(&addr, sizeof(addr));
 89     addr.sin_family = AF_INET;
 90     addr.sin_addr.s_addr = inet_addr(ip);
 91     addr.sin_port = htons(port);
 92
 93     socklen_t len = sizeof(addr);
 94     for (int i = 0; i < num; ++i)
 95     {
 96         sleep(1);
 97         int fd = socket(AF_INET, SOCK_STREAM, 0);
 98         if (fd == -1)
 99         {
100             printf("socket failed, errno[%d], error[%s]\n", errno, strerror(errno));
101             continue;
102         }
103
104         if (connect(fd, (struct sockaddr*)&addr, len) == 0)
105         {
106             printf("build connection %d\n", i);
107             add_fd(epoll_fd, fd);
108         }
109         else
110         {
111             printf("connect failed, errno[%d], error[%s]\n", errno, strerror(errno));
112             continue;
113         }
114     }
115 }
116
117 int close_conn(int epoll_fd, int fd)
118 {
119     epoll_ctl(epoll_fd, EPOLL_CTL_DEL, fd, NULL);
120
121     printf("close socket[%d]\n", fd);
122     close(fd);
123 }
124
125 int main(int argc, char **argv)
126 {
127     if (argc != 4)
128     {
129         printf("usage:stress_test ip port num\n");
130         exit(-1);
131     }
132
133     int epoll_fd = epoll_create(10000);
134     start_conn(epoll_fd, argv[1], atoi(argv[2]), atoi(argv[3]));
135     epoll_event events[10000];
136     char buffer[2048];
137     while (1)
138     {
139         int fds = epoll_wait(epoll_fd, events, 10000, 2000);
140         for (int i = 0; i < fds; ++i)
141         {
142             int sockfd = events[i].data.fd;
143             if (events[i].events & EPOLLIN)
144             {
145                 if (!read_once(sockfd, buffer, 2048))
146                 {
147                     close_conn(epoll_fd, sockfd);
148                 }
149
150                 struct epoll_event event;
151                 event.events = EPOLLOUT | EPOLLET | EPOLLERR;
152                 event.data.fd = sockfd;
153                 epoll_ctl(epoll_fd, EPOLL_CTL_MOD, sockfd, &event);
154             }
155             else if (events[i].events & EPOLLOUT)
156             {
157                 if (!write_nbytes(sockfd, request, strlen(request)))
158                 {
159                     close_conn(epoll_fd, sockfd);
160                 }
161
162                 struct epoll_event event;
163                 event.events = EPOLLIN | EPOLLET | EPOLLERR;
164                 event.data.fd = sockfd;
165                 epoll_ctl(epoll_fd, EPOLL_CTL_MOD, sockfd, &event);
166             }
167             else if (events[i].events & EPOLLERR)
168             {
169                 close_conn(epoll_fd, sockfd);
170             }
171         }
172     }
173
174     return 0;
175 }