/* SPDX-License-Identifier: GPL-2.0 */
/*
 * A demo sched_ext user space scheduler which provides vruntime semantics
 * using a simple ordered-list implementation.
 *
 * Each CPU in the system resides in a single, global domain. This precludes
 * the need to do any load balancing between domains. The scheduler could
 * easily be extended to support multiple domains, with load balancing
 * happening in user space.
 *
 * Any task which has any CPU affinity is scheduled entirely in BPF. This
 * program only schedules tasks which may run on any CPU.
 *
 * Copyright (c) 2022 Meta Platforms, Inc. and affiliates.
 * Copyright (c) 2022 Tejun Heo <tj@kernel.org>
 * Copyright (c) 2022 David Vernet <dvernet@meta.com>
 */
#include <stdio.h>
#include <unistd.h>
#include <sched.h>
#include <signal.h>
#include <assert.h>
#include <libgen.h>
#include <pthread.h>
#include <bpf/bpf.h>
#include <sys/mman.h>
#include <sys/queue.h>
#include <sys/syscall.h>

#include <scx/common.h>
#include "scx_userland.h"
#include "scx_userland.bpf.skel.h"

const char help_fmt[] =
"A minimal userland sched_ext scheduler.\n"
"\n"
"See the top-level comment in .bpf.c for more details.\n"
"\n"
"Try to reduce `sysctl kernel.pid_max` if this program triggers OOMs.\n"
"\n"
"Usage: %s [-b BATCH] [-v]\n"
"\n"
"  -b BATCH      The number of tasks to batch when dispatching (default: 8)\n"
"  -v            Print libbpf debug messages\n"
"  -h            Display this help and exit\n";

/* Defined in UAPI */
#define SCHED_EXT 7

/* Number of tasks to batch when dispatching to user space. */
static __u32 batch_size = 8;

static bool verbose;
static volatile int exit_req;
static int enqueued_fd, dispatched_fd;

static pthread_t stats_printer;
static struct scx_userland *skel;
static struct bpf_link *ops_link;

/* Stats collected in user space. */
static __u64 nr_vruntime_enqueues, nr_vruntime_dispatches, nr_vruntime_failed;

/* Number of tasks currently enqueued. */
static __u64 nr_curr_enqueued;

/* The data structure containing tasks that are enqueued in user space. */
struct enqueued_task {
	LIST_ENTRY(enqueued_task) entries;
	__u64 sum_exec_runtime;
	double vruntime;
};

/*
 * Use a vruntime-sorted list to store tasks. This could easily be extended to
 * a more optimal data structure, such as an rbtree as is done in CFS. We
 * currently elect to use a sorted list to simplify the example for
 * illustrative purposes.
 */
LIST_HEAD(listhead, enqueued_task);

/*
 * A vruntime-sorted list of tasks. The head of the list contains the task with
 * the lowest vruntime. That is, the task that has the "highest" claim to be
 * scheduled.
 */
static struct listhead vruntime_head = LIST_HEAD_INITIALIZER(vruntime_head);

/*
 * The main array of tasks. The array is allocated all at once during
 * initialization, based on /proc/sys/kernel/pid_max, to avoid having to
 * dynamically allocate memory on the enqueue path, which could cause a
 * deadlock. A more substantive user space scheduler could e.g. provide a hook
 * for newly enabled tasks that are passed to the scheduler from the
 * .prep_enable() callback to allows the scheduler to allocate on safe paths.
 */
struct enqueued_task *tasks;
static int pid_max;

static double min_vruntime;

static int libbpf_print_fn(enum libbpf_print_level level, const char *format, va_list args)
{
	if (level == LIBBPF_DEBUG && !verbose)
		return 0;
	return vfprintf(stderr, format, args);
}

static void sigint_handler(int userland)
{
	exit_req = 1;
}

static int get_pid_max(void)
{
	FILE *fp;
	int pid_max;

	fp = fopen("/proc/sys/kernel/pid_max", "r");
	if (fp == NULL) {
		fprintf(stderr, "Error opening /proc/sys/kernel/pid_max\n");
		return -1;
	}
	if (fscanf(fp, "%d", &pid_max) != 1) {
		fprintf(stderr, "Error reading from /proc/sys/kernel/pid_max\n");
		fclose(fp);
		return -1;
	}
	fclose(fp);

	return pid_max;
}

static int init_tasks(void)
{
	pid_max = get_pid_max();
	if (pid_max < 0)
		return pid_max;

	tasks = calloc(pid_max, sizeof(*tasks));
	if (!tasks) {
		fprintf(stderr, "Error allocating tasks array\n");
		return -ENOMEM;
	}

	return 0;
}

static __u32 task_pid(const struct enqueued_task *task)
{
	return ((uintptr_t)task - (uintptr_t)tasks) / sizeof(*task);
}

