// SPDX-License-Identifier: GPL-2.0-only
#include <pthread.h>
#include <sys/ioctl.h>
#include <sys/mman.h>

#include <linux/sizes.h>
#include <linux/time64.h>
#include <linux/vfio.h>

#include <libvfio.h>

#include "../kselftest_harness.h"

static char **device_bdfs;
static int nr_devices;

struct thread_args {
	struct iommu *iommu;
	int device_index;
	struct timespec start;
	struct timespec end;
	pthread_barrier_t *barrier;
};

FIXTURE(vfio_pci_device_init_perf_test) {
	pthread_t *threads;
	pthread_barrier_t barrier;
	struct thread_args *thread_args;
	struct iommu *iommu;
};

FIXTURE_VARIANT(vfio_pci_device_init_perf_test) {
	const char *iommu_mode;
};

#define FIXTURE_VARIANT_ADD_IOMMU_MODE(_iommu_mode)			\
FIXTURE_VARIANT_ADD(vfio_pci_device_init_perf_test, _iommu_mode) {	\
	.iommu_mode = #_iommu_mode,					\
}

FIXTURE_VARIANT_ADD_ALL_IOMMU_MODES();

FIXTURE_SETUP(vfio_pci_device_init_perf_test)
{
	int i;

	self->iommu = iommu_init(variant->iommu_mode);
	self->threads = calloc(nr_devices, sizeof(self->threads[0]));
	self->thread_args = calloc(nr_devices, sizeof(self->thread_args[0]));

	pthread_barrier_init(&self->barrier, NULL, nr_devices);

	for (i = 0; i < nr_devices; i++) {
		self->thread_args[i].iommu = self->iommu;
		self->thread_args[i].barrier = &self->barrier;
		self->thread_args[i].device_index = i;
	}
}

FIXTURE_TEARDOWN(vfio_pci_device_init_perf_test)
{
	iommu_cleanup(self->iommu);
	free(self->threads);
	free(self->thread_args);
}

static s64 to_ns(struct timespec ts)
{
	return (s64)ts.tv_nsec + NSEC_PER_SEC * (s64)ts.tv_sec;
}

static struct timespec to_timespec(s64 ns)
{
	struct timespec ts = {
		.tv_nsec = ns % NSEC_PER_SEC,
		.tv_sec = ns / NSEC_PER_SEC,
	};

	return ts;
}

static struct timespec timespec_sub(struct timespec a, struct timespec b)
{
	return to_timespec(to_ns(a) - to_ns(b));
}

static struct timespec timespec_min(struct timespec a, struct timespec b)
{
	return to_ns(a) < to_ns(b) ? a : b;
}

static struct timespec timespec_max(struct timespec a, struct timespec b)
{
	return to_ns(a) > to_ns(b) ? a : b;
}

static void *thread_main(void *__args)
{
	struct thread_args *args = __args;
	struct vfio_pci_device *device;

	pthread_barrier_wait(args->barrier);

	clock_gettime(CLOCK_MONOTONIC, &args->start);
	device = vfio_pci_device_init(device_bdfs[args->device_index], args->iommu);
	clock_gettime(CLOCK_MONOTONIC, &args->end);

	pthread_barrier_wait(args->barrier);

	vfio_pci_device_cleanup(device);
	return NULL;
}

TEST_F(vfio_pci_device_init_perf_test, init)
{
	struct timespec start = to_timespec(INT64_MAX), end = {};
	struct timespec min = to_timespec(INT64_MAX);
	struct timespec max = {};
	struct timespec avg = {};
	struct timespec wall_time;
	s64 thread_ns = 0;
	int i;

	for (i = 0; i < nr_devices; i++) {
		pthread_create(&self->threads[i], NULL, thread_main,
			       &self->thread_args[i]);
	}

	for (i = 0; i < nr_devices; i++) {
		struct thread_args *args = &self->thread_args[i];
		struct timespec init_time;

		pthread_join(self->threads[i], NULL);

		start = timespec_min(start, args->start);
		end = timespec_max(end, args->end);

		init_time = timespec_sub(args->end, args->start);
		min = timespec_min(min, init_time);
		max = timespec_max(max, init_time);
		thread_ns += to_ns(init_time);
	}

	avg = to_timespec(thread_ns / nr_devices);
	wall_time = timespec_sub(end, start);

	printf("Wall time: %lu.%09lus\n",
	       wall_time.tv_sec, wall_time.tv_nsec);
	printf("Min init time (per device): %lu.%09lus\n",
	       min.tv_sec, min.tv_nsec);
	printf("Max init time (per device): %lu.%09lus\n",
	       max.tv_sec, max.tv_nsec);
	printf("Avg init time (per device): %lu.%09lus\n",
	       avg.tv_sec, avg.tv_nsec);
}

int main(int argc, char *argv[])
{
	int i;

	device_bdfs = vfio_selftests_get_bdfs(&argc, argv, &nr_devices);

	printf("Testing parallel initialization of %d devices:\n", nr_devices);
	for (i = 0; i < nr_devices; i++)
		printf("    %s\n", device_bdfs[i]);

	return test_harness_run(argc, argv);
}