/* SPDX-License-Identifier: GPL-2.0-only OR GPL-3.0-only */
/* Copyright (c) 2024 Brett Sheffield <bacs@librecast.net> */

/* router bandwidth test */

#include "test.h"
#include <librecast/router.h>
#include <pthread.h>
#include <semaphore.h>
#include <sys/poll.h>
#include <sys/socket.h>
#include <time.h>
#include <unistd.h>

#define USE_POLL 0
#define NANO 1000000000ULL
#define DATALEN 512 * 1024 * 1024 /* 512 MiB */
#define PAYLOAD 1024

enum {
	SOCK_SEND,
	SOCK_RECV,
};

static struct timespec start, end;
static char *buf;

static char *generate_source_object(size_t F)
{
	char *obj = malloc(F);
	assert(obj);
	arc4random_buf(obj, F);
	test_log("source object generated (%zu bytes)\n", F);
	return obj;
}

static void print_stats(size_t bytes, struct timespec *start, struct timespec *end)
{
	double Bs;
	double bps;
	double s = (end->tv_sec * NANO + end->tv_nsec);
	s -= start->tv_sec * NANO + start->tv_nsec;
	s /= NANO;
	Bs = bytes / s;
	bps = Bs * 8;
	test_log("%zu bytes sent in %lfs (%lf Gbps)\n", bytes, s, bps / 1000 / 1000 / 1000);
}

static void *thread_recv(void *arg)
{
	struct iovec iov;
	struct msghdr msg = {0};
#if USE_POLL
	struct pollfd fds = {0};
	int rc;
#endif
	size_t bytes_recv = 0;
	ssize_t byt;
	int sock = *(int *)arg;

	msg.msg_iov = &iov;
	msg.msg_iovlen = 1;
	iov.iov_base = buf;
	iov.iov_len = DATALEN;
#if USE_POLL
	fds.fd = sock;
	fds.events = POLLIN;
#endif
	test_log("recving on %i\n", sock);
	while (bytes_recv < DATALEN) {
#if USE_POLL
		rc = poll(&fds, 1, 10000);
		if (rc == -1) {
			perror("poll");
			break;
		}
		byt = recvmsg(sock, &msg, MSG_DONTWAIT);
#else
		byt = recvmsg(sock, &msg, 0);
#endif
		if (byt > 0) {
			//test_log("read %zi bytes (tot = %zu)\n", byt, bytes_recv);
			bytes_recv += byt;
			iov.iov_base = (char *)iov.iov_base + byt;
			iov.iov_len -= byt;
		}
		else if (byt == -1) {
			perror("recvmsg");
			break;
		}
	}
	__atomic_thread_fence(__ATOMIC_SEQ_CST);
	clock_gettime(CLOCK_REALTIME, &end);
	return NULL;
}

static void measure_bandwidth_memcpy(char *data, size_t len)
{
	char *dst = malloc(len);
	assert(dst);
	test_log("memcpy\n");
	clock_gettime(CLOCK_REALTIME, &start);
	__atomic_thread_fence(__ATOMIC_SEQ_CST);
	memcpy(dst, data, len);
	__atomic_thread_fence(__ATOMIC_SEQ_CST);
	clock_gettime(CLOCK_REALTIME, &end);
	free(dst);
	print_stats(len, &start, &end);
}

static void measure_bandwidth(char *data, size_t len, int sock_send, int sock_recv)
{
	char *ptr;
	struct iovec iov;
	struct msghdr msg = {0};
#if USE_POLL
	struct pollfd fds = {0};
	int rc;
#endif
	ssize_t byt;
	size_t bytes_sent = 0;
	pthread_t tid_recv;

	buf = malloc(DATALEN);
	if (pthread_create(&tid_recv, NULL, &thread_recv, &sock_recv) == -1) {
		perror("pthread_create");
		exit(EXIT_FAILURE);
	}

#if USE_POLL
	fds.fd = sock_send;
	fds.events = POLLOUT;
#endif
	msg.msg_iov = &iov;
	msg.msg_iovlen = 1;
	iov.iov_len = PAYLOAD;
	ptr = data;
	test_log("sending on %i\n", sock_send);
	clock_gettime(CLOCK_REALTIME, &start);
	__atomic_thread_fence(__ATOMIC_SEQ_CST);
	while (bytes_sent < len) {
		iov.iov_base = ptr;
#if USE_POLL
		rc = poll(&fds, 1, 10000);
		if (rc == -1) {
			perror("poll");
			break;
		}
		byt = sendmsg(sock_send, &msg, MSG_DONTWAIT);
#else
		byt = sendmsg(sock_send, &msg, 0);
#endif
		if (byt == -1) {
			if (errno == ENOBUFS) continue;
			perror("sendmsg");
			break;
		}
		else if (byt > 0) {
			//test_log("send %zi bytes (tot = %zu)\n", byt, bytes_sent);
			ptr += byt;
			bytes_sent += byt;
		}
	}
	pthread_join(tid_recv, NULL);
	test_assert(memcmp(buf, data, len) == 0, "data matches");
	free(buf);
	print_stats(bytes_sent, &start, &end);
}

static void *router_ready(void *arg)
{
	test_log("%s\n", __func__);
	sem_post(arg);
	return NULL;
}

static void router_tests(char *data, size_t len)
{
	/* now create a router */
	lc_ctx_t *lctx;
	lc_router_t *r;
	lc_socket_t *s[4];
	sem_t sem_ready;
	int rc;

	/* create context and router */
	lctx = lc_ctx_new();
	if (!test_assert(lctx != NULL, "lc_ctx_new()")) return;
	r = lc_router_new(lctx, 2, LC_ROUTER_FLAG_FIXED);
	if (!test_assert(r != NULL, "lc_router_new() - router created"))
		goto err_ctx_free;

	/* use one thread to preserve data ordering */
	int nthreads = 1;
	rc = lc_ctx_qpool_init(lctx, nthreads);

	/* create socketpairs and connect to router */
	/* socketpair (router) socketpair */
	/* s[1]-s[0]-(p0[r]p1)-s[2]-s[3] */
	for (int i = 0; i < 4; i += 2) {
		rc = lc_socketpair(lctx, &s[i]);
		test_assert(rc == 0, "lc_socketpair()[%i]", i);
		rc = lc_router_socket_add(r, s[i]);
		test_assert(rc == 0, "lc_router_socket_add()[%i]", i);
	}

	/* bring up router ports, enable forwarding and start router */
	lc_router_port_up(r, -1);
	lc_router_port_set(r, -1, LC_PORT_FWD);
	lc_router_port_set(r, -1, LC_PORT_FLOOD);
	sem_init(&sem_ready, 0, 0);
	lc_router_onready(r, &router_ready, &sem_ready);
	rc = lc_router_start(r);
	if (!test_assert(rc == 0, "lc_router_start()")) goto err_ctx_free;
	sem_wait(&sem_ready); /* wait until router ready */
	sem_destroy(&sem_ready);
	measure_bandwidth(data, len, lc_socket_raw(s[1]), lc_socket_raw(s[3]));
err_ctx_free:
	lc_ctx_free(lctx);
}

int main(void)
{
	char *data;
	int sv[2];
	int rc;
	test_name("router bandwidth test");

	/* generate random data */
	data = generate_source_object(DATALEN);

	/* measure memcpy for reference */
	measure_bandwidth_memcpy(data, DATALEN);

	/* measure bandwidth for a straight socketpair */
	rc = socketpair(AF_UNIX, SOCK_DGRAM, 0, sv);
	test_assert(rc == 0, "socketpair()");
	test_log("sockpair created %i/%i\n", sv[0], sv[1]);
	measure_bandwidth(data, DATALEN, sv[SOCK_SEND], sv[SOCK_RECV]);
	close(sv[0]); close(sv[1]);

	/* now run router tests */
	router_tests(data, DATALEN);

	free(data);
	return test_status;
}