tcpp_server.c revision 208873 
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: head/tools/tools/netrate/tcpp/tcpp_server.c 208873 2010-06-06 15:27:08Z rwatson $
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
64struct 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
73static pid_t			*pid_list;
74static int			 kq;
75
76static struct connection *
77tcpp_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
110static void
111tcpp_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
126static u_char buffer[256*1024];	/* Buffer in which to sink data. */
127static void
128tcpp_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
193static void
194tcpp_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
277void
278tcpp_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