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 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 }