// SPDX-License-Identifier: (GPL-2.0 OR BSD-3-Clause)
/*
 * Copyright (c) 2023 MediaTek Inc.
 * Author Mark-PK Tsai <mark-pk.tsai@mediatek.com>
 *
 * MTK Auto Nice Daemon
 */

#define pr_fmt(fmt) KBUILD_MODNAME ": " fmt

#include <linux/bitmap.h>
#include <linux/cpuset.h>
#include <linux/freezer.h>
#include <linux/kernel.h>
#include <linux/kthread.h>
#include <linux/module.h>
#include <linux/sched.h>
#include <linux/sched/deadline.h>
#include <linux/sched/rt.h>
#include <linux/sched/signal.h>
#include <linux/sched/task.h>
#include <linux/slab.h>
#include <linux/string_helpers.h>
#include <linux/tracepoint.h>
#include <linux/version.h>

#include <trace/events/sched.h>

#include <uapi/linux/sched/types.h>

#define MTK_AND_MAX_COMM 16
#define DEFAULT_SCAN_PERIOD_MS 5000
#define DEFAULT_MONITOR_NICE 10

#define IGNORE_KTHREAD 0
#define AND_FLAG_BIT_SCHED 0
#define AND_FLAG_BIT_AFF 1

struct mtk_and_conf {
	char *comm;
	char *tname;
	int policy;
	int sched_priority;
	cpumask_var_t cpumask;
	unsigned long flags;
	bool enable;
};

struct mtk_and_data {
	struct mtk_and_conf *conf;
	struct task_struct *scan_task;
};

static int nr_elem;
static char *conf[MTK_AND_MAX_COMM];
static int poll_delay = DEFAULT_SCAN_PERIOD_MS;
static bool force_nice;
static bool th_eff = true;
static bool tp_enable;
module_param_array(conf, charp, &nr_elem, 0444);
module_param(poll_delay, int, 0644);
module_param(force_nice, bool, 0644);
module_param(th_eff, bool, 0644);
module_param(tp_enable, bool, 0444); /* experimental */

struct tracepoint *task_rename;
struct mtk_and_data *and_data;

static bool cmdline_match(struct task_struct *tsk, char *target)
{
	char *cmdline;
	int ret;

	/* might sleep */
	cmdline = kstrdup_quotable_cmdline(tsk, GFP_KERNEL);
	if (!cmdline)
		return false;

	ret = strncmp(cmdline, target, strlen(target));
	kfree(cmdline);

	return !ret;
}

static bool is_rt_policy(int policy)
{
	return policy == SCHED_FIFO || policy == SCHED_RR;
}

static bool is_fair_policy(int policy)
{
	return policy == SCHED_NORMAL || policy == SCHED_BATCH;
}

/* caller should hold rcu read lock or task ref */
static int valid_and_set_sched(struct task_struct *task, struct mtk_and_conf *target)
{
	int priority = target->sched_priority;
	unsigned int policy = target->policy;

	if (task->policy != policy || is_rt_policy(policy)) {
		struct sched_param params;

		params.sched_priority = is_rt_policy(policy) ? priority : 0;

		sched_setscheduler_nocheck(task,
					   policy,
					   &params);
	}

	if (is_fair_policy(policy) && task_nice(task) != priority)
		set_user_nice(task, priority);

	pr_debug("set task %s with policy %d prio %d\n",
		task->comm, target->policy, target->sched_priority);

	return 0;
}

static int valid_and_set_cpumask(struct task_struct *task, struct mtk_and_conf *target)
{
	int ret;

	ret = set_cpus_allowed_ptr(task, target->cpumask);
	if (ret)
		pr_err("failed to set cpu aff, ret = %d\n", ret);

	return ret;
}

static int do_set_sched_or_aff(struct task_struct *task, struct mtk_and_conf *target)
{
	int ret = 0;

	if (test_bit(AND_FLAG_BIT_SCHED, &target->flags)) {
		pr_debug("set sched\n");
		ret = valid_and_set_sched(task, target);
	}

	if (test_bit(AND_FLAG_BIT_AFF, &target->flags)) {
		pr_debug("set aff\n");
		ret |= valid_and_set_cpumask(task, target);
	}

	return ret;
}

