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