1// SPDX-License-Identifier: GPL-2.0-only
2/*
3 * vsock_perf - benchmark utility for vsock.
4 *
5 * Copyright (C) 2022 SberDevices.
6 *
7 * Author: Arseniy Krasnov <AVKrasnov@sberdevices.ru>
8 */
9#include <getopt.h>
10#include <stdio.h>
11#include <stdlib.h>
12#include <stdbool.h>
13#include <string.h>
14#include <errno.h>
15#include <unistd.h>
16#include <time.h>
17#include <stdint.h>
18#include <poll.h>
19#include <sys/socket.h>
20#include <linux/vm_sockets.h>
21
22#define DEFAULT_BUF_SIZE_BYTES	(128 * 1024)
23#define DEFAULT_TO_SEND_BYTES	(64 * 1024)
24#define DEFAULT_VSOCK_BUF_BYTES (256 * 1024)
25#define DEFAULT_RCVLOWAT_BYTES	1
26#define DEFAULT_PORT		1234
27
28#define BYTES_PER_GB		(1024 * 1024 * 1024ULL)
29#define NSEC_PER_SEC		(1000000000ULL)
30
31static unsigned int port = DEFAULT_PORT;
32static unsigned long buf_size_bytes = DEFAULT_BUF_SIZE_BYTES;
33static unsigned long vsock_buf_bytes = DEFAULT_VSOCK_BUF_BYTES;
34
35static void error(const char *s)
36{
37	perror(s);
38	exit(EXIT_FAILURE);
39}
40
41static time_t current_nsec(void)
42{
43	struct timespec ts;
44
45	if (clock_gettime(CLOCK_REALTIME, &ts))
46		error("clock_gettime");
47
48	return (ts.tv_sec * NSEC_PER_SEC) + ts.tv_nsec;
49}
50
51/* From lib/cmdline.c. */
52static unsigned long memparse(const char *ptr)
53{
54	char *endptr;
55
56	unsigned long long ret = strtoull(ptr, &endptr, 0);
57
58	switch (*endptr) {
59	case 'E':
60	case 'e':
61		ret <<= 10;
62	case 'P':
63	case 'p':
64		ret <<= 10;
65	case 'T':
66	case 't':
67		ret <<= 10;
68	case 'G':
69	case 'g':
70		ret <<= 10;
71	case 'M':
72	case 'm':
73		ret <<= 10;
74	case 'K':
75	case 'k':
76		ret <<= 10;
77		endptr++;
78	default:
79		break;
80	}
81
82	return ret;
83}
84
85static void vsock_increase_buf_size(int fd)
86{
87	if (setsockopt(fd, AF_VSOCK, SO_VM_SOCKETS_BUFFER_MAX_SIZE,
88		       &vsock_buf_bytes, sizeof(vsock_buf_bytes)))
89		error("setsockopt(SO_VM_SOCKETS_BUFFER_MAX_SIZE)");
90
91	if (setsockopt(fd, AF_VSOCK, SO_VM_SOCKETS_BUFFER_SIZE,
92		       &vsock_buf_bytes, sizeof(vsock_buf_bytes)))
93		error("setsockopt(SO_VM_SOCKETS_BUFFER_SIZE)");
94}
95
96static int vsock_connect(unsigned int cid, unsigned int port)
97{
98	union {
99		struct sockaddr sa;
100		struct sockaddr_vm svm;
101	} addr = {
102		.svm = {
103			.svm_family = AF_VSOCK,
104			.svm_port = port,
105			.svm_cid = cid,
106		},
107	};
108	int fd;
109
110	fd = socket(AF_VSOCK, SOCK_STREAM, 0);
111
112	if (fd < 0) {
113		perror("socket");
114		return -1;
115	}
116
117	if (connect(fd, &addr.sa, sizeof(addr.svm)) < 0) {
118		perror("connect");
119		close(fd);
120		return -1;
121	}
122
123	return fd;
124}
125
126static float get_gbps(unsigned long bits, time_t ns_delta)
127{
128	return ((float)bits / 1000000000ULL) /
129	       ((float)ns_delta / NSEC_PER_SEC);
130}
131
132static void run_receiver(unsigned long rcvlowat_bytes)
133{
134	unsigned int read_cnt;
135	time_t rx_begin_ns;
136	time_t in_read_ns;
137	size_t total_recv;
138	int client_fd;
139	char *data;
140	int fd;
141	union {
142		struct sockaddr sa;
143		struct sockaddr_vm svm;
144	} addr = {
145		.svm = {
146			.svm_family = AF_VSOCK,
147			.svm_port = port,
148			.svm_cid = VMADDR_CID_ANY,
149		},
150	};
151	union {
152		struct sockaddr sa;
153		struct sockaddr_vm svm;
154	} clientaddr;
155
156	socklen_t clientaddr_len = sizeof(clientaddr.svm);
157
158	printf("Run as receiver\n");
159	printf("Listen port %u\n", port);
160	printf("RX buffer %lu bytes\n", buf_size_bytes);
161	printf("vsock buffer %lu bytes\n", vsock_buf_bytes);
162	printf("SO_RCVLOWAT %lu bytes\n", rcvlowat_bytes);
163
164	fd = socket(AF_VSOCK, SOCK_STREAM, 0);
165
166	if (fd < 0)
167		error("socket");
168
169	if (bind(fd, &addr.sa, sizeof(addr.svm)) < 0)
170		error("bind");
171
172	if (listen(fd, 1) < 0)
173		error("listen");
174
175	client_fd = accept(fd, &clientaddr.sa, &clientaddr_len);
176
177	if (client_fd < 0)
178		error("accept");
179
180	vsock_increase_buf_size(client_fd);
181
182	if (setsockopt(client_fd, SOL_SOCKET, SO_RCVLOWAT,
183		       &rcvlowat_bytes,
184		       sizeof(rcvlowat_bytes)))
185		error("setsockopt(SO_RCVLOWAT)");
186
187	data = malloc(buf_size_bytes);
188
189	if (!data) {
190		fprintf(stderr, "'malloc()' failed\n");
191		exit(EXIT_FAILURE);
192	}
193
194	read_cnt = 0;
195	in_read_ns = 0;
196	total_recv = 0;
197	rx_begin_ns = current_nsec();
198
199	while (1) {
200		struct pollfd fds = { 0 };
201
202		fds.fd = client_fd;
203		fds.events = POLLIN | POLLERR |
204			     POLLHUP | POLLRDHUP;
205
206		if (poll(&fds, 1, -1) < 0)
207			error("poll");
208
209		if (fds.revents & POLLERR) {
210			fprintf(stderr, "'poll()' error\n");
211			exit(EXIT_FAILURE);
212		}
213
214		if (fds.revents & POLLIN) {
215			ssize_t bytes_read;
216			time_t t;
217
218			t = current_nsec();
219			bytes_read = read(fds.fd, data, buf_size_bytes);
220			in_read_ns += (current_nsec() - t);
221			read_cnt++;
222
223			if (!bytes_read)
224				break;
225
226			if (bytes_read < 0) {
227				perror("read");
228				exit(EXIT_FAILURE);
229			}
230
231			total_recv += bytes_read;
232		}
233
234		if (fds.revents & (POLLHUP | POLLRDHUP))
235			break;
236	}
237
238	printf("total bytes received: %zu\n", total_recv);
239	printf("rx performance: %f Gbits/s\n",
240	       get_gbps(total_recv * 8, current_nsec() - rx_begin_ns));
241	printf("total time in 'read()': %f sec\n", (float)in_read_ns / NSEC_PER_SEC);
242	printf("average time in 'read()': %f ns\n", (float)in_read_ns / read_cnt);
243	printf("POLLIN wakeups: %i\n", read_cnt);
244
245	free(data);
246	close(client_fd);
247	close(fd);
248}
249
250static void run_sender(int peer_cid, unsigned long to_send_bytes)
251{
252	time_t tx_begin_ns;
253	time_t tx_total_ns;
254	size_t total_send;
255	void *data;
256	int fd;
257
258	printf("Run as sender\n");
259	printf("Connect to %i:%u\n", peer_cid, port);
260	printf("Send %lu bytes\n", to_send_bytes);
261	printf("TX buffer %lu bytes\n", buf_size_bytes);
262
263	fd = vsock_connect(peer_cid, port);
264
265	if (fd < 0)
266		exit(EXIT_FAILURE);
267
268	data = malloc(buf_size_bytes);
269
270	if (!data) {
271		fprintf(stderr, "'malloc()' failed\n");
272		exit(EXIT_FAILURE);
273	}
274
275	memset(data, 0, buf_size_bytes);
276	total_send = 0;
277	tx_begin_ns = current_nsec();
278
279	while (total_send < to_send_bytes) {
280		ssize_t sent;
281
282		sent = write(fd, data, buf_size_bytes);
283
284		if (sent <= 0)
285			error("write");
286
287		total_send += sent;
288	}
289
290	tx_total_ns = current_nsec() - tx_begin_ns;
291
292	printf("total bytes sent: %zu\n", total_send);
293	printf("tx performance: %f Gbits/s\n",
294	       get_gbps(total_send * 8, tx_total_ns));
295	printf("total time in 'write()': %f sec\n",
296	       (float)tx_total_ns / NSEC_PER_SEC);
297
298	close(fd);
299	free(data);
300}
301
302static const char optstring[] = "";
303static const struct option longopts[] = {
304	{
305		.name = "help",
306		.has_arg = no_argument,
307		.val = 'H',
308	},
309	{
310		.name = "sender",
311		.has_arg = required_argument,
312		.val = 'S',
313	},
314	{
315		.name = "port",
316		.has_arg = required_argument,
317		.val = 'P',
318	},
319	{
320		.name = "bytes",
321		.has_arg = required_argument,
322		.val = 'M',
323	},
324	{
325		.name = "buf-size",
326		.has_arg = required_argument,
327		.val = 'B',
328	},
329	{
330		.name = "vsk-size",
331		.has_arg = required_argument,
332		.val = 'V',
333	},
334	{
335		.name = "rcvlowat",
336		.has_arg = required_argument,
337		.val = 'R',
338	},
339	{},
340};
341
342static void usage(void)
343{
344	printf("Usage: ./vsock_perf [--help] [options]\n"
345	       "\n"
346	       "This is benchmarking utility, to test vsock performance.\n"
347	       "It runs in two modes: sender or receiver. In sender mode, it\n"
348	       "connects to the specified CID and starts data transmission.\n"
349	       "\n"
350	       "Options:\n"
351	       "  --help			This message\n"
352	       "  --sender   <cid>		Sender mode (receiver default)\n"
353	       "                                <cid> of the receiver to connect to\n"
354	       "  --port     <port>		Port (default %d)\n"
355	       "  --bytes    <bytes>KMG		Bytes to send (default %d)\n"
356	       "  --buf-size <bytes>KMG		Data buffer size (default %d). In sender mode\n"
357	       "                                it is the buffer size, passed to 'write()'. In\n"
358	       "                                receiver mode it is the buffer size passed to 'read()'.\n"
359	       "  --vsk-size <bytes>KMG		Socket buffer size (default %d)\n"
360	       "  --rcvlowat <bytes>KMG		SO_RCVLOWAT value (default %d)\n"
361	       "\n", DEFAULT_PORT, DEFAULT_TO_SEND_BYTES,
362	       DEFAULT_BUF_SIZE_BYTES, DEFAULT_VSOCK_BUF_BYTES,
363	       DEFAULT_RCVLOWAT_BYTES);
364	exit(EXIT_FAILURE);
365}
366
367static long strtolx(const char *arg)
368{
369	long value;
370	char *end;
371
372	value = strtol(arg, &end, 10);
373
374	if (end != arg + strlen(arg))
375		usage();
376
377	return value;
378}
379
380int main(int argc, char **argv)
381{
382	unsigned long to_send_bytes = DEFAULT_TO_SEND_BYTES;
383	unsigned long rcvlowat_bytes = DEFAULT_RCVLOWAT_BYTES;
384	int peer_cid = -1;
385	bool sender = false;
386
387	while (1) {
388		int opt = getopt_long(argc, argv, optstring, longopts, NULL);
389
390		if (opt == -1)
391			break;
392
393		switch (opt) {
394		case 'V': /* Peer buffer size. */
395			vsock_buf_bytes = memparse(optarg);
396			break;
397		case 'R': /* SO_RCVLOWAT value. */
398			rcvlowat_bytes = memparse(optarg);
399			break;
400		case 'P': /* Port to connect to. */
401			port = strtolx(optarg);
402			break;
403		case 'M': /* Bytes to send. */
404			to_send_bytes = memparse(optarg);
405			break;
406		case 'B': /* Size of rx/tx buffer. */
407			buf_size_bytes = memparse(optarg);
408			break;
409		case 'S': /* Sender mode. CID to connect to. */
410			peer_cid = strtolx(optarg);
411			sender = true;
412			break;
413		case 'H': /* Help. */
414			usage();
415			break;
416		default:
417			usage();
418		}
419	}
420
421	if (!sender)
422		run_receiver(rcvlowat_bytes);
423	else
424		run_sender(peer_cid, to_send_bytes);
425
426	return 0;
427}
428