tools/tools/netrate/tcpp/tcpp_server.c

   1 /*-
   2  * Copyright (c) 2008-2009 Robert N. M. Watson
   3  * All rights reserved.
   4  *
   5  * Redistribution and use in source and binary forms, with or without
   6  * modification, are permitted provided that the following conditions
   7  * are met:
   8  * 1. Redistributions of source code must retain the above copyright
   9  *    notice, this list of conditions and the following disclaimer.
  10  * 2. Redistributions in binary form must reproduce the above copyright
  11  *    notice, this list of conditions and the following disclaimer in the
  12  *    documentation and/or other materials provided with the distribution.
  13  *
  14  * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
  15  * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  16  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  17  * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
  18  * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
  19  * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
  20  * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
  21  * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
  22  * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
  23  * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
  24  * SUCH DAMAGE.
  25  *
  26  * $FreeBSD$
  27  */
  28
  29 #include <sys/types.h>
  30 #include <sys/endian.h>
  31 #include <sys/event.h>
  32 #include <sys/resource.h>
  33 #include <sys/sched.h>
  34 #include <sys/socket.h>
  35 #include <sys/sysctl.h>
  36 #include <sys/time.h>
  37 #include <sys/wait.h>
  38
  39 #include <netinet/in.h>
  40 #include <netinet/tcp.h>
  41
  42 #include <err.h>
  43 #include <fcntl.h>
  44 #include <inttypes.h>
  45 #include <signal.h>
  46 #include <stdio.h>
  47 #include <stdlib.h>
  48 #include <string.h>
  49 #include <unistd.h>
  50
  51 #include "tcpp.h"
  52
  53 /*
  54  * Server side -- create a pool of processes, each listening on its own TCP
  55  * port number for new connections.  The first 8 bytes of each connection
  56  * will be a network byte order length, then there will be that number of
  57  * bytes of data.  We use non-blocking sockets with kqueue to to avoid the
  58  * overhead of threading or more than one process per processor, which makes
  59  * things a bit awkward when dealing with data we care about.  As such, we
  60  * read into a small character buffer which we then convert to a length once
  61  * we have all the data.
  62  */
  63 #define CONNECTION_MAGIC        0x6392af27
  64 struct connection {
  65         uint32_t        conn_magic;             /* Just magic. */
  66         int             conn_fd;
  67         struct tcpp_header      conn_header;    /* Header buffer. */
  68         u_int           conn_header_len;        /* Bytes so far. */
  69         u_int64_t       conn_data_len;          /* How much to sink. */
  70         u_int64_t       conn_data_received;     /* How much so far. */
  71 };
  72
  73 static pid_t                    *pid_list;
  74 static int                       kq;
  75
  76 static struct connection *
  77 tcpp_server_newconn(int listen_fd)
  78 {
  79         struct connection *conn;
  80         struct kevent kev;
  81         int fd;
  82
  83         fd = accept(listen_fd, NULL, NULL);
  84         if (fd < 0) {
  85                 warn("accept");
  86                 return (NULL);
  87         }
  88
  89         if (fcntl(fd, F_SETFL, O_NONBLOCK) < 0)
  90                 err(-1, "fcntl");
  91
  92         conn = malloc(sizeof(*conn));
  93         if (conn == NULL)
  94                 return (NULL);
  95         bzero(conn, sizeof(*conn));
  96         conn->conn_magic = CONNECTION_MAGIC;
  97         conn->conn_fd = fd;
  98
  99         /*
 100          * Register to read on the socket, and set our conn pointer as the
 101          * udata so we can find it quickly in the future.
 102          */
 103         EV_SET(&kev, fd, EVFILT_READ, EV_ADD, 0, 0, conn);
 104         if (kevent(kq, &kev, 1, NULL, 0, NULL) < 0)
 105                 err(-1, "kevent");
 106
 107         return (conn);
 108 }
 109
 110 static void
 111 tcpp_server_closeconn(struct connection *conn)
 112 {
 113
 114         /*
 115          * Kqueue cleans up after itself once we close the socket, and since
 116          * we are processing only one kevent at a time, we don't need to
 117          * worry about watching out for future kevents referring to it.
 118          *
 119          * ... right?
 120          */
 121         close(conn->conn_fd);
 122         bzero(conn, sizeof(*conn));
 123         free(conn);
 124 }
 125
 126 static u_char buffer[256*1024]; /* Buffer in which to sink data. */
 127 static void
 128 tcpp_server_handleconn(struct kevent *kev)
 129 {
 130         struct connection *conn;
 131         ssize_t len;
 132
 133         conn = kev->udata;
 134         if (conn->conn_magic != CONNECTION_MAGIC)
 135                 errx(-1, "tcpp_server_handleconn: magic");
 136
 137         if (conn->conn_header_len < sizeof(conn->conn_header)) {
 138                 len = read(conn->conn_fd,
 139                     ((u_char *)&conn->conn_header) + conn->conn_header_len,
 140                     sizeof(conn->conn_header) - conn->conn_header_len);
 141                 if (len < 0) {
 142                         warn("tcpp_server_handleconn: header read");
 143                         tcpp_server_closeconn(conn);
 144                         return;
 145                 }
 146                 if (len == 0) {
 147                         warnx("tcpp_server_handleconn: header premature eof");
 148                         tcpp_server_closeconn(conn);
 149                         return;
 150                 }
 151                 conn->conn_header_len += len;
 152                 if (conn->conn_header_len == sizeof(conn->conn_header)) {
 153                         tcpp_header_decode(&conn->conn_header);
 154                         if (conn->conn_header.th_magic != TCPP_MAGIC) {
 155                                 warnx("tcpp_server_handleconn: bad magic");
 156                                 tcpp_server_closeconn(conn);
 157                                 return;
 158                         }
 159                 }
 160         } else {
 161                 /*
 162                  * Drain up to a buffer from the connection, so that we pay
 163                  * attention to other connections too.
 164                  */
 165                 len = read(conn->conn_fd, buffer, sizeof(buffer));
 166                 if (len < 0) {
 167                         warn("tcpp_server_handleconn: data bad read");
 168                         tcpp_server_closeconn(conn);
 169                         return;
 170                 }
 171                 if (len == 0 && conn->conn_data_received <
 172                     conn->conn_header.th_len) {
 173                         warnx("tcpp_server_handleconn: data premature eof");
 174                         tcpp_server_closeconn(conn);
 175                         return;
 176                 }
 177                 conn->conn_data_received += len;
 178                 if (conn->conn_data_received > conn->conn_header.th_len) {
 179                         warnx("tcpp_server_handleconn: too much data");
 180                         tcpp_server_closeconn(conn);
 181                         return;
 182                 }
 183                 if (conn->conn_data_received == conn->conn_header.th_len) {
 184                         /*
 185                          * All is well.
 186                          */
 187                         tcpp_server_closeconn(conn);
 188                         return;
 189                 }
 190         }
 191 }
 192
 193 static void
 194 tcpp_server_worker(int workernum)
 195 {
 196         int i, listen_sock, numevents;
 197         struct kevent kev, *kev_array;
 198         int kev_bytes;
 199 #if defined(CPU_SETSIZE) && 0
 200         cpu_set_t mask;
 201         int ncpus;
 202         ssize_t len;
 203
 204         if (Pflag) {
 205                 len = sizeof(ncpus);
 206                 if (sysctlbyname(SYSCTLNAME_CPUS, &ncpus, &len, NULL, 0) < 0)
 207                         err(-1, "sysctlbyname: %s", SYSCTLNAME_CPUS);
 208                 if (len != sizeof(ncpus))
 209                         errx(-1, "sysctlbyname: %s: len %jd", SYSCTLNAME_CPUS,
 210                             (intmax_t)len);
 211
 212                 CPU_ZERO(&mask);
 213                 CPU_SET(workernum % ncpus, &mask);
 214                 if (sched_setaffinity(0, CPU_SETSIZE, &mask) < 0)
 215                         err(-1, "sched_setaffinity");
 216         }
 217 #endif
 218         setproctitle("tcpp_server %d", workernum);
 219
 220         /* Allow an extra kevent for the listen socket. */
 221         kev_bytes = sizeof(*kev_array) * (mflag + 1);
 222         kev_array = malloc(kev_bytes);
 223         if (kev_array == NULL)
 224                 err(-1, "malloc");
 225         bzero(kev_array, kev_bytes);
 226
 227         /* XXXRW: Want to set and pin the CPU here. */
 228
 229         /*
 230          * Add the worker number to the local port.
 231          */
 232         localipbase.sin_port = htons(rflag + workernum);
 233
 234         listen_sock = socket(PF_INET, SOCK_STREAM, 0);
 235         if (listen_sock < 0)
 236                 err(-1, "socket");
 237         i = 1;
 238         if (setsockopt(listen_sock, SOL_SOCKET, SO_NOSIGPIPE, &i, sizeof(i))
 239             < 0)
 240                 err(-1, "setsockopt");
 241         i = 1;
 242         if (setsockopt(listen_sock, SOL_SOCKET, SO_REUSEPORT, &i, sizeof(i))
 243             < 0)
 244                 err(-1, "setsockopt");
 245         i = 1;
 246         if (setsockopt(listen_sock, IPPROTO_TCP, TCP_NODELAY, &i, sizeof(i))
 247             < 0)
 248                 err(-1, "setsockopt");
 249         if (bind(listen_sock, (struct sockaddr *)&localipbase,
 250             sizeof(localipbase)) < 0)
 251                 err(-1, "bind");
 252         if (listen(listen_sock, 16384))
 253                 err(-1, "listen");
 254         if (fcntl(listen_sock, F_SETFL, O_NONBLOCK) < 0)
 255                 err(-1, "fcntl");
 256
 257         kq = kqueue();
 258         if (kq < 0)
 259                 err(-1, "kqueue");
 260
 261         EV_SET(&kev, listen_sock, EVFILT_READ, EV_ADD, 0, 0, NULL);
 262         if (kevent(kq, &kev, 1, NULL, 0, NULL) < 0)
 263                 err(-1, "kevent");
 264
 265         while ((numevents = kevent(kq, NULL, 0, kev_array, mflag + 1, NULL))
 266             > 0) {
 267                 for (i = 0; i < numevents; i++) {
 268                         if (kev_array[i].ident == (u_int)listen_sock)
 269                                 (void)tcpp_server_newconn(listen_sock);
 270                         else
 271                                 tcpp_server_handleconn(&kev_array[i]);
 272                 }
 273         }
 274         printf("Worker %d done\n", workernum);
 275 }
 276
 277 void
 278 tcpp_server(void)
 279 {
 280 #if 0
 281         long cp_time_last[CPUSTATES], cp_time_now[CPUSTATES], ticks;
 282         size_t size;
 283 #endif
 284         pid_t pid;
 285         int i;
 286
 287         pid_list = malloc(sizeof(*pid_list) * pflag);
 288         if (pid_list == NULL)
 289                 err(-1, "malloc pid_list");
 290         bzero(pid_list, sizeof(*pid_list) * pflag);
 291
 292         /*
 293          * Start workers.
 294          */
 295         for (i = 0; i < pflag; i++) {
 296                 pid = fork();
 297                 if (pid < 0) {
 298                         warn("fork");
 299                         for (i = 0; i < pflag; i++) {
 300                                 if (pid_list[i] != 0)
 301                                         (void)kill(pid_list[i], SIGKILL);
 302                         }
 303                         exit(-1);
 304                 }
 305                 if (pid == 0) {
 306                         tcpp_server_worker(i);
 307                         exit(0);
 308                 }
 309                 pid_list[i] = pid;
 310         }
 311
 312 #if 0
 313                 size = sizeof(cp_time_last);
 314                 if (sysctlbyname(SYSCTLNAME_CPTIME, &cp_time_last, &size,
 315                     NULL, 0) < 0)
 316                         err(-1, "sysctlbyname: %s", SYSCTLNAME_CPTIME);
 317                 while (1) {
 318                         sleep(10);
 319                         size = sizeof(cp_time_last);
 320                         if (sysctlbyname(SYSCTLNAME_CPTIME, &cp_time_now,
 321                             &size, NULL, 0) < 0)
 322                                 err(-1, "sysctlbyname: %s",
 323                                     SYSCTLNAME_CPTIME);
 324                         ticks = 0;
 325                         for (i = 0; i < CPUSTATES; i++) {
 326                                 cp_time_last[i] = cp_time_now[i] -
 327                                     cp_time_last[i];
 328                                 ticks += cp_time_last[i];
 329                         }
 330                         printf("user%% %lu nice%% %lu sys%% %lu intr%% %lu "
 331                             "idle%% %lu\n",
 332                             (100 * cp_time_last[CP_USER]) / ticks,
 333                             (100 * cp_time_last[CP_NICE]) / ticks,
 334                             (100 * cp_time_last[CP_SYS]) / ticks,
 335                             (100 * cp_time_last[CP_INTR]) / ticks,
 336                             (100 * cp_time_last[CP_IDLE]) / ticks);
 337                         bcopy(cp_time_now, cp_time_last, sizeof(cp_time_last));
 338                 }
 339 #endif
 340
 341         /*
 342          * GC workers.
 343          */
 344         for (i = 0; i < pflag; i++) {
 345                 if (pid_list[i] != 0) {
 346                         while (waitpid(pid_list[i], NULL, 0) != pid_list[i]);
 347                 }
 348         }
 349 }