static int filter_and_renice_thread(struct task_struct *task, struct mtk_and_conf *target)
{
	struct pid *tgid;
	struct task_struct *p;

	tgid = task_tgid(task);
	rcu_read_lock();

	do_each_pid_thread(tgid, PIDTYPE_TGID, p) {
		pr_debug("%s, pname = %s, tname = %s\n", __func__, p->comm, target->tname);
		if (strncmp(p->comm, target->tname, strlen(target->tname)))
			continue;

		do_set_sched_or_aff(p, target);
	} while_each_pid_thread(tgid, PIDTYPE_TGID, p);

	rcu_read_unlock();
	return 0;
}

static void filter_and_renice(struct mtk_and_conf *data, struct task_struct *task)
{
	int i = 0;
	struct pid *tgid;
	struct task_struct *p;
	char *comm = task->comm;
	int ret;

	for (; i < nr_elem; i++) {
		if (!strncmp(comm, data[i].comm, strlen(data[i].comm)))
			break;

		if (cmdline_match(task, data[i].comm))
			break;
	}

	if (i >= nr_elem)
		return;

	pr_info("%s %d pname = %s\n", __func__, __LINE__, data[i].comm);
	if (data[i].tname) {
		pr_info("%s %d , tname = %s\n", __func__, __LINE__, data[i].tname);
		filter_and_renice_thread(task, &data[i]);
		return;
	}

	ret = do_set_sched_or_aff(task, &data[i]);
	if (ret) {
		pr_err("failed to set prio\n");
		return;
	}

	if (!th_eff)
		return;

	tgid = task_tgid(task);
	rcu_read_lock();

	do_each_pid_thread(tgid, PIDTYPE_TGID, p) {
		ret = do_set_sched_or_aff(p, &data[i]);
	} while_each_pid_thread(tgid, PIDTYPE_TGID, p);

	rcu_read_unlock();
}

static int mtk_and_scan_func(void *config)
{
	struct task_struct *p, **t;
	size_t arr_size = PAGE_SIZE;

	pr_info("monitor thread start\n");
	t = kmalloc(arr_size, GFP_KERNEL);
	if (!t) {
		pr_err("failed to alloc thread pool\n");
		return -ENOMEM;
	}

	set_user_nice(current, DEFAULT_MONITOR_NICE);
	set_freezable();
	while (!kthread_should_stop()) {
		signed long timeout = msecs_to_jiffies(poll_delay);
		ktime_t start;
		int todo, i, need_realloc = 0;

		todo = 0;
		start = ktime_get();

		rcu_read_lock();

		for_each_process(p) {
#if IGNORE_KTHREAD
			if (p->flags & PF_KTHREAD)
				continue;
#endif

			if (dl_task(p))
				continue;

			if (todo * sizeof(p) >= arr_size) {
				pr_info("need realloc\n");
				need_realloc = 1;
				break;
			}

			get_task_struct(p);
			t[todo++] = p;
		}

		rcu_read_unlock();

		for (i = 0; i < todo; i++) {
			filter_and_renice((struct mtk_and_conf *)config, t[i]);
			put_task_struct(t[i]);
		}

		pr_debug("scan took %lld us\n",  ktime_us_delta(ktime_get(), start));

		if (need_realloc) {
			struct task_struct **new;
			size_t new_size = arr_size + PAGE_SIZE;

			new = krealloc(t, new_size, GFP_KERNEL);
			if (!new) {
				pr_warn("failed to realloc size %zx\n", arr_size);
			} else {
				t = new;
				arr_size = new_size;
				need_realloc = 0;
				pr_info("increase arr size to %zx\n", arr_size);
			}
		}


		/* wait before the next scan */
#if LINUX_VERSION_CODE > KERNEL_VERSION(6, 6, 0)
		while (timeout && !kthread_should_stop()) {
			__set_current_state(TASK_INTERRUPTIBLE | TASK_FREEZABLE);
			timeout = schedule_timeout(timeout);
		}
#else
		while (timeout && !kthread_should_stop())
			timeout = freezable_schedule_timeout_interruptible(timeout);
#endif

		try_to_freeze();
	}

	kfree(t);
	pr_info("monitor thread ended\n");

	return 0;
}

