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