[RFC PATCH v2 8/8] test/vsock: vsock rx zerocopy utility

[Arseniy Krasnov]

This adds simple util for zerocopy benchmarking.

Signed-off-by: Arseniy Krasnov <AVKrasnov@xxxxxxxxxxxxxx>
---
 tools/testing/vsock/Makefile      |   1 +
 tools/testing/vsock/rx_zerocopy.c | 356 ++++++++++++++++++++++++++++++
 2 files changed, 357 insertions(+)
 create mode 100644 tools/testing/vsock/rx_zerocopy.c

diff --git a/tools/testing/vsock/Makefile b/tools/testing/vsock/Makefile
index f8293c6910c9..2cb5820ca2f3 100644
--- a/tools/testing/vsock/Makefile
+++ b/tools/testing/vsock/Makefile
@@ -3,6 +3,7 @@ all: test
 test: vsock_test vsock_diag_test
 vsock_test: vsock_test.o timeout.o control.o util.o
 vsock_diag_test: vsock_diag_test.o timeout.o control.o util.o
+rx_zerocopy: rx_zerocopy.o timeout.o control.o util.o
 
 CFLAGS += -g -O2 -Werror -Wall -I. -I../../include -I../../../usr/include -Wno-pointer-sign -fno-strict-overflow -fno-strict-aliasing -fno-common -MMD -U_FORTIFY_SOURCE -D_GNU_SOURCE
 .PHONY: all test clean