static void probe_new_app(void *ignore, struct task_struct *p, pid_t old_pid,
			  struct linux_binprm *bprm)
{
	if (!p) {
		pr_err("wtf\n");
		return;
	}

	if (p->flags & PF_KTHREAD)
		return;

	if (dl_task(p) || rt_task(p))
		return;

	if (task_pid_nr(p) != task_tgid_nr(p))
		return;

	pr_info("probe new app comm = %s, filename = %s\n", p->comm, bprm->filename);

}

static void probe_app_exit(void *data, struct task_struct *p)
{
	if (!p) {
		pr_err("app died but not task struct, wtf?\n");
		return;
	}

	if (p->flags & PF_KTHREAD)
		return;

	if (dl_task(p) || rt_task(p))
		return;

	if (task_pid_nr(p) != task_tgid_nr(p))
		return;

	pr_info("app died comm = %s\n", p->comm);
}

static void register_tracepoints(struct tracepoint *tp, void *ignore)
{
	/*check_trace_callback_type_console(probe_console);*/
	if (!strcmp(tp->name, "sched_process_exec"))
		WARN_ON(tracepoint_probe_register(tp, probe_new_app, NULL));

	if (!strcmp(tp->name, "sched_process_free"))
		WARN_ON(tracepoint_probe_register(tp, probe_app_exit, NULL));
}

static void unregister_tracepoints(struct tracepoint *tp, void *ignore)
{
	if (!strcmp(tp->name, "sched_process_exec"))
		tracepoint_probe_unregister(tp, probe_new_app, NULL);

	if (!strcmp(tp->name, "sched_process_free"))
		tracepoint_probe_unregister(tp, probe_app_exit, NULL);
}

static int register_tp(void)
{
	for_each_kernel_tracepoint(register_tracepoints, NULL);
	return 0;
}

static void show_usage(void)
{
	pr_warn("Usage:\n"
		"Use conf={name}={OPT}:{value}\n"
		"when OPT = S, {value} definition is {sched policy}{prio}\n"
		"e.g. conf=logd=S:N6,audioserver=S:R10\n"
		"     This set logd=(sched_policy=normal, nice=6) and audioserver=(sched_policy=RR, prio=10)\n\n"
		"when OPT = A, {value} definition is {cpumask}\n"
		"e.g. conf=logd=A:e,audioserver=A:1\n"
		"     This set logd=(cpu affinity = 0xE) and audioserver=(cpu affinity = 0x1)\n");
}

static int parse_sched_option(struct mtk_and_conf *config, char *arg)
{
	long prio;
	int policy;

	if (!arg)
		return -EINVAL;

	pr_err("%s, %d, line = %s\n", __func__, __LINE__, arg);

	/* policy */
	switch (*(arg++)) {
	case 'F':
		policy = SCHED_FIFO;
		break;
	case 'R':
		policy = SCHED_RR;
		break;
	case 'N':
		policy = SCHED_NORMAL;
		break;
	default:
		return -EINVAL;
	}

	if (!(*arg))
		return -EINVAL;

	/* prio */
	if (kstrtol(arg, 0, &prio)) {
		pr_err("failed to parse prio \"%s\"\n", arg);
		return -EINVAL;
	}

	config->policy = policy;
	config->sched_priority = (int)prio;
	config->flags |= BIT(AND_FLAG_BIT_SCHED);

	return 0;
}

