tools/tools/netrate/accept_connect/kq_accept_server/kq_accept_server.c

   1 #include <sys/types.h>
   2 #include <sys/event.h>
   3 #include <sys/ioctl.h>
   4 #include <sys/socket.h>
   5 #include <sys/sysctl.h>
   6 #include <sys/time.h>
   7
   8 #include <arpa/inet.h>
   9 #include <netinet/in.h>
  10
  11 #include <errno.h>
  12 #include <stdio.h>
  13 #include <stdlib.h>
  14 #include <string.h>
  15 #include <unistd.h>
  16
  17 #define EVENT_MAX       128
  18
  19 static void     mainloop(int, const struct sockaddr_in *);
  20
  21 static void
  22 usage(const char *cmd)
  23 {
  24         fprintf(stderr, "%s -p port [-i n_instance] [-r]\n", cmd);
  25         exit(1);
  26 }
  27
  28 static int
  29 create_socket(const struct sockaddr_in *in, int reuseport)
  30 {
  31         int serv_s, on;
  32
  33         serv_s = socket(AF_INET, SOCK_STREAM, 0);
  34         if (serv_s < 0) {
  35                 fprintf(stderr, "socket failed: %d\n", errno);
  36                 exit(1);
  37         }
  38
  39         on = 1;
  40         if (!reuseport) {
  41                 if (setsockopt(serv_s, SOL_SOCKET, SO_REUSEADDR,
  42                     &on, sizeof(on)) < 0) {
  43                         fprintf(stderr, "setsockopt(REUSEADDR) failed: %d\n",
  44                             errno);
  45                         exit(1);
  46                 }
  47         } else {
  48                 if (setsockopt(serv_s, SOL_SOCKET, SO_REUSEPORT,
  49                     &on, sizeof(on)) < 0) {
  50                         fprintf(stderr, "setsockopt(REUSEPORT) failed: %d\n",
  51                             errno);
  52                         exit(1);
  53                 }
  54         }
  55
  56         on = 1;
  57         if (ioctl(serv_s, FIONBIO, &on, sizeof(on)) < 0) {
  58                 fprintf(stderr, "ioctl(FIONBIO) failed: %d\n", errno);
  59                 exit(1);
  60         }
  61
  62         if (bind(serv_s, (const struct sockaddr *)in, sizeof(*in)) < 0) {
  63                 fprintf(stderr, "bind failed: %d\n", errno);
  64                 exit(1);
  65         }
  66
  67         if (listen(serv_s, -1) < 0) {
  68                 fprintf(stderr, "listen failed: %d\n", errno);
  69                 exit(1);
  70         }
  71         return serv_s;
  72 }
  73
  74 int
  75 main(int argc, char *argv[])
  76 {
  77         struct sockaddr_in in;
  78         int opt, ninst, serv_s, i, reuseport;
  79         size_t prm_len;
  80
  81         prm_len = sizeof(ninst);
  82         if (sysctlbyname("hw.ncpu", &ninst, &prm_len, NULL, 0) != 0) {
  83                 fprintf(stderr, "sysctl hw.ncpu failed: %d\n", errno);
  84                 exit(2);
  85         }
  86
  87         memset(&in, 0, sizeof(in));
  88         in.sin_family = AF_INET;
  89         in.sin_addr.s_addr = INADDR_ANY;
  90
  91         reuseport = 0;
  92
  93         while ((opt = getopt(argc, argv, "p:i:r")) != -1) {
  94                 switch (opt) {
  95                 case 'p':
  96                         in.sin_port = htons(atoi(optarg));
  97                         break;
  98
  99                 case 'i':
 100                         ninst = atoi(optarg);
 101                         break;
 102
 103                 case 'r':
 104                         reuseport = 1;
 105                         break;
 106
 107                 default:
 108                         usage(argv[0]);
 109                 }
 110         }
 111
 112         if (ninst < 1 || in.sin_port == 0)
 113                 usage(argv[0]);
 114
 115         serv_s = -1;
 116         if (!reuseport)
 117                 serv_s = create_socket(&in, 0);
 118
 119         for (i = 1; i < ninst; ++i) {
 120                 pid_t pid;
 121
 122                 pid = fork();
 123                 if (pid == 0) {
 124                         mainloop(serv_s, &in);
 125                         exit(0);
 126                 } else if (pid < 0) {
 127                         fprintf(stderr, "fork failed: %d\n", errno);
 128                 }
 129         }
 130
 131         mainloop(serv_s, &in);
 132         exit(0);
 133 }
 134
 135 static void
 136 mainloop(int serv_s, const struct sockaddr_in *in)
 137 {
 138         struct kevent change_evt0[EVENT_MAX];
 139         int kq, nchange;
 140
 141         if (serv_s < 0)
 142                 serv_s = create_socket(in, 1);
 143
 144         kq = kqueue();
 145         if (kq < 0) {
 146                 fprintf(stderr, "kqueue failed: %d\n", errno);
 147                 exit(1);
 148         }
 149
 150         EV_SET(&change_evt0[0], serv_s, EVFILT_READ, EV_ADD, 0, 0, NULL);
 151         nchange = 1;
 152
 153         for (;;) {
 154                 const struct kevent *change_evt = NULL;
 155                 struct kevent evt[EVENT_MAX];
 156                 int i, nevt;
 157
 158                 if (nchange > 0)
 159                         change_evt = change_evt0;
 160
 161                 nevt = kevent(kq, change_evt, nchange, evt, EVENT_MAX, NULL);
 162                 if (nevt < 0) {
 163                         fprintf(stderr, "kevent failed: %d\n", errno);
 164                         exit(1);
 165                 }
 166                 nchange = 0;
 167
 168                 for (i = 0; i < nevt; ++i) {
 169                         if (evt[i].ident == (u_int)serv_s) {
 170                                 while (nchange < EVENT_MAX) {
 171                                         int s;
 172
 173                                         s = accept(serv_s, NULL, NULL);
 174                                         if (s < 0)
 175                                                 break;
 176                                         EV_SET(&change_evt0[nchange], s,
 177                                             EVFILT_READ, EV_ADD, 0, 0, NULL);
 178                                         ++nchange;
 179                                 }
 180                         } else {
 181                                 close(evt[i].ident);
 182                         }
 183                 }
 184         }
 185         /* NEVER REACHED */
 186 }