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 }