// SPDX-License-Identifier: LGPL-2.1
#define _GNU_SOURCE
#include <assert.h>
#include <pthread.h>
#include <sched.h>
#include <signal.h>
#include <stdbool.h>
#include <stdio.h>
#include <string.h>
#include <syscall.h>
#include <unistd.h>

#include <linux/prctl.h>
#include <sys/prctl.h>
#include <sys/time.h>

#include "rseq.h"

#include "../kselftest_harness.h"

#ifndef __NR_rseq_slice_yield
# define __NR_rseq_slice_yield	471
#endif

#define BITS_PER_INT	32
#define BITS_PER_BYTE	8

#ifndef PR_RSEQ_SLICE_EXTENSION
# define PR_RSEQ_SLICE_EXTENSION		79
#  define PR_RSEQ_SLICE_EXTENSION_GET		1
#  define PR_RSEQ_SLICE_EXTENSION_SET		2
#  define PR_RSEQ_SLICE_EXT_ENABLE		0x01
#endif

#ifndef RSEQ_SLICE_EXT_REQUEST_BIT
# define RSEQ_SLICE_EXT_REQUEST_BIT	0
# define RSEQ_SLICE_EXT_GRANTED_BIT	1
#endif

#ifndef asm_inline
# define asm_inline	asm __inline
#endif

#define NSEC_PER_SEC	1000000000L
#define NSEC_PER_USEC	      1000L

struct noise_params {
	int64_t	noise_nsecs;
	int64_t	sleep_nsecs;
	int64_t	run;
};

FIXTURE(slice_ext)
{
	pthread_t		noise_thread;
	struct noise_params	noise_params;
};

FIXTURE_VARIANT(slice_ext)
{
	int64_t	total_nsecs;
	int64_t	slice_nsecs;
	int64_t	noise_nsecs;
	int64_t	sleep_nsecs;
	bool	no_yield;
};

FIXTURE_VARIANT_ADD(slice_ext, n2_2_50)
{
	.total_nsecs	=  5LL * NSEC_PER_SEC,
	.slice_nsecs	=  2LL * NSEC_PER_USEC,
	.noise_nsecs    =  2LL * NSEC_PER_USEC,
	.sleep_nsecs	= 50LL * NSEC_PER_USEC,
};

FIXTURE_VARIANT_ADD(slice_ext, n50_2_50)
{
	.total_nsecs	=  5LL * NSEC_PER_SEC,
	.slice_nsecs	= 50LL * NSEC_PER_USEC,
	.noise_nsecs    =  2LL * NSEC_PER_USEC,
	.sleep_nsecs	= 50LL * NSEC_PER_USEC,
};

FIXTURE_VARIANT_ADD(slice_ext, n2_2_50_no_yield)
{
	.total_nsecs	=  5LL * NSEC_PER_SEC,
	.slice_nsecs	=  2LL * NSEC_PER_USEC,
	.noise_nsecs    =  2LL * NSEC_PER_USEC,
	.sleep_nsecs	= 50LL * NSEC_PER_USEC,
	.no_yield	= true,
};


static inline bool elapsed(struct timespec *start, struct timespec *now,
			   int64_t span)
{
	int64_t delta = now->tv_sec - start->tv_sec;

	delta *= NSEC_PER_SEC;
	delta += now->tv_nsec - start->tv_nsec;
	return delta >= span;
}

static void *noise_thread(void *arg)
{
	struct noise_params *p = arg;

	while (RSEQ_READ_ONCE(p->run)) {
		struct timespec ts_start, ts_now;

		clock_gettime(CLOCK_MONOTONIC, &ts_start);
		do {
			clock_gettime(CLOCK_MONOTONIC, &ts_now);
		} while (!elapsed(&ts_start, &ts_now, p->noise_nsecs));

		ts_start.tv_sec = 0;
		ts_start.tv_nsec = p->sleep_nsecs;
		clock_nanosleep(CLOCK_MONOTONIC, 0, &ts_start, NULL);
	}
	return NULL;
}

FIXTURE_SETUP(slice_ext)
{
	cpu_set_t affinity;

	if (__rseq_register_current_thread(true, false))
		SKIP(return, "RSEQ not supported\n");

	if (prctl(PR_RSEQ_SLICE_EXTENSION, PR_RSEQ_SLICE_EXTENSION_SET,
		  PR_RSEQ_SLICE_EXT_ENABLE, 0, 0))
		SKIP(return, "Time slice extension not supported\n");

	ASSERT_EQ(sched_getaffinity(0, sizeof(affinity), &affinity), 0);

	/* Pin it on a single CPU. Avoid CPU 0 */
	for (int i = 1; i < CPU_SETSIZE; i++) {
		if (!CPU_ISSET(i, &affinity))
			continue;

		CPU_ZERO(&affinity);
		CPU_SET(i, &affinity);
		ASSERT_EQ(sched_setaffinity(0, sizeof(affinity), &affinity), 0);
		break;
	}

	self->noise_params.noise_nsecs = variant->noise_nsecs;
	self->noise_params.sleep_nsecs = variant->sleep_nsecs;
	self->noise_params.run = 1;

	ASSERT_EQ(pthread_create(&self->noise_thread, NULL, noise_thread, &self->noise_params), 0);
}

FIXTURE_TEARDOWN(slice_ext)
{
	self->noise_params.run = 0;
	pthread_join(self->noise_thread, NULL);
}

TEST_F(slice_ext, slice_test)
{
	unsigned long success = 0, yielded = 0, scheduled = 0, raced = 0;
	unsigned long total = 0, aborted = 0;
	struct rseq_abi *rs = rseq_get_abi();
	struct timespec ts_start, ts_now;

	ASSERT_NE(rs, NULL);

	clock_gettime(CLOCK_MONOTONIC, &ts_start);
	do {
		struct timespec ts_cs;
		bool req = false;

		clock_gettime(CLOCK_MONOTONIC, &ts_cs);

		total++;
		RSEQ_WRITE_ONCE(rs->slice_ctrl.request, 1);
		do {
			clock_gettime(CLOCK_MONOTONIC, &ts_now);
		} while (!elapsed(&ts_cs, &ts_now, variant->slice_nsecs));

		/*
		 * request can be cleared unconditionally, but for making
		 * the stats work this is actually checking it first
		 */
		if (RSEQ_READ_ONCE(rs->slice_ctrl.request)) {
			RSEQ_WRITE_ONCE(rs->slice_ctrl.request, 0);
			/* Race between check and clear! */
			req = true;
			success++;
		}

		if (RSEQ_READ_ONCE(rs->slice_ctrl.granted)) {
			/* The above raced against a late grant */
			if (req)
				success--;
			if (variant->no_yield) {
				syscall(__NR_getpid);
				aborted++;
			} else {
				yielded++;
				if (!syscall(__NR_rseq_slice_yield))
					raced++;
			}
		} else {
			if (!req)
				scheduled++;
		}

		clock_gettime(CLOCK_MONOTONIC, &ts_now);
	} while (!elapsed(&ts_start, &ts_now, variant->total_nsecs));

	printf("# Total     %12ld\n", total);
	printf("# Success   %12ld\n", success);
	printf("# Yielded   %12ld\n", yielded);
	printf("# Aborted   %12ld\n", aborted);
	printf("# Scheduled %12ld\n", scheduled);
	printf("# Raced     %12ld\n", raced);
}

TEST_HARNESS_MAIN