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