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