1 #include <sys/types.h>
2 #include <sys/event.h>
3 #include <sys/mman.h>
4 #include <sys/socket.h>
5 #include <sys/time.h>
6 #include <sys/wait.h>
7 
8 #include <arpa/inet.h>
9 #include <netinet/in.h>
10 
11 #include <errno.h>
12 #include <fcntl.h>
13 #include <signal.h>
14 #include <stdio.h>
15 #include <stdlib.h>
16 #include <string.h>
17 #include <unistd.h>
18 
19 static void	mainloop(const struct sockaddr_in *, int, int, long, u_long *);
20 
21 static void
22 usage(const char *cmd)
23 {
24 	fprintf(stderr, "%s -4 inet4 [-4 inet4_1] -p port "
25 	    "[-i n_instance] [-c conn_max] [-l duration]\n", cmd);
26 	exit(1);
27 }
28 
29 int
30 main(int argc, char *argv[])
31 {
32 	struct sockaddr_in *in, *tmp;
33 	int opt, ninst, nconn, i, in_max, in_cnt;
34 	long dur;
35 	u_long *result, sum;
36 	u_short port;
37 
38 	ninst = 1;
39 	nconn = 8;
40 	dur = 10;
41 	port = 0;
42 
43 	in_max = 8;
44 	in = calloc(in_max, sizeof(struct sockaddr_in));
45 	if (in == NULL) {
46 		fprintf(stderr, "calloc failed\n");
47 		exit(1);
48 	}
49 	in_cnt = 0;
50 
51 	while ((opt = getopt(argc, argv, "4:p:i:l:c:")) != -1) {
52 		switch (opt) {
53 		case '4':
54 			if (in_cnt >= in_max) {
55 				struct sockaddr_in *old_in = in;
56 				int old_in_max = in_max;
57 
58 				in_max <<= 1;
59 				in = calloc(in_max, sizeof(struct sockaddr_in));
60 				if (in == NULL) {
61 					fprintf(stderr, "calloc failed\n");
62 					exit(1);
63 				}
64 
65 				memcpy(in, old_in,
66 				    old_in_max * sizeof(struct sockaddr_in));
67 				free(old_in);
68 			}
69 
70 			tmp = &in[in_cnt];
71 			if (inet_pton(AF_INET, optarg, &tmp->sin_addr) <= 0) {
72 				fprintf(stderr, "invalid inet address %s\n",
73 				    optarg);
74 				usage(argv[0]);
75 			}
76 			++in_cnt;
77 			break;
78 
79 		case 'p':
80 			port = htons(atoi(optarg));
81 			break;
82 
83 		case 'i':
84 			ninst = atoi(optarg);
85 			break;
86 
87 		case 'c':
88 			nconn = atoi(optarg);
89 			break;
90 
91 		case 'l':
92 			dur = strtol(optarg, NULL, 10);
93 			break;
94 
95 		default:
96 			usage(argv[0]);
97 		}
98 	}
99 
100 	if (ninst < 1 || dur < 1 || nconn < 1 || port == 0 || in_cnt == 0)
101 		usage(argv[0]);
102 
103 	for (i = 0; i < in_cnt; ++i) {
104 		tmp = &in[i];
105 		tmp->sin_family = AF_INET;
106 		tmp->sin_port = port;
107 	}
108 
109 	result = mmap(NULL, ninst * sizeof(u_long), PROT_READ | PROT_WRITE,
110 	    MAP_ANON | MAP_SHARED, -1, 0);
111 	if (result == MAP_FAILED) {
112 		fprintf(stderr, "mmap failed: %d\n", errno);
113 		exit(1);
114 	}
115 	memset(result, 0, ninst * sizeof(u_long));
116 
117 	for (i = 0; i < ninst; ++i) {
118 		pid_t pid;
119 
120 		pid = fork();
121 		if (pid == 0) {
122 			mainloop(in, in_cnt, nconn, dur, &result[i]);
123 			exit(0);
124 		} else if (pid < 0) {
125 			fprintf(stderr, "fork failed: %d\n", errno);
126 		}
127 	}
128 
129 	for (i = 0; i < ninst; ++i) {
130 		pid_t pid;
131 
132 		pid = waitpid(-1, NULL, 0);
133 		if (pid < 0) {
134 			fprintf(stderr, "waitpid failed: %d\n", errno);
135 			exit(1);
136 		}
137 	}
138 
139 	sum = 0;
140 	for (i = 0; i < ninst; ++i)
141 		sum += result[i];
142 	printf("%.2f\n", (double)sum / (double)dur);
143 
144 	exit(0);
145 }
146 
147 static void
148 mainloop(const struct sockaddr_in *in, int in_cnt, int nconn_max,
149     long dur, u_long *res)
150 {
151 	struct kevent *evt_change0, *evt;
152 	int kq, nchange = 0, nconn = 0, nevt_max;
153 	u_long count = 0;
154 	u_int in_idx = 0;
155 
156 	kq = kqueue();
157 	if (kq < 0) {
158 		fprintf(stderr, "kqueue failed: %d\n", errno);
159 		return;
160 	}
161 
162 	nevt_max = nconn_max + 1; /* timer */
163 
164 	evt_change0 = malloc(nevt_max * sizeof(struct kevent));
165 	if (evt_change0 == NULL) {
166 		fprintf(stderr, "malloc evt_change failed\n");
167 		return;
168 	}
169 
170 	evt = malloc(nevt_max * sizeof(struct kevent));
171 	if (evt == NULL) {
172 		fprintf(stderr, "malloc evt failed\n");
173 		return;
174 	}
175 
176 	EV_SET(&evt_change0[0], 0, EVFILT_TIMER, EV_ADD | EV_ONESHOT, 0,
177 	    dur * 1000L, NULL);
178 	nchange = 1;
179 
180 	for (;;) {
181 		struct kevent *evt_change = NULL;
182 		int n, i, done = 0;
183 
184 		while (nconn < nconn_max) {
185 			const struct sockaddr_in *tmp;
186 			int s, flags;
187 
188 			s = socket(AF_INET, SOCK_STREAM, 0);
189 			if (s < 0) {
190 				fprintf(stderr, "socket failed: %d\n", errno);
191 				return;
192 			}
193 
194 			flags = fcntl(s, F_GETFL, 0);
195 			if (flags < 0) {
196 				fprintf(stderr, "fcntl(GETFL) failed: %d\n",
197 				    errno);
198 				return;
199 			}
200 			if (fcntl(s, F_SETFL, flags | O_NONBLOCK) < 0) {
201 				fprintf(stderr, "fcntl(SETFL) failed: %d\n",
202 				    errno);
203 				return;
204 			}
205 
206 			tmp = &in[in_idx % in_cnt];
207 			++in_idx;
208 
209 			n = connect(s, (const struct sockaddr *)tmp,
210 			    sizeof(*tmp));
211 			if (n == 0) {
212 				++count;
213 				close(s);
214 				continue;
215 			} else {
216 			    	int error = errno;
217 
218 				if (error != EINPROGRESS) {
219 					fprintf(stderr, "connect failed: %d\n",
220 					    error);
221 					return;
222 				}
223 			}
224 			++nconn;
225 
226 			if (nchange >= nevt_max) {
227 				fprintf(stderr, "invalid nchange %d, max %d\n",
228 				    nchange, nevt_max);
229 				abort();
230 			}
231 			EV_SET(&evt_change0[nchange], s, EVFILT_WRITE, EV_ADD,
232 			    0, 0, NULL);
233 			++nchange;
234 		}
235 
236 		if (nchange)
237 			evt_change = evt_change0;
238 
239 		n = kevent(kq, evt_change, nchange, evt, nevt_max, NULL);
240 		if (n < 0) {
241 			fprintf(stderr, "kevent failed: %d\n", errno);
242 			return;
243 		}
244 		nchange = 0;
245 
246 		for (i = 0; i < n; ++i) {
247 			struct kevent *e = &evt[i];
248 
249 			if (e->filter == EVFILT_TIMER) {
250 				done = 1;
251 				continue;
252 			}
253 
254 			if ((e->flags & EV_EOF) && e->fflags) {
255 				/* Error, don't count */
256 			} else {
257 				++count;
258 			}
259 			close(e->ident);
260 			--nconn;
261 		}
262 		if (done)
263 			break;
264 	}
265 	*res = count;
266 }
267