/* SPDX-License-Identifier: GPL-2.0 */

#include <linux/kernel.h>
#include <linux/uaccess.h>
#include <linux/sched/deadline.h>
#include <asm/syscall.h>
#include <uapi/linux/sched/types.h>
#include <trace/events/sched.h>

/*
 * Dummy values if not available
 */
#ifndef __NR_sched_setscheduler
#define __NR_sched_setscheduler -__COUNTER__
#endif
#ifndef __NR_sched_setattr
#define __NR_sched_setattr -__COUNTER__
#endif

extern struct rv_monitor rv_deadline;
/* Initialised when registering the deadline container */
extern struct sched_class *rv_ext_sched_class;

/*
 * If both have dummy values, the syscalls are not supported and we don't even
 * need to register the handler.
 */
static inline bool should_skip_syscall_handle(void)
{
	return __NR_sched_setattr < 0 && __NR_sched_setscheduler < 0;
}

/*
 * is_supported_type - return true if @type is supported by the deadline monitors
 */
static inline bool is_supported_type(u8 type)
{
	return type == DL_TASK || type == DL_SERVER_FAIR || type == DL_SERVER_EXT;
}

/*
 * is_server_type - return true if @type is a supported server
 */
static inline bool is_server_type(u8 type)
{
	return is_supported_type(type) && type != DL_TASK;
}

/*
 * Use negative numbers for the server.
 * Currently only one fair server per CPU, may change in the future.
 */
#define fair_server_id(cpu) (-cpu)
#define ext_server_id(cpu) (-cpu - num_possible_cpus())
#define NO_SERVER_ID (-2 * num_possible_cpus())
/*
 * Get a unique id used for dl entities
 *
 * The cpu is not required for tasks as the pid is used there, if this function
 * is called on a dl_se that for sure corresponds to a task, DL_TASK can be
 * used in place of cpu.
 * We need the cpu for servers as it is provided in the tracepoint and we
 * cannot easily retrieve it from the dl_se (requires the struct rq definition).
 */
static inline int get_entity_id(struct sched_dl_entity *dl_se, int cpu, u8 type)
{
	if (dl_server(dl_se) && type != DL_TASK) {
		if (type == DL_SERVER_FAIR)
			return fair_server_id(cpu);
		if (type == DL_SERVER_EXT)
			return ext_server_id(cpu);
		return NO_SERVER_ID;
	}
	return dl_task_of(dl_se)->pid;
}

static inline bool task_is_scx_enabled(struct task_struct *tsk)
{
	return IS_ENABLED(CONFIG_SCHED_CLASS_EXT) &&
	       tsk->sched_class == rv_ext_sched_class;
}

/* Expand id and target as arguments for da functions */
#define EXPAND_ID(dl_se, cpu, type) get_entity_id(dl_se, cpu, type), dl_se
#define EXPAND_ID_TASK(tsk) get_entity_id(&tsk->dl, task_cpu(tsk), DL_TASK), &tsk->dl

static inline u8 get_server_type(struct task_struct *tsk)
{
	if (tsk->policy == SCHED_NORMAL || tsk->policy == SCHED_EXT ||
	    tsk->policy == SCHED_BATCH || tsk->policy == SCHED_IDLE)
		return task_is_scx_enabled(tsk) ? DL_SERVER_EXT : DL_SERVER_FAIR;
	return DL_OTHER;
}

static inline int extract_params(struct pt_regs *regs, long id, pid_t *pid_out)
{
	size_t size = offsetofend(struct sched_attr, sched_flags);
	struct sched_attr __user *uattr, attr;
	int new_policy = -1, ret;
	unsigned long args[6];

	switch (id) {
	case __NR_sched_setscheduler:
		syscall_get_arguments(current, regs, args);
		*pid_out = args[0];
		new_policy = args[1];
		break;
	case __NR_sched_setattr:
		syscall_get_arguments(current, regs, args);
		*pid_out = args[0];
		uattr = (struct sched_attr __user *)args[1];
		/*
		 * Just copy up to sched_flags, we are not interested after that
		 */
		ret = copy_struct_from_user(&attr, size, uattr, size);
		if (ret)
			return ret;
		if (attr.sched_flags & SCHED_FLAG_KEEP_POLICY)
			return -EINVAL;
		new_policy = attr.sched_policy;
		break;
	default:
		return -EINVAL;
	}

	return new_policy & ~SCHED_RESET_ON_FORK;
}

/* Helper functions requiring DA/HA utilities */
#ifdef RV_MON_TYPE

/*
 * get_fair_server - get the fair server associated to a task
 *
 * If the task is a boosted task, the server is available in the task_struct,
 * otherwise grab the dl entity saved for the CPU where the task is enqueued.
 * This function assumes the task is enqueued somewhere.
 */
static inline struct sched_dl_entity *get_server(struct task_struct *tsk, u8 type)
{
	if (tsk->dl_server && get_server_type(tsk) == type)
		return tsk->dl_server;
	if (type == DL_SERVER_FAIR)
		return da_get_target_by_id(fair_server_id(task_cpu(tsk)));
	if (type == DL_SERVER_EXT)
		return da_get_target_by_id(ext_server_id(task_cpu(tsk)));
	return NULL;
}

/*
 * Initialise monitors for all tasks and pre-allocate the storage for servers.
 * This is necessary since we don't have access to the servers here and
 * allocation can cause deadlocks from their tracepoints. We can only fill
 * pre-initialised storage from there.
 */
static inline int init_storage(bool skip_tasks)
{
	struct task_struct *g, *p;
	int cpu;

	for_each_possible_cpu(cpu) {
		if (!da_create_empty_storage(fair_server_id(cpu)))
			goto fail;
		if (IS_ENABLED(CONFIG_SCHED_CLASS_EXT) &&
		    !da_create_empty_storage(ext_server_id(cpu)))
			goto fail;
	}

	if (skip_tasks)
		return 0;

	read_lock(&tasklist_lock);
	for_each_process_thread(g, p) {
		if (p->policy == SCHED_DEADLINE) {
			if (!da_create_storage(EXPAND_ID_TASK(p), NULL)) {
				read_unlock(&tasklist_lock);
				goto fail;
			}
		}
	}
	read_unlock(&tasklist_lock);
	return 0;

fail:
	da_monitor_destroy();
	return -ENOMEM;
}

static void __maybe_unused handle_newtask(void *data, struct task_struct *task, u64 flags)
{
	/* Might be superfluous as tasks are not started with this policy.. */
	if (task->policy == SCHED_DEADLINE)
		da_create_storage(EXPAND_ID_TASK(task), NULL);
}

static void __maybe_unused handle_exit(void *data, struct task_struct *p, bool group_dead)
{
	if (p->policy == SCHED_DEADLINE)
		da_destroy_storage(get_entity_id(&p->dl, DL_TASK, DL_TASK));
}

#endif