static int dispatch_task(__s32 pid)
{
	int err;

	err = bpf_map_update_elem(dispatched_fd, NULL, &pid, 0);
	if (err) {
		__atomic_add_fetch(&nr_vruntime_failed, 1, __ATOMIC_RELAXED);
	} else {
		__atomic_add_fetch(&nr_vruntime_dispatches, 1, __ATOMIC_RELAXED);
	}

	return err;
}

static struct enqueued_task *get_enqueued_task(__s32 pid)
{
	if (pid >= pid_max)
		return NULL;

	return &tasks[pid];
}

static double calc_vruntime_delta(__u64 weight, __u64 delta)
{
	double weight_f = (double)weight / 100.0;
	double delta_f = (double)delta;

	return delta_f / weight_f;
}

static void update_enqueued(struct enqueued_task *enqueued, const struct scx_userland_enqueued_task *bpf_task)
{
	__u64 delta;

	delta = bpf_task->sum_exec_runtime - enqueued->sum_exec_runtime;

	enqueued->vruntime += calc_vruntime_delta(bpf_task->weight, delta);
	if (min_vruntime > enqueued->vruntime)
		enqueued->vruntime = min_vruntime;
	enqueued->sum_exec_runtime = bpf_task->sum_exec_runtime;
}

static int vruntime_enqueue(const struct scx_userland_enqueued_task *bpf_task)
{
	struct enqueued_task *curr, *enqueued, *prev;

	curr = get_enqueued_task(bpf_task->pid);
	if (!curr)
		return ENOENT;

	update_enqueued(curr, bpf_task);
	__atomic_add_fetch(&nr_vruntime_enqueues, 1, __ATOMIC_RELAXED);
	__atomic_add_fetch(&nr_curr_enqueued, 1, __ATOMIC_RELAXED);

	/*
	 * Enqueue the task in a vruntime-sorted list. A more optimal data
	 * structure such as an rbtree could easily be used as well. We elect
	 * to use a list here simply because it's less code, and thus the
	 * example is less convoluted and better serves to illustrate what a
	 * user space scheduler could look like.
	 */

	if (LIST_EMPTY(&vruntime_head)) {
		LIST_INSERT_HEAD(&vruntime_head, curr, entries);
		return 0;
	}

	LIST_FOREACH(enqueued, &vruntime_head, entries) {
		if (curr->vruntime <= enqueued->vruntime) {
			LIST_INSERT_BEFORE(enqueued, curr, entries);
			return 0;
		}
		prev = enqueued;
	}

	LIST_INSERT_AFTER(prev, curr, entries);

	return 0;
}

static void drain_enqueued_map(void)
{
	while (1) {
		struct scx_userland_enqueued_task task;
		int err;

		if (bpf_map_lookup_and_delete_elem(enqueued_fd, NULL, &task)) {
			skel->bss->nr_queued = 0;
			skel->bss->nr_scheduled = nr_curr_enqueued;
			return;
		}

		err = vruntime_enqueue(&task);
		if (err) {
			fprintf(stderr, "Failed to enqueue task %d: %s\n",
				task.pid, strerror(err));
			exit_req = 1;
			return;
		}
	}
}

static void dispatch_batch(void)
{
	__u32 i;

	for (i = 0; i < batch_size; i++) {
		struct enqueued_task *task;
		int err;
		__s32 pid;

		task = LIST_FIRST(&vruntime_head);
		if (!task)
			break;

		min_vruntime = task->vruntime;
		pid = task_pid(task);
		LIST_REMOVE(task, entries);
		err = dispatch_task(pid);
		if (err) {
			/*
			 * If we fail to dispatch, put the task back to the
			 * vruntime_head list and stop dispatching additional
			 * tasks in this batch.
			 */
			LIST_INSERT_HEAD(&vruntime_head, task, entries);
			break;
		}
		__atomic_sub_fetch(&nr_curr_enqueued, 1, __ATOMIC_RELAXED);
	}
	skel->bss->nr_scheduled = __atomic_load_n(&nr_curr_enqueued, __ATOMIC_RELAXED);
}

static void *run_stats_printer(void *arg)
{
	while (!exit_req) {
		__u64 nr_failed_enqueues, nr_kernel_enqueues, nr_user_enqueues, total;

		nr_failed_enqueues = skel->bss->nr_failed_enqueues;
		nr_kernel_enqueues = skel->bss->nr_kernel_enqueues;
		nr_user_enqueues = skel->bss->nr_user_enqueues;
		total = nr_failed_enqueues + nr_kernel_enqueues + nr_user_enqueues;

		printf("o-----------------------o\n");
		printf("| BPF ENQUEUES          |\n");
		printf("|-----------------------|\n");
		printf("|  kern:     %10llu |\n", nr_kernel_enqueues);
		printf("|  user:     %10llu |\n", nr_user_enqueues);
		printf("|  failed:   %10llu |\n", nr_failed_enqueues);
		printf("|  -------------------- |\n");
		printf("|  total:    %10llu |\n", total);
		printf("|                       |\n");
		printf("|-----------------------|\n");
		printf("| VRUNTIME / USER       |\n");
		printf("|-----------------------|\n");
		printf("|  enq:      %10llu |\n", __atomic_load_n(&nr_vruntime_enqueues, __ATOMIC_RELAXED));
		printf("|  disp:     %10llu |\n", __atomic_load_n(&nr_vruntime_dispatches, __ATOMIC_RELAXED));
		printf("|  failed:   %10llu |\n", __atomic_load_n(&nr_vruntime_failed, __ATOMIC_RELAXED));
		printf("o-----------------------o\n");
		printf("\n\n");
		fflush(stdout);
		sleep(1);
	}

	return NULL;
}

