1 /*
2  * Copyright 2018 Google Inc.
3  * Author: Eric Dumazet (edumazet@google.com)
4  *
5  * Reference program demonstrating tcp mmap() usage,
6  * and SO_RCVLOWAT hints for receiver.
7  *
8  * Note : NIC with header split is needed to use mmap() on TCP :
9  * Each incoming frame must be a multiple of PAGE_SIZE bytes of TCP payload.
10  *
11  * How to use on loopback interface :
12  *
13  *  ifconfig lo mtu 61512  # 15*4096 + 40 (ipv6 header) + 32 (TCP with TS option header)
14  *  tcp_mmap -s -z &
15  *  tcp_mmap -H ::1 -z
16  *
17  *  Or leave default lo mtu, but use -M option to set TCP_MAXSEG option to (4096 + 12)
18  *      (4096 : page size on x86, 12: TCP TS option length)
19  *  tcp_mmap -s -z -M $((4096+12)) &
20  *  tcp_mmap -H ::1 -z -M $((4096+12))
21  *
22  * Note: -z option on sender uses MSG_ZEROCOPY, which forces a copy when packets go through loopback interface.
23  *       We might use sendfile() instead, but really this test program is about mmap(), for receivers ;)
24  *
25  * $ ./tcp_mmap -s &                                 # Without mmap()
26  * $ for i in {1..4}; do ./tcp_mmap -H ::1 -z ; done
27  * received 32768 MB (0 % mmap'ed) in 14.1157 s, 19.4732 Gbit
28  *   cpu usage user:0.057 sys:7.815, 240.234 usec per MB, 65531 c-switches
29  * received 32768 MB (0 % mmap'ed) in 14.6833 s, 18.7204 Gbit
30  *  cpu usage user:0.043 sys:8.103, 248.596 usec per MB, 65524 c-switches
31  * received 32768 MB (0 % mmap'ed) in 11.143 s, 24.6682 Gbit
32  *   cpu usage user:0.044 sys:6.576, 202.026 usec per MB, 65519 c-switches
33  * received 32768 MB (0 % mmap'ed) in 14.9056 s, 18.4413 Gbit
34  *   cpu usage user:0.036 sys:8.193, 251.129 usec per MB, 65530 c-switches
35  * $ kill %1   # kill tcp_mmap server
36  *
37  * $ ./tcp_mmap -s -z &                              # With mmap()
38  * $ for i in {1..4}; do ./tcp_mmap -H ::1 -z ; done
39  * received 32768 MB (99.9939 % mmap'ed) in 6.73792 s, 40.7956 Gbit
40  *   cpu usage user:0.045 sys:2.827, 87.6465 usec per MB, 65532 c-switches
41  * received 32768 MB (99.9939 % mmap'ed) in 7.26732 s, 37.8238 Gbit
42  *   cpu usage user:0.037 sys:3.087, 95.3369 usec per MB, 65532 c-switches
43  * received 32768 MB (99.9939 % mmap'ed) in 7.61661 s, 36.0893 Gbit
44  *   cpu usage user:0.046 sys:3.559, 110.016 usec per MB, 65529 c-switches
45  * received 32768 MB (99.9939 % mmap'ed) in 7.43764 s, 36.9577 Gbit
46  *   cpu usage user:0.035 sys:3.467, 106.873 usec per MB, 65530 c-switches
47  *
48  * License (GPLv2):
49  *
50  * This program is free software; you can redistribute it and/or modify it
51  * under the terms and conditions of the GNU General Public License,
52  * version 2, as published by the Free Software Foundation.
53  *
54  * This program is distributed in the hope it will be useful, but WITHOUT
55  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
56  * FITNESS FOR A PARTICULAR PURPOSE. * See the GNU General Public License for
57  * more details.
58  *
59  * You should have received a copy of the GNU General Public License along with
60  * this program; if not, write to the Free Software Foundation, Inc.,
61  * 51 Franklin St - Fifth Floor, Boston, MA 02110-1301 USA.
62  */
63 #define _GNU_SOURCE
64 #include <pthread.h>
65 #include <sys/types.h>
66 #include <fcntl.h>
67 #include <error.h>
68 #include <sys/socket.h>
69 #include <sys/mman.h>
70 #include <sys/resource.h>
71 #include <unistd.h>
72 #include <string.h>
73 #include <stdlib.h>
74 #include <stdio.h>
75 #include <errno.h>
76 #include <time.h>
77 #include <sys/time.h>
78 #include <netinet/in.h>
79 #include <arpa/inet.h>
80 #include <poll.h>
81 #include <linux/tcp.h>
82 #include <assert.h>
83 
84 #ifndef MSG_ZEROCOPY
85 #define MSG_ZEROCOPY    0x4000000
86 #endif
87 
88 #define FILE_SZ (1UL << 35)
89 static int cfg_family = AF_INET6;
90 static socklen_t cfg_alen = sizeof(struct sockaddr_in6);
91 static int cfg_port = 8787;
92 
93 static int rcvbuf; /* Default: autotuning.  Can be set with -r <integer> option */
94 static int sndbuf; /* Default: autotuning.  Can be set with -w <integer> option */
95 static int zflg; /* zero copy option. (MSG_ZEROCOPY for sender, mmap() for receiver */
96 static int xflg; /* hash received data (simple xor) (-h option) */
97 static int keepflag; /* -k option: receiver shall keep all received file in memory (no munmap() calls) */
98 
99 static int chunk_size  = 512*1024;
100 
101 unsigned long htotal;
102 
103 static inline void prefetch(const void *x)
104 {
105 #if defined(__x86_64__)
106 	asm volatile("prefetcht0 %P0" : : "m" (*(const char *)x));
107 #endif
108 }
109 
110 void hash_zone(void *zone, unsigned int length)
111 {
112 	unsigned long temp = htotal;
113 
114 	while (length >= 8*sizeof(long)) {
115 		prefetch(zone + 384);
116 		temp ^= *(unsigned long *)zone;
117 		temp ^= *(unsigned long *)(zone + sizeof(long));
118 		temp ^= *(unsigned long *)(zone + 2*sizeof(long));
119 		temp ^= *(unsigned long *)(zone + 3*sizeof(long));
120 		temp ^= *(unsigned long *)(zone + 4*sizeof(long));
121 		temp ^= *(unsigned long *)(zone + 5*sizeof(long));
122 		temp ^= *(unsigned long *)(zone + 6*sizeof(long));
123 		temp ^= *(unsigned long *)(zone + 7*sizeof(long));
124 		zone += 8*sizeof(long);
125 		length -= 8*sizeof(long);
126 	}
127 	while (length >= 1) {
128 		temp ^= *(unsigned char *)zone;
129 		zone += 1;
130 		length--;
131 	}
132 	htotal = temp;
133 }
134 
135 void *child_thread(void *arg)
136 {
137 	unsigned long total_mmap = 0, total = 0;
138 	struct tcp_zerocopy_receive zc;
139 	unsigned long delta_usec;
140 	int flags = MAP_SHARED;
141 	struct timeval t0, t1;
142 	char *buffer = NULL;
143 	void *addr = NULL;
144 	double throughput;
145 	struct rusage ru;
146 	int lu, fd;
147 
148 	fd = (int)(unsigned long)arg;
149 
150 	gettimeofday(&t0, NULL);
151 
152 	fcntl(fd, F_SETFL, O_NDELAY);
153 	buffer = malloc(chunk_size);
154 	if (!buffer) {
155 		perror("malloc");
156 		goto error;
157 	}
158 	if (zflg) {
159 		addr = mmap(NULL, chunk_size, PROT_READ, flags, fd, 0);
160 		if (addr == (void *)-1)
161 			zflg = 0;
162 	}
163 	while (1) {
164 		struct pollfd pfd = { .fd = fd, .events = POLLIN, };
165 		int sub;
166 
167 		poll(&pfd, 1, 10000);
168 		if (zflg) {
169 			socklen_t zc_len = sizeof(zc);
170 			int res;
171 
172 			zc.address = (__u64)addr;
173 			zc.length = chunk_size;
174 			zc.recv_skip_hint = 0;
175 			res = getsockopt(fd, IPPROTO_TCP, TCP_ZEROCOPY_RECEIVE,
176 					 &zc, &zc_len);
177 			if (res == -1)
178 				break;
179 
180 			if (zc.length) {
181 				assert(zc.length <= chunk_size);
182 				total_mmap += zc.length;
183 				if (xflg)
184 					hash_zone(addr, zc.length);
185 				total += zc.length;
186 			}
187 			if (zc.recv_skip_hint) {
188 				assert(zc.recv_skip_hint <= chunk_size);
189 				lu = read(fd, buffer, zc.recv_skip_hint);
190 				if (lu > 0) {
191 					if (xflg)
192 						hash_zone(buffer, lu);
193 					total += lu;
194 				}
195 			}
196 			continue;
197 		}
198 		sub = 0;
199 		while (sub < chunk_size) {
200 			lu = read(fd, buffer + sub, chunk_size - sub);
201 			if (lu == 0)
202 				goto end;
203 			if (lu < 0)
204 				break;
205 			if (xflg)
206 				hash_zone(buffer + sub, lu);
207 			total += lu;
208 			sub += lu;
209 		}
210 	}
211 end:
212 	gettimeofday(&t1, NULL);
213 	delta_usec = (t1.tv_sec - t0.tv_sec) * 1000000 + t1.tv_usec - t0.tv_usec;
214 
215 	throughput = 0;
216 	if (delta_usec)
217 		throughput = total * 8.0 / (double)delta_usec / 1000.0;
218 	getrusage(RUSAGE_THREAD, &ru);
219 	if (total > 1024*1024) {
220 		unsigned long total_usec;
221 		unsigned long mb = total >> 20;
222 		total_usec = 1000000*ru.ru_utime.tv_sec + ru.ru_utime.tv_usec +
223 			     1000000*ru.ru_stime.tv_sec + ru.ru_stime.tv_usec;
224 		printf("received %lg MB (%lg %% mmap'ed) in %lg s, %lg Gbit\n"
225 		       "  cpu usage user:%lg sys:%lg, %lg usec per MB, %lu c-switches\n",
226 				total / (1024.0 * 1024.0),
227 				100.0*total_mmap/total,
228 				(double)delta_usec / 1000000.0,
229 				throughput,
230 				(double)ru.ru_utime.tv_sec + (double)ru.ru_utime.tv_usec / 1000000.0,
231 				(double)ru.ru_stime.tv_sec + (double)ru.ru_stime.tv_usec / 1000000.0,
232 				(double)total_usec/mb,
233 				ru.ru_nvcsw);
234 	}
235 error:
236 	free(buffer);
237 	close(fd);
238 	if (zflg)
239 		munmap(addr, chunk_size);
240 	pthread_exit(0);
241 }
242 
243 static void apply_rcvsnd_buf(int fd)
244 {
245 	if (rcvbuf && setsockopt(fd, SOL_SOCKET,
246 				 SO_RCVBUF, &rcvbuf, sizeof(rcvbuf)) == -1) {
247 		perror("setsockopt SO_RCVBUF");
248 	}
249 
250 	if (sndbuf && setsockopt(fd, SOL_SOCKET,
251 				 SO_SNDBUF, &sndbuf, sizeof(sndbuf)) == -1) {
252 		perror("setsockopt SO_SNDBUF");
253 	}
254 }
255 
256 
257 static void setup_sockaddr(int domain, const char *str_addr,
258 			   struct sockaddr_storage *sockaddr)
259 {
260 	struct sockaddr_in6 *addr6 = (void *) sockaddr;
261 	struct sockaddr_in *addr4 = (void *) sockaddr;
262 
263 	switch (domain) {
264 	case PF_INET:
265 		memset(addr4, 0, sizeof(*addr4));
266 		addr4->sin_family = AF_INET;
267 		addr4->sin_port = htons(cfg_port);
268 		if (str_addr &&
269 		    inet_pton(AF_INET, str_addr, &(addr4->sin_addr)) != 1)
270 			error(1, 0, "ipv4 parse error: %s", str_addr);
271 		break;
272 	case PF_INET6:
273 		memset(addr6, 0, sizeof(*addr6));
274 		addr6->sin6_family = AF_INET6;
275 		addr6->sin6_port = htons(cfg_port);
276 		if (str_addr &&
277 		    inet_pton(AF_INET6, str_addr, &(addr6->sin6_addr)) != 1)
278 			error(1, 0, "ipv6 parse error: %s", str_addr);
279 		break;
280 	default:
281 		error(1, 0, "illegal domain");
282 	}
283 }
284 
285 static void do_accept(int fdlisten)
286 {
287 	if (setsockopt(fdlisten, SOL_SOCKET, SO_RCVLOWAT,
288 		       &chunk_size, sizeof(chunk_size)) == -1) {
289 		perror("setsockopt SO_RCVLOWAT");
290 	}
291 
292 	apply_rcvsnd_buf(fdlisten);
293 
294 	while (1) {
295 		struct sockaddr_in addr;
296 		socklen_t addrlen = sizeof(addr);
297 		pthread_t th;
298 		int fd, res;
299 
300 		fd = accept(fdlisten, (struct sockaddr *)&addr, &addrlen);
301 		if (fd == -1) {
302 			perror("accept");
303 			continue;
304 		}
305 		res = pthread_create(&th, NULL, child_thread,
306 				     (void *)(unsigned long)fd);
307 		if (res) {
308 			errno = res;
309 			perror("pthread_create");
310 			close(fd);
311 		}
312 	}
313 }
314 
315 int main(int argc, char *argv[])
316 {
317 	struct sockaddr_storage listenaddr, addr;
318 	unsigned int max_pacing_rate = 0;
319 	unsigned long total = 0;
320 	char *host = NULL;
321 	int fd, c, on = 1;
322 	char *buffer;
323 	int sflg = 0;
324 	int mss = 0;
325 
326 	while ((c = getopt(argc, argv, "46p:svr:w:H:zxkP:M:")) != -1) {
327 		switch (c) {
328 		case '4':
329 			cfg_family = PF_INET;
330 			cfg_alen = sizeof(struct sockaddr_in);
331 			break;
332 		case '6':
333 			cfg_family = PF_INET6;
334 			cfg_alen = sizeof(struct sockaddr_in6);
335 			break;
336 		case 'p':
337 			cfg_port = atoi(optarg);
338 			break;
339 		case 'H':
340 			host = optarg;
341 			break;
342 		case 's': /* server : listen for incoming connections */
343 			sflg++;
344 			break;
345 		case 'r':
346 			rcvbuf = atoi(optarg);
347 			break;
348 		case 'w':
349 			sndbuf = atoi(optarg);
350 			break;
351 		case 'z':
352 			zflg = 1;
353 			break;
354 		case 'M':
355 			mss = atoi(optarg);
356 			break;
357 		case 'x':
358 			xflg = 1;
359 			break;
360 		case 'k':
361 			keepflag = 1;
362 			break;
363 		case 'P':
364 			max_pacing_rate = atoi(optarg) ;
365 			break;
366 		default:
367 			exit(1);
368 		}
369 	}
370 	if (sflg) {
371 		int fdlisten = socket(cfg_family, SOCK_STREAM, 0);
372 
373 		if (fdlisten == -1) {
374 			perror("socket");
375 			exit(1);
376 		}
377 		apply_rcvsnd_buf(fdlisten);
378 		setsockopt(fdlisten, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on));
379 
380 		setup_sockaddr(cfg_family, host, &listenaddr);
381 
382 		if (mss &&
383 		    setsockopt(fdlisten, IPPROTO_TCP, TCP_MAXSEG,
384 			       &mss, sizeof(mss)) == -1) {
385 			perror("setsockopt TCP_MAXSEG");
386 			exit(1);
387 		}
388 		if (bind(fdlisten, (const struct sockaddr *)&listenaddr, cfg_alen) == -1) {
389 			perror("bind");
390 			exit(1);
391 		}
392 		if (listen(fdlisten, 128) == -1) {
393 			perror("listen");
394 			exit(1);
395 		}
396 		do_accept(fdlisten);
397 	}
398 	buffer = mmap(NULL, chunk_size, PROT_READ | PROT_WRITE,
399 			      MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
400 	if (buffer == (char *)-1) {
401 		perror("mmap");
402 		exit(1);
403 	}
404 
405 	fd = socket(cfg_family, SOCK_STREAM, 0);
406 	if (fd == -1) {
407 		perror("socket");
408 		exit(1);
409 	}
410 	apply_rcvsnd_buf(fd);
411 
412 	setup_sockaddr(cfg_family, host, &addr);
413 
414 	if (mss &&
415 	    setsockopt(fd, IPPROTO_TCP, TCP_MAXSEG, &mss, sizeof(mss)) == -1) {
416 		perror("setsockopt TCP_MAXSEG");
417 		exit(1);
418 	}
419 	if (connect(fd, (const struct sockaddr *)&addr, cfg_alen) == -1) {
420 		perror("connect");
421 		exit(1);
422 	}
423 	if (max_pacing_rate &&
424 	    setsockopt(fd, SOL_SOCKET, SO_MAX_PACING_RATE,
425 		       &max_pacing_rate, sizeof(max_pacing_rate)) == -1)
426 		perror("setsockopt SO_MAX_PACING_RATE");
427 
428 	if (zflg && setsockopt(fd, SOL_SOCKET, SO_ZEROCOPY,
429 			       &on, sizeof(on)) == -1) {
430 		perror("setsockopt SO_ZEROCOPY, (-z option disabled)");
431 		zflg = 0;
432 	}
433 	while (total < FILE_SZ) {
434 		long wr = FILE_SZ - total;
435 
436 		if (wr > chunk_size)
437 			wr = chunk_size;
438 		/* Note : we just want to fill the pipe with 0 bytes */
439 		wr = send(fd, buffer, wr, zflg ? MSG_ZEROCOPY : 0);
440 		if (wr <= 0)
441 			break;
442 		total += wr;
443 	}
444 	close(fd);
445 	munmap(buffer, chunk_size);
446 	return 0;
447 }
448