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