static int spawn_stats_thread(void)
{
	return pthread_create(&stats_printer, NULL, run_stats_printer, NULL);
}

static void pre_bootstrap(int argc, char **argv)
{
	int err;
	__u32 opt;
	struct sched_param sched_param = {
		.sched_priority = sched_get_priority_max(SCHED_EXT),
	};

	err = init_tasks();
	if (err)
		exit(err);

	libbpf_set_print(libbpf_print_fn);
	signal(SIGINT, sigint_handler);
	signal(SIGTERM, sigint_handler);

	/*
	 * Enforce that the user scheduler task is managed by sched_ext. The
	 * task eagerly drains the list of enqueued tasks in its main work
	 * loop, and then yields the CPU. The BPF scheduler only schedules the
	 * user space scheduler task when at least one other task in the system
	 * needs to be scheduled.
	 */
	err = syscall(__NR_sched_setscheduler, getpid(), SCHED_EXT, &sched_param);
	SCX_BUG_ON(err, "Failed to set scheduler to SCHED_EXT");

	while ((opt = getopt(argc, argv, "b:vh")) != -1) {
		switch (opt) {
		case 'b':
			batch_size = strtoul(optarg, NULL, 0);
			break;
		case 'v':
			verbose = true;
			break;
		default:
			fprintf(stderr, help_fmt, basename(argv[0]));
			exit(opt != 'h');
		}
	}

	/*
	 * It's not always safe to allocate in a user space scheduler, as an
	 * enqueued task could hold a lock that we require in order to be able
	 * to allocate.
	 */
	err = mlockall(MCL_CURRENT | MCL_FUTURE);
	SCX_BUG_ON(err, "Failed to prefault and lock address space");
}

static void bootstrap(char *comm)
{
	exit_req = 0;
	min_vruntime = 0.0;
	__atomic_store_n(&nr_vruntime_enqueues, 0, __ATOMIC_RELAXED);
	__atomic_store_n(&nr_vruntime_dispatches, 0, __ATOMIC_RELAXED);
	__atomic_store_n(&nr_vruntime_failed, 0, __ATOMIC_RELAXED);
	__atomic_store_n(&nr_curr_enqueued, 0, __ATOMIC_RELAXED);
	memset(tasks, 0, pid_max * sizeof(*tasks));
	LIST_INIT(&vruntime_head);

	skel = SCX_OPS_OPEN(userland_ops, scx_userland);

	skel->rodata->num_possible_cpus = libbpf_num_possible_cpus();
	assert(skel->rodata->num_possible_cpus > 0);
	skel->rodata->usersched_pid = getpid();
	assert(skel->rodata->usersched_pid > 0);

	SCX_OPS_LOAD(skel, userland_ops, scx_userland, uei);

	enqueued_fd = bpf_map__fd(skel->maps.enqueued);
	dispatched_fd = bpf_map__fd(skel->maps.dispatched);
	assert(enqueued_fd > 0);
	assert(dispatched_fd > 0);

	SCX_BUG_ON(spawn_stats_thread(), "Failed to spawn stats thread");

	ops_link = SCX_OPS_ATTACH(skel, userland_ops, scx_userland);
}

static void sched_main_loop(void)
{
	while (!exit_req) {
		/*
		 * Perform the following work in the main user space scheduler
		 * loop:
		 *
		 * 1. Drain all tasks from the enqueued map, and enqueue them
		 *    to the vruntime sorted list.
		 *
		 * 2. Dispatch a batch of tasks from the vruntime sorted list
		 *    down to the kernel.
		 *
		 * 3. Yield the CPU back to the system. The BPF scheduler will
		 *    reschedule the user space scheduler once another task has
		 *    been enqueued to user space.
		 */
		drain_enqueued_map();
		dispatch_batch();
		sched_yield();
	}
}

int main(int argc, char **argv)
{
	__u64 ecode;

	pre_bootstrap(argc, argv);
restart:
	bootstrap(argv[0]);
	sched_main_loop();

	exit_req = 1;
	bpf_link__destroy(ops_link);
	pthread_join(stats_printer, NULL);
	ecode = UEI_REPORT(skel, uei);
	scx_userland__destroy(skel);

	if (UEI_ECODE_RESTART(ecode))
		goto restart;
	return 0;
}