6b1957c4aeb5ecda5246ecd98e550b1c5ca11b63
[dragonfly.git] / tools / tools / netrate / accept_connect / kq_connect_client / kq_connect_client.c
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 }