static int parse_aff_option(struct mtk_and_conf *config, char *arg)
{
	if (!arg)
		return -EINVAL;

	if (!alloc_cpumask_var(&config->cpumask, GFP_KERNEL))
		return -ENOMEM;

	if (cpumask_parse(arg, config->cpumask)) {
		free_cpumask_var(config->cpumask);
		return -EINVAL;
	}

	pr_debug("%s, %d, line = %s, CPUs %*pbl",
		__func__, __LINE__, arg, cpumask_pr_args(config->cpumask));

	config->flags |= BIT(AND_FLAG_BIT_AFF);
	return 0;
}

static int __init and_init(struct mtk_and_data *data)
{
	int ret, i;
	struct mtk_and_conf *config;

	if (!nr_elem) {
		pr_warn("skip and init\n");
		return 0;
	}

	config = kcalloc(nr_elem, sizeof(struct mtk_and_conf), GFP_KERNEL);
	if (!config)
		return -ENOMEM;

	pr_debug("nr_elem = %d\n", nr_elem);
	for (i = 0; i < nr_elem; i++) {
		char *name, *opt, *tname;
		char *ptr, *tmp_ptr;

		ptr = conf[i];
		pr_crit("line %d: %s\n", i, ptr);
		name = strsep(&ptr, "=");
		if (!name || !ptr) {
			pr_err("failed to parse line \"%s\"\n", conf[i]);
			ret = -EINVAL;
			goto err;
		}

		opt = strsep(&ptr, ":");
		if (!opt || !ptr) {
			pr_err("failed to parse line \"%s\"\n", conf[i]);
			ret = -EINVAL;
			goto err;
		}

		// threadname%processname
		tmp_ptr = name;
		tname = strsep(&tmp_ptr, "%");
		if (tname) {
			pr_info("tname = %s, pname = %s\n", tname, tmp_ptr);
			name = tmp_ptr;
		} else
			pr_info("no tname found, keep name %s\n", name);

		switch (*opt) {
		case 'S':
			ret = parse_sched_option(&config[i], ptr);
			if (ret) {
				pr_err("failed to parse sched opt with line \"%s\"\n", conf[i]);
				goto err;
			}
			break;
		case 'A':
			ret = parse_aff_option(&config[i], ptr);
			if (ret) {
				pr_err("failed to parse cpu aff opt with line \"%s\"\n", conf[i]);
				goto err;
			}
			break;
		default:
			pr_err("not support opt \"%s\"\n", opt);
			ret = -EINVAL;
			goto err;
		}

		// if ok set name
		config[i].comm = name;
		config[i].tname = tname;
		pr_info("config[%d] (%s, %d, %d)\n", i,
			config[i].comm, config[i].policy, config[i].sched_priority);
	}

	and_data->conf = config;

	if (tp_enable) {
		ret = register_tp();
		if (ret)
			goto err;
	}

	and_data->scan_task = kthread_run(mtk_and_scan_func, and_data->conf, "mtk_and_scan");
	if (IS_ERR(and_data->scan_task)) {
		if (tp_enable) {
			for_each_kernel_tracepoint(unregister_tracepoints, NULL);
			tracepoint_synchronize_unregister();
		}

		ret = PTR_ERR(and_data->scan_task);
		pr_err("failed to create kthread\n");
		goto err;
	}

	return 0;
err:
	kfree(config);
	return ret;
}

static int __init mtk_and_init(void)
{
	int ret;

	if (!nr_elem) {
		show_usage();
		return -EINVAL;
	}

	and_data = kzalloc(sizeof(struct mtk_and_data), GFP_KERNEL);
	if (!and_data) {
		ret = -ENOMEM;
		goto err;
	}

	ret = and_init(and_data);
	if (ret)
		goto err;

	return 0;
err:
	kfree(and_data);
	return ret;
}

static void __exit mtk_and_exit(void)
{
	kthread_stop(and_data->scan_task);

	if (tp_enable) {
		for_each_kernel_tracepoint(unregister_tracepoints, NULL);
		tracepoint_synchronize_unregister();
	}

	kfree(and_data->conf);
	kfree(and_data);
}

module_init(mtk_and_init);
module_exit(mtk_and_exit);

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Mark-PK Tsai");
MODULE_DESCRIPTION("MTK Auto Nice Daemon Driver");