diff --git a/tools/testing/vsock/rx_zerocopy.c b/tools/testing/vsock/rx_zerocopy.c
new file mode 100644
index 000000000000..55deaa665752
--- /dev/null
+++ b/tools/testing/vsock/rx_zerocopy.c
@@ -0,0 +1,356 @@
+// SPDX-License-Identifier: GPL-2.0-only
+/*
+ * rx_zerocopy - benchmark utility for zerocopy
+ * receive.
+ *
+ * Copyright (C) 2022 SberDevices.
+ *
+ * Author: Arseniy Krasnov <AVKrasnov@xxxxxxxxxxxxxx>
+ */
+#include <getopt.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <stdbool.h>
+#include <string.h>
+#include <errno.h>
+#include <unistd.h>
+#include <time.h>
+#include <sys/mman.h>
+#include <stdint.h>
+#include <poll.h>
+#include <uapi/linux/virtio_vsock.h>
+#include <uapi/linux/vm_sockets.h>
+
+#include "util.h"
+
+#define PAGE_SIZE		4096
+
+#define DEFAULT_TX_SIZE		128
+#define DEFAULT_RX_SIZE		128
+#define DEFAULT_PORT		1234
+
+static int client_mode = 1;
+static int peer_cid = -1;
+static int port = DEFAULT_PORT;
+static unsigned long tx_buf_size;
+static unsigned long rx_buf_size;
+static unsigned long mb_to_send = 40;
+
+static time_t current_nsec(void)
+{
+	struct timespec ts;
+
+	if (clock_gettime(CLOCK_REALTIME, &ts)) {
+		perror("clock_gettime");
+		exit(EXIT_FAILURE);
+	}
+
+	return (ts.tv_sec * 1000000000ULL) + ts.tv_nsec;
+}
+
+/* Server accepts connection and */
+static void run_server(void)
+{
+	int fd;
+	char *data;
+	int client_fd;
+	union {
+		struct sockaddr sa;
+		struct sockaddr_vm svm;
+	} addr = {
+		.svm = {
+			.svm_family = AF_VSOCK,
+			.svm_port = port,
+			.svm_cid = VMADDR_CID_ANY,
+		},
+	};
+	union {
+		struct sockaddr sa;
+		struct sockaddr_vm svm;
+	} clientaddr;
+
+	socklen_t clientaddr_len = sizeof(clientaddr.svm);
+	time_t tx_begin_ns;
+	ssize_t total_send = 0;
+
+	fprintf(stderr, "Running server, listen %i, mb %lu tx buf %lu\n",
+			port, mb_to_send, tx_buf_size);
+
+	fd = socket(AF_VSOCK, SOCK_STREAM, 0);
+
+	if (fd < 0) {
+		perror("socket");
+		exit(EXIT_FAILURE);
+	}
+
+	if (bind(fd, &addr.sa, sizeof(addr.svm)) < 0) {
+		perror("bind");
+		exit(EXIT_FAILURE);
+	}
+
+	if (listen(fd, 1) < 0) {
+		perror("listen");
+		exit(EXIT_FAILURE);
+	}
+
+	client_fd = accept(fd, &clientaddr.sa, &clientaddr_len);
+
+	if (client_fd < 0) {
+		perror("accept");
+		exit(EXIT_FAILURE);
+	}
+
+	data = malloc(tx_buf_size);
+
+	if (data == NULL) {
+		fprintf(stderr, "malloc failed\n");
+		exit(EXIT_FAILURE);
+	}
+
+	memset(data, 0, tx_buf_size);
+	tx_begin_ns = current_nsec();
+
+	while (1) {
+		ssize_t sent;
+
+		if (total_send > mb_to_send * 1024 * 1024ULL)
+			break;
+
+		sent = write(client_fd, data, tx_buf_size);
+
+		if (sent <= 0) {
+			perror("write");
+			exit(EXIT_FAILURE);
+		}
+
+		total_send += sent;
+	}
+
+	free(data);
+
+	fprintf(stderr, "Total %zi MB, time %f\n", mb_to_send,
+			(float)(current_nsec() - tx_begin_ns)/1000.0/1000.0/1000.0);
+
+	close(fd);
+	close(client_fd);
+}
+
+static void run_client(int zerocopy)
+{
+	int fd;
+	union {
+		struct sockaddr sa;
+		struct sockaddr_vm svm;
+	} addr = {
+		.svm = {
+			.svm_family = AF_VSOCK,
+			.svm_port = port,
+			.svm_cid = peer_cid,
+		},
+	};
+	unsigned long sum = 0;
+	void *rx_va = NULL;
+	unsigned long zc_on = 1;
+
+	printf("Running client, %s mode, peer %i:%i, rx buf %lu\n",
+		zerocopy ? "zerocopy" : "copy", peer_cid, port,
+		rx_buf_size);
+
+	fd = socket(AF_VSOCK, SOCK_STREAM, 0);
+
+	if (fd < 0) {
+		perror("socket");
+		exit(EXIT_FAILURE);
+	}
+
+	if (connect(fd, &addr.sa, sizeof(addr.svm))) {
+		perror("connect");
+		exit(EXIT_FAILURE);
+	}
+
+	if (setsockopt(fd, AF_VSOCK, SO_VM_SOCKETS_ZEROCOPY,
+				(void *)&zc_on, sizeof(zc_on))) {
+		perror("setsockopt");
+		exit(EXIT_FAILURE);
+	}
+
+	if (zerocopy) {
+		rx_va = mmap(NULL, rx_buf_size,
+				PROT_READ, MAP_SHARED, fd, 0);
+
+		if (rx_va == MAP_FAILED) {
+			perror("mmap");
+			exit(EXIT_FAILURE);
+		}
+	}
+
+	while (1) {
+		struct pollfd fds = { 0 };
+		int done = 0;
+
+		fds.fd = fd;
+		fds.events = POLLIN | POLLERR | POLLHUP |
+			     POLLRDHUP | POLLNVAL;
+
+		if (poll(&fds, 1, -1) < 0) {
+			perror("poll");
+			exit(EXIT_FAILURE);
+		}
+
+		if (fds.revents & (POLLHUP | POLLRDHUP))
+			done = 1;
+
+		if (fds.revents & POLLERR) {
+			fprintf(stderr, "Done error\n");
+			break;
+		}
+
+		if (fds.revents & POLLIN) {
+			if (zerocopy) {
+				struct virtio_vsock_usr_hdr *hdr;
+				uintptr_t tmp_rx_va = (uintptr_t)rx_va;
+				socklen_t len = sizeof(tmp_rx_va);
+
+				if (getsockopt(fd, AF_VSOCK, SO_VM_SOCKETS_MAP_RX,
+						&tmp_rx_va, &len) < 0) {
+					perror("getsockopt");
+					exit(EXIT_FAILURE);
+				}
+
+				hdr = (struct virtio_vsock_usr_hdr *)tmp_rx_va;
+
+				if (!hdr->len) {
+					if (done) {
+						fprintf(stderr, "Done, sum %lu\n", sum);
+						break;
+					}
+				}
+
+				tmp_rx_va += PAGE_SIZE;
+
+				if (madvise((void *)rx_va, rx_buf_size,
+						MADV_DONTNEED)) {
+					perror("madvise");
+					exit(EXIT_FAILURE);
+				}
+			} else {
+				char data[rx_buf_size - PAGE_SIZE];
+				ssize_t bytes_read;
+
+				bytes_read = read(fd, data, sizeof(data));
+
+				if (bytes_read <= 0)
+					break;
+			}
+		}
+	}
+}
+
+static const char optstring[] = "";
+static const struct option longopts[] = {
+	{
+		.name = "mode",
+		.has_arg = required_argument,
+		.val = 'm',
+	},
+	{
+		.name = "zerocopy",
+		.has_arg = no_argument,
+		.val = 'z',
+	},
+	{
+		.name = "cid",
+		.has_arg = required_argument,
+		.val = 'c',
+	},
+	{
+		.name = "port",
+		.has_arg = required_argument,
+		.val = 'p',
+	},
+	{
+		.name = "mb",
+		.has_arg = required_argument,
+		.val = 's',
+	},
+	{
+		.name = "tx",
+		.has_arg = required_argument,
+		.val = 't',
+	},
+	{
+		.name = "rx",
+		.has_arg = required_argument,
+		.val = 'r',
+	},
+	{
+		.name = "help",
+		.has_arg = no_argument,
+		.val = '?',
+	},
+	{},
+};
+
+int main(int argc, char **argv)
+{
+	int zerocopy = 0;
+
+	for (;;) {
+		int opt = getopt_long(argc, argv, optstring, longopts, NULL);
+
+		if (opt == -1)
+			break;
+
+		switch (opt) {
+		case 's':
+			mb_to_send = atoi(optarg);
+			break;
+		case 'c':
+			peer_cid = atoi(optarg);
+			break;
+		case 'p':
+			port = atoi(optarg);
+			break;
+		case 'r':
+			rx_buf_size = atoi(optarg);
+			break;
+		case 't':
+			tx_buf_size = atoi(optarg);
+			break;
+		case 'm':
+			if (strcmp(optarg, "client") == 0)
+				client_mode = 1;
+			else if (strcmp(optarg, "server") == 0)
+				client_mode = 0;
+			else {
+				fprintf(stderr, "--mode must be \"client\" or \"server\"\n");
+				return EXIT_FAILURE;
+			}
+			break;
+		case 'z':
+			zerocopy = 1;
+			break;
+		default:
+			break;
+		}
+
+	}
+
+	if (!tx_buf_size)
+		tx_buf_size = DEFAULT_TX_SIZE;
+
+	if (!rx_buf_size)
+		rx_buf_size = DEFAULT_RX_SIZE;
+
+	tx_buf_size *= PAGE_SIZE;
+	rx_buf_size *= PAGE_SIZE;
+
+	srand(time(NULL));
+
+	if (client_mode)
+		run_client(zerocopy);
+	else
+		run_server();
+
+	return 0;
+}
-- 
2.25.1