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