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         int nblock = 1;
 205
 206         if (bindcpu) {
 207                 int cpu = inst % cpucnt;
 208
 209                 usched_set(getpid(), USCHED_SET_CPU, &cpu, sizeof(cpu));
 210         }
 211
 212         kq = kqueue();
 213         if (kq < 0)
 214                 err(1, "kqueue failed");
 215
 216         nevt_max = nconn_max + 1; /* timer */
 217
 218         evt_change0 = malloc(nevt_max * sizeof(struct kevent));
 219         if (evt_change0 == NULL)
 220                 err(1, "malloc evt_change failed");
 221
 222         evt = malloc(nevt_max * sizeof(struct kevent));
 223         if (evt == NULL)
 224                 err(1, "malloc evt failed");
 225
 226         EV_SET(&evt_change0[0], 0, EVFILT_TIMER, EV_ADD | EV_ONESHOT, 0,
 227             dur * 1000L, NULL);
 228         nchange = 1;
 229
 230         clock_gettime(CLOCK_MONOTONIC_PRECISE, &start);
 231         for (;;) {
 232                 struct kevent *evt_change = NULL;
 233                 int n, i, done = 0;
 234
 235                 while (nconn < nconn_max) {
 236                         const struct sockaddr_in *tmp;
 237                         int s;
 238
 239                         tmp = &in[in_idx % in_cnt];
 240                         ++in_idx;
 241
 242                         if (do_udp)
 243                                 udp_send(tmp);
 244
 245                         s = socket(AF_INET, SOCK_STREAM, 0);
 246                         if (s < 0)
 247                                 err(1, "socket failed");
 248
 249                         if (ioctl(s, FIONBIO, &nblock, sizeof(nblock)) < 0)
 250                                 err(1, "ioctl failed");
 251
 252                         n = connect(s, (const struct sockaddr *)tmp,
 253                             sizeof(*tmp));
 254                         if (n == 0) {
 255                                 ++count;
 256                                 close(s);
 257                                 continue;
 258                         } else {
 259                                 int error = errno;
 260
 261                                 if (error != EINPROGRESS)
 262                                         errc(1, error, "connect failed");
 263                         }
 264                         ++nconn;
 265
 266                         if (nchange >= nevt_max) {
 267                                 fprintf(stderr, "invalid nchange %d, max %d\n",
 268                                     nchange, nevt_max);
 269                                 abort();
 270                         }
 271                         EV_SET(&evt_change0[nchange], s, EVFILT_WRITE, EV_ADD,
 272                             0, 0, NULL);
 273                         ++nchange;
 274                 }
 275
 276                 if (nchange)
 277                         evt_change = evt_change0;
 278
 279                 n = kevent(kq, evt_change, nchange, evt, nevt_max, NULL);
 280                 if (n < 0)
 281                         err(1, "kevent failed");
 282                 nchange = 0;
 283
 284                 for (i = 0; i < n; ++i) {
 285                         struct kevent *e = &evt[i];
 286
 287                         if (e->filter == EVFILT_TIMER) {
 288                                 done = 1;
 289                                 continue;
 290                         }
 291
 292                         if ((e->flags & EV_EOF) && e->fflags) {
 293                                 /* Error, don't count */
 294                         } else {
 295                                 ++count;
 296                         }
 297                         close(e->ident);
 298                         --nconn;
 299                 }
 300                 if (done)
 301                         break;
 302         }
 303         clock_gettime(CLOCK_MONOTONIC_PRECISE, &end);
 304
 305         timespecsub(&end, &start);
 306         time_us = ((double)end.tv_sec * 1000000.0) +
 307             ((double)end.tv_nsec / 1000.0);
 308
 309         *res = ((double)count * 1000000.0) / time_us;
 310 }