/*
 * Oracle Linux DTrace.
 * Copyright (c) 2023, 2025, Oracle and/or its affiliates. All rights reserved.
 * Licensed under the Universal Permissive License v 1.0 as shown at
 * http://oss.oracle.com/licenses/upl.
 *
 * The 'sched' SDT provider for DTrace-specific probes.
 */
#include <assert.h>
#include <errno.h>

#include "dt_dctx.h"
#include "dt_cg.h"
#include "dt_provider_sdt.h"
#include "dt_probe.h"

static const char		prvname[] = "sched";
static const char		modname[] = "vmlinux";

static probe_dep_t	probes[] = {
	{ "dequeue",
	  DTRACE_PROBESPEC_NAME,	"fbt::dequeue_task_*:entry" },
	{ "enqueue",
	  DTRACE_PROBESPEC_NAME,	"fbt::enqueue_task_*:entry" },
	{ "off-cpu",
	  DTRACE_PROBESPEC_NAME,	"rawtp:sched::sched_switch" },
	{ "on-cpu",
	  DTRACE_PROBESPEC_NAME,	"fbt:vmlinux:finish_task_switch:return" },
	{ "on-cpu",
	  DTRACE_PROBESPEC_NAME,	"rawfbt:vmlinux:finish_task_switch.*:return" },
	{ "surrender",
	  DTRACE_PROBESPEC_NAME,	"fbt::do_sched_yield:entry" },
	{ "tick",
	  DTRACE_PROBESPEC_NAME,	"fbt::scheduler_tick:entry" },
	{ "wakeup",
	  DTRACE_PROBESPEC_NAME,	"rawtp:sched::sched_wakeup" },
	{ NULL, }
};

static probe_arg_t probe_args[] = {
#if 0
	{ "change-pri", 0, { 0, 0, "struct task_struct *", "lwpsinfo_t *" } },
	{ "change-pri", 1, { 0, 0, "struct task_struct *", "psinfo_t *" } },
	{ "change-pri", 2, { 1, 0, "int", } },
#endif
	{ "dequeue", 0, { 0, 0, "struct task_struct *", "lwpsinfo_t *" } },
	{ "dequeue", 1, { 0, 0, "struct task_struct *", "psinfo_t *" } },
	{ "dequeue", 2, { 1, 0, "cpuinfo_t *", } },
	{ "enqueue", 0, { 0, 0, "struct task_struct *", "lwpsinfo_t *" } },
	{ "enqueue", 1, { 0, 0, "struct task_struct *", "psinfo_t *" } },
	{ "enqueue", 2, { 1, 0, "cpuinfo_t *", } },
	{ "enqueue", 3, { 2, 0, "int", } },
	{ "off-cpu", 0, { 0, 0, "struct task_struct *", "lwpsinfo_t *" } },
	{ "off-cpu", 1, { 0, 0, "struct task_struct *", "psinfo_t *" } },
	{ "on-cpu", },
#if 0
	{ "preempt", },
	{ "remain-cpu", },
	{ "sleep", },
#endif
	{ "surrender", 0, { 0, 0, "struct task_struct *", "lwpsinfo_t *" } },
	{ "surrender", 1, { 0, 0, "struct task_struct *", "psinfo_t *" } },
	{ "tick", 0, { 0, 0, "struct task_struct *", "lwpsinfo_t *" } },
	{ "tick", 1, { 0, 0, "struct task_struct *", "psinfo_t *" } },
	{ "wakeup", 0, { 0, 0, "struct task_struct *", "lwpsinfo_t *" } },
	{ "wakeup", 1, { 0, 0, "struct task_struct *", "psinfo_t *" } },
	{ NULL, }
};

static const dtrace_pattr_t	pattr = {
{ DTRACE_STABILITY_EVOLVING, DTRACE_STABILITY_EVOLVING, DTRACE_CLASS_ISA },
{ DTRACE_STABILITY_PRIVATE, DTRACE_STABILITY_PRIVATE, DTRACE_CLASS_UNKNOWN },
{ DTRACE_STABILITY_PRIVATE, DTRACE_STABILITY_PRIVATE, DTRACE_CLASS_UNKNOWN },
{ DTRACE_STABILITY_EVOLVING, DTRACE_STABILITY_EVOLVING, DTRACE_CLASS_ISA },
{ DTRACE_STABILITY_EVOLVING, DTRACE_STABILITY_EVOLVING, DTRACE_CLASS_ISA },
};

/*
 * Provide all the "sched" SDT probes.
 */
static int populate(dtrace_hdl_t *dtp)
{
	return dt_sdt_populate(dtp, prvname, modname, &dt_sched, &pattr,
			       probe_args, probes);
}

/*
 * Get a pointer to the cpuinfo_t structure for the CPU associated
 * with the runqueue that is in arg0.
 *
 * Clobbers %r1 through %r5
 * Stores pointer to cpuinfo_t struct in %r0
 */
static void get_cpuinfo(dtrace_hdl_t *dtp, dt_irlist_t *dlp, uint_t exitlbl)
{
	dt_ident_t	*idp = dt_dlib_get_map(dtp, "cpuinfo");

	assert(idp != NULL);

	/* Put the runqueue pointer from mst->arg0 into %r3. */
	emit(dlp, BPF_LOAD(BPF_DW, BPF_REG_3, BPF_REG_7, DMST_ARG(0)));

	/* Turn it into a pointer to its cpu member. */
	emit(dlp, BPF_ALU64_IMM(BPF_ADD, BPF_REG_3, dt_cg_ctf_offsetof("struct rq", "cpu", NULL, NULL, 1)));

	/* Call bpf_probe_read_kernel(%fp + DT_TRAMP_SP_SLOT[0], sizeof(int), %r3) */
	emit(dlp, BPF_MOV_IMM(BPF_REG_2, (int) sizeof(int)));
	emit(dlp, BPF_MOV_REG(BPF_REG_1, BPF_REG_FP));
	emit(dlp, BPF_ALU64_IMM(BPF_ADD, BPF_REG_1, DT_TRAMP_SP_SLOT(0)));
	emit(dlp, BPF_CALL_HELPER(BPF_FUNC_probe_read_kernel));
	emit(dlp, BPF_BRANCH_IMM(BPF_JNE, BPF_REG_0, 0, exitlbl));

	/* Now look up the corresponding cpuinfo_t. */
	dt_cg_xsetx(dlp, idp, DT_LBL_NONE, BPF_REG_1, idp->di_id);
	emit(dlp, BPF_MOV_REG(BPF_REG_2, BPF_REG_FP));
	emit(dlp, BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, DT_TRAMP_SP_SLOT(0)));
	emit(dlp, BPF_CALL_HELPER(BPF_FUNC_map_lookup_elem));
	emit(dlp, BPF_BRANCH_IMM(BPF_JEQ, BPF_REG_0, 0, exitlbl));
}

/*
 * Generate a BPF trampoline for a SDT probe.
 *
 * The trampoline function is called when a SDT probe triggers, and it must
 * satisfy the following prototype:
 *
 *	int dt_sched(void *data)
 *
 * The trampoline will populate a dt_dctx_t struct and then call the function
 * that implements the compiled D clause.  It returns the value that it gets
 * back from that function.
 */
static int trampoline(dt_pcb_t *pcb, uint_t exitlbl)
{
	dtrace_hdl_t	*dtp = pcb->pcb_hdl;
	dt_irlist_t	*dlp = &pcb->pcb_ir;
	dt_probe_t	*prp = pcb->pcb_probe;

	if (strcmp(prp->desc->prb, "dequeue") == 0) {
		/*
		 * Get the runqueue from arg0 and place its cpuinfo_t* into %r0.
		 */
		get_cpuinfo(dtp, dlp, exitlbl);

		/*
		 * Copy arg1 into arg0.
		 */
		emit(dlp, BPF_LOAD(BPF_DW, BPF_REG_3, BPF_REG_7, DMST_ARG(1)));
		emit(dlp, BPF_STORE(BPF_DW, BPF_REG_7, DMST_ARG(0), BPF_REG_3));

		/* Store the cpuinfo_t* in %r0 into arg1. */
		emit(dlp, BPF_STORE(BPF_DW, BPF_REG_7, DMST_ARG(1), BPF_REG_0));
	} else if (strcmp(prp->desc->prb, "enqueue") == 0) {
		/*
		 * Get the runqueue from arg0 and place its cpuinfo_t* into %r0.
		 */
		get_cpuinfo(dtp, dlp, exitlbl);

		/*
		 * Copy arg1 into arg0.
		 */
		emit(dlp, BPF_LOAD(BPF_DW, BPF_REG_3, BPF_REG_7, DMST_ARG(1)));
		emit(dlp, BPF_STORE(BPF_DW, BPF_REG_7, DMST_ARG(0), BPF_REG_3));

		/* Store the cpuinfo_t* in %r0 into arg1. */
		emit(dlp, BPF_STORE(BPF_DW, BPF_REG_7, DMST_ARG(1), BPF_REG_0));

/*
 * This is ugly but necessary...  enqueue_task() takes a flags argument and the
 * ENQUEUE_HEAD flag is used to indicate that the task is to be placed at the
 * head of the queue.  We need to be able to pass this special case as arg2
 * in the enqueue probe.
 *
 * The flag values are found in kernel/sched/sched.h which is not exposed
 * outside the kernel source tree.
 */
#define ENQUEUE_HEAD	0x10
		emit(dlp, BPF_LOAD(BPF_DW, BPF_REG_0, BPF_REG_7, DMST_ARG(2)));
		emit(dlp, BPF_ALU64_IMM(BPF_AND, BPF_REG_0, ENQUEUE_HEAD));
		emit(dlp, BPF_STORE(BPF_DW, BPF_REG_7, DMST_ARG(2), BPF_REG_0));
	} else if (strcmp(prp->desc->prb, "off-cpu") == 0) {
		emit(dlp, BPF_LOAD(BPF_DW, BPF_REG_0, BPF_REG_7, DMST_ARG(2)));
		emit(dlp, BPF_STORE(BPF_DW, BPF_REG_7, DMST_ARG(0), BPF_REG_0));
	} else if (strcmp(prp->desc->prb, "surrender") == 0 ||
		   strcmp(prp->desc->prb, "tick") == 0) {
		emit(dlp, BPF_CALL_HELPER(BPF_FUNC_get_current_task));
		emit(dlp, BPF_STORE(BPF_DW, BPF_REG_7, DMST_ARG(0), BPF_REG_0));
	}

	return 0;
}

dt_provimpl_t	dt_sched = {
	.name		= prvname,
	.prog_type	= BPF_PROG_TYPE_UNSPEC,
	.populate	= &populate,
	.enable		= &dt_sdt_enable,
	.load_prog	= &dt_bpf_prog_load,
	.trampoline	= &trampoline,
	.probe_info	= &dt_sdt_probe_info,
	.destroy	= &dt_sdt_destroy,
};
