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/time.h>
   7 #include <sys/wait.h>
   8
   9 #include <arpa/inet.h>
  10 #include <netinet/in.h>
  11
  12 #include <errno.h>
  13 #include <signal.h>
  14 #include <stdio.h>
  15 #include <stdlib.h>
  16 #include <string.h>
  17 #include <unistd.h>
  18
  19 static void     mainloop(const struct sockaddr_in *, int, int, long, u_long *,
  20                     int);
  21
  22 static void
  23 usage(const char *cmd)
  24 {
  25         fprintf(stderr, "%s -4 inet4 [-4 inet4_1] -p port "
  26             "[-i n_instance] [-c conn_max] [-l duration] [-u]\n", cmd);
  27         exit(1);
  28 }
  29
  30 int
  31 main(int argc, char *argv[])
  32 {
  33         struct sockaddr_in *in, *tmp;
  34         int opt, ninst, nconn, i, in_max, in_cnt, do_udp;
  35         long dur;
  36         u_long *result, sum;
  37         u_short port;
  38
  39         ninst = 1;
  40         nconn = 8;
  41         dur = 10;
  42         port = 0;
  43         do_udp = 0;
  44
  45         in_max = 8;
  46         in = calloc(in_max, sizeof(struct sockaddr_in));
  47         if (in == NULL) {
  48                 fprintf(stderr, "calloc failed\n");
  49                 exit(1);
  50         }
  51         in_cnt = 0;
  52
  53         while ((opt = getopt(argc, argv, "4:p:i:l:c:u")) != -1) {
  54                 switch (opt) {
  55                 case '4':
  56                         if (in_cnt >= in_max) {
  57                                 struct sockaddr_in *old_in = in;
  58                                 int old_in_max = in_max;
  59
  60                                 in_max <<= 1;
  61                                 in = calloc(in_max, sizeof(struct sockaddr_in));
  62                                 if (in == NULL) {
  63                                         fprintf(stderr, "calloc failed\n");
  64                                         exit(1);
  65                                 }
  66
  67                                 memcpy(in, old_in,
  68                                     old_in_max * sizeof(struct sockaddr_in));
  69                                 free(old_in);
  70                         }
  71
  72                         tmp = &in[in_cnt];
  73                         if (inet_pton(AF_INET, optarg, &tmp->sin_addr) <= 0) {
  74                                 fprintf(stderr, "invalid inet address %s\n",
  75                                     optarg);
  76                                 usage(argv[0]);
  77                         }
  78                         ++in_cnt;
  79                         break;
  80
  81                 case 'p':
  82                         port = htons(atoi(optarg));
  83                         break;
  84
  85                 case 'i':
  86                         ninst = atoi(optarg);
  87                         break;
  88
  89                 case 'c':
  90                         nconn = atoi(optarg);
  91                         break;
  92
  93                 case 'l':
  94                         dur = strtol(optarg, NULL, 10);
  95                         break;
  96
  97                 case 'u':
  98                         do_udp = 1;
  99                         break;
 100
 101                 default:
 102                         usage(argv[0]);
 103                 }
 104         }
 105
 106         if (ninst < 1 || dur < 1 || nconn < 1 || port == 0 || in_cnt == 0)
 107                 usage(argv[0]);
 108
 109         for (i = 0; i < in_cnt; ++i) {
 110                 tmp = &in[i];
 111                 tmp->sin_family = AF_INET;
 112                 tmp->sin_port = port;
 113         }
 114
 115         result = mmap(NULL, ninst * sizeof(u_long), PROT_READ | PROT_WRITE,
 116             MAP_ANON | MAP_SHARED, -1, 0);
 117         if (result == MAP_FAILED) {
 118                 fprintf(stderr, "mmap failed: %d\n", errno);
 119                 exit(1);
 120         }
 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                         fprintf(stderr, "fork failed: %d\n", errno);
 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                         fprintf(stderr, "waitpid failed: %d\n", errno);
 141                         exit(1);
 142                 }
 143         }
 144
 145         sum = 0;
 146         for (i = 0; i < ninst; ++i)
 147                 sum += result[i];
 148         printf("%.2f\n", (double)sum / (double)dur);
 149
 150         exit(0);
 151 }
 152
 153 static void
 154 udp_send(const struct sockaddr_in *in)
 155 {
 156         uint8_t d[18];
 157         int s;
 158
 159         s = socket(AF_INET, SOCK_DGRAM, 0);
 160         if (s < 0) {
 161                 fprintf(stderr, "socket DGRAM failed: %d\n", errno);
 162                 return;
 163         }
 164
 165         if (connect(s, (const struct sockaddr *)in, sizeof(*in)) < 0) {
 166                 fprintf(stderr, "connect DGRAM failed: %d\n", errno);
 167                 goto done;
 168         }
 169
 170         write(s, d, sizeof(d));
 171 done:
 172         close(s);
 173 }
 174
 175 static void
 176 mainloop(const struct sockaddr_in *in, int in_cnt, int nconn_max,
 177     long dur, u_long *res, int do_udp)
 178 {
 179         struct kevent *evt_change0, *evt;
 180         int kq, nchange = 0, nconn = 0, nevt_max;
 181         u_long count = 0;
 182         u_int in_idx = 0;
 183         int nblock = 1;
 184
 185         kq = kqueue();
 186         if (kq < 0) {
 187                 fprintf(stderr, "kqueue failed: %d\n", errno);
 188                 return;
 189         }
 190
 191         nevt_max = nconn_max + 1; /* timer */
 192
 193         evt_change0 = malloc(nevt_max * sizeof(struct kevent));
 194         if (evt_change0 == NULL) {
 195                 fprintf(stderr, "malloc evt_change failed\n");
 196                 return;
 197         }
 198
 199         evt = malloc(nevt_max * sizeof(struct kevent));
 200         if (evt == NULL) {
 201                 fprintf(stderr, "malloc evt failed\n");
 202                 return;
 203         }
 204
 205         EV_SET(&evt_change0[0], 0, EVFILT_TIMER, EV_ADD | EV_ONESHOT, 0,
 206             dur * 1000L, NULL);
 207         nchange = 1;
 208
 209         for (;;) {
 210                 struct kevent *evt_change = NULL;
 211                 int n, i, done = 0;
 212
 213                 while (nconn < nconn_max) {
 214                         const struct sockaddr_in *tmp;
 215                         int s;
 216
 217                         tmp = &in[in_idx % in_cnt];
 218                         ++in_idx;
 219
 220                         if (do_udp)
 221                                 udp_send(tmp);
 222
 223                         s = socket(AF_INET, SOCK_STREAM, 0);
 224                         if (s < 0) {
 225                                 fprintf(stderr, "socket failed: %d\n", errno);
 226                                 return;
 227                         }
 228
 229                         if (ioctl(s, FIONBIO, &nblock, sizeof(nblock)) < 0) {
 230                                 fprintf(stderr, "ioctl failed: %d\n", errno);
 231                                 return;
 232                         }
 233
 234                         n = connect(s, (const struct sockaddr *)tmp,
 235                             sizeof(*tmp));
 236                         if (n == 0) {
 237                                 ++count;
 238                                 close(s);
 239                                 continue;
 240                         } else {
 241                                 int error = errno;
 242
 243                                 if (error != EINPROGRESS) {
 244                                         fprintf(stderr, "connect failed: %d\n",
 245                                             error);
 246                                         return;
 247                                 }
 248                         }
 249                         ++nconn;
 250
 251                         if (nchange >= nevt_max) {
 252                                 fprintf(stderr, "invalid nchange %d, max %d\n",
 253                                     nchange, nevt_max);
 254                                 abort();
 255                         }
 256                         EV_SET(&evt_change0[nchange], s, EVFILT_WRITE, EV_ADD,
 257                             0, 0, NULL);
 258                         ++nchange;
 259                 }
 260
 261                 if (nchange)
 262                         evt_change = evt_change0;
 263
 264                 n = kevent(kq, evt_change, nchange, evt, nevt_max, NULL);
 265                 if (n < 0) {
 266                         fprintf(stderr, "kevent failed: %d\n", errno);
 267                         return;
 268                 }
 269                 nchange = 0;
 270
 271                 for (i = 0; i < n; ++i) {
 272                         struct kevent *e = &evt[i];
 273
 274                         if (e->filter == EVFILT_TIMER) {
 275                                 done = 1;
 276                                 continue;
 277                         }
 278
 279                         if ((e->flags & EV_EOF) && e->fflags) {
 280                                 /* Error, don't count */
 281                         } else {
 282                                 ++count;
 283                         }
 284                         close(e->ident);
 285                         --nconn;
 286                 }
 287                 if (done)
 288                         break;
 289         }
 290         *res = count;
 291 }