// SPDX-License-Identifier: GPL-2.0
/*
 * scx_k2 - minimal cid-form scheduler pair for sub-scheduler attach
 * testing (single binary, two modes).
 *
 * Root mode (default): private FIFO DSQ; when a sub-scheduler attaches,
 * grants grant_caps (default SCX_CAP_ENQ_IMMED|SCX_CAP_PERF, full cid
 * mask) over the cgroup given with -g. Never writes cpuperf itself.
 * -R <ticks> issues scx_bpf_sub_revoke() of SCX_CAP_PERF on tick N
 * (0 = off).
 *
 * Sub mode (-s): ops.tick() writes perf_target via
 * scx_bpf_cidperf_set(this_cid).
 *
 * cid-form notes: must use SCX_OPS_CID_DEFINE/OPEN; every program must
 * touch the arena (k2_arena_touch) or bpf_prog_arena() comes back NULL
 * and programs get rejected; sub_cgroup_id must be written to both
 * rodata and struct_ops; the minimal cap set is ENQ_IMMED; a private
 * DSQ is required.
 */
#include <scx/common.bpf.h>

char _license[] SEC("license") = "GPL";

const volatile u64 grant_cgid;	/* root: cgid to grant/revoke caps over */
const volatile u64 grant_caps;	/* root: cap bits to grant */
const volatile u64 revoke_at;	/* root: revoke PERF on tick N (0 = off) */
const volatile u32 restore_on_detach; /* root: parent writes neutral values on sub_detach */
const volatile u64 sub_cgroup_id; /* sub mode: dual-written with ops->sub_cgroup_id */
const volatile u32 is_sub;	/* 1: sub mode */
const volatile u32 perf_target;	/* sub: target written every tick */

UEI_DEFINE(uei);

#define SHARED_DSQ 0

enum {
	K2_GRANT_RET = 0,	/* sub_grant return value (errno stored as 8192-errno) */
	K2_REVOKE_DONE = 1,	/* number of revokes performed */
	K2_SETCNT = 2,		/* number of sub tick writes */
	K2_SET_ERR = 3,		/* last cidperf_set return value (errno as 8192-errno) */
	K2_ROOTTICKS = 4,	/* root ticks on cpu0 (ownership discriminator) */
	K2_DSQDEPTH = 7,	/* SHARED_DSQ depth sampled in root tick (starvation probe) */
	K2_DISP0 = 8,		/* ROOT's cid0 dispatch invocations */
	K2_DISP0_MOVE = 9,	/* ROOT's cid0 successful move_to_local */
	K2_SUBDISP0 = 10,	/* SUB's cid0 dispatch invocations (core routing signal) */
	K2_SUBDISP0_MOVE = 11,	/* SUB's cid0 successful move_to_local */
	K2_CNT_NR = 12,
};

struct {
	__uint(type, BPF_MAP_TYPE_ARRAY);
	__uint(max_entries, K2_CNT_NR);
	__type(key, u32);
	__type(value, u64);
} counters SEC(".maps");

/* cid-form requires an arena map (verified by the kernel) */
struct {
	__uint(type, BPF_MAP_TYPE_ARENA);
	__uint(map_flags, BPF_F_MMAPABLE);
	__uint(max_entries, 1 << 16);	/* pages */
} arena SEC(".maps");

struct scx_cmask __arena cap_mask;
u64 __arena k2_ticks;

static void bump(u32 idx)
{
	u64 *cnt = bpf_map_lookup_elem(&counters, &idx);

	if (cnt)
		(*cnt)++;
}

static void setcnt(u32 idx, s64 v)
{
	u64 val = v < 0 ? 8192 + (-v) : v;

	bpf_map_update_elem(&counters, &idx, &val, BPF_ANY);
}

/* every scx program must touch the arena, or bpf_prog_arena() == NULL -> EINVAL */
u64 __arena k2_arena_sink;
static void arena_touch(void)
{
	k2_arena_sink += cap_mask.alloc_words;
}

s32 BPF_STRUCT_OPS(k2_select_cid, struct task_struct *p,
		   s32 prev_cid, u64 wake_flags)
{
	arena_touch();
	return prev_cid;
}

void BPF_STRUCT_OPS(k2_enqueue, struct task_struct *p, u64 enq_flags)
{
	arena_touch();
	scx_bpf_dsq_insert(p, SHARED_DSQ, SCX_SLICE_DFL, enq_flags);
}

void BPF_STRUCT_OPS(k2_dispatch, s32 cid, struct task_struct *prev)
{
	arena_touch();

	if (is_sub) {
		/* the sub's own dispatch turn: consume our SHARED DSQ */
		if (cid == 0)
			bump(K2_SUBDISP0);
		if (scx_bpf_dsq_move_to_local(SHARED_DSQ, 0))
			bump(K2_SUBDISP0_MOVE);
		return;
	}

	/*
	 * Root's dispatch turn: consume our own SHARED DSQ. (A/B tested:
	 * handing the turn down to the sub first via
	 * scx_bpf_sub_dispatch() does not rescue the claimed task.)
	 */

	if (cid == 0) {
		bump(K2_DISP0);
		if (scx_bpf_dsq_move_to_local(SHARED_DSQ, 0))
			bump(K2_DISP0_MOVE);
	} else {
		scx_bpf_dsq_move_to_local(SHARED_DSQ, 0);
	}
}

/* root: grant caps when a sub attaches (full cid mask; test harness
 * deliberately over-grants) */
s32 BPF_STRUCT_OPS(k2_sub_attach, struct scx_sub_attach_args *args)
{
	u32 nr;
	s32 ret;

	arena_touch();
	if (!grant_cgid || !grant_caps)
		return 0;

	nr = scx_bpf_nr_cids();
	cmask_init(&cap_mask, 0, nr);
	cmask_set_range(&cap_mask, 0, nr);
	ret = scx_bpf_sub_grant(grant_cgid, grant_caps, &cap_mask, NULL);
	setcnt(K2_GRANT_RET, ret);
	return 0;
}

/* parent-side restore per the last-writer-wins ownership model: once
 * the sub is inert (caps cleared), the parent writes back neutral values */
void BPF_STRUCT_OPS(k2_sub_detach, struct scx_sub_detach_args *args)
{
	u32 i, nr;

	arena_touch();
	if (!restore_on_detach)
		return;

	nr = scx_bpf_nr_cids();
	bpf_for(i, 0, nr) {
		if (restore_on_detach == 2 || i == 0)
			scx_bpf_cidperf_set(i, SCX_CPUPERF_ONE);
	}
}

/* root: revoke PERF on tick revoke_at; sub: write target every tick */
void BPF_STRUCT_OPS(k2_tick, struct task_struct *p)
{
	s32 cid;

	if (is_sub) {
		cid = scx_bpf_this_cid();
		if (cid >= 0)
			setcnt(K2_SET_ERR,
			       scx_bpf_cidperf_set(cid, perf_target));
		bump(K2_SETCNT);
		return;
	}

	/* count only cpu0 root ticks: discriminates who owns the cpu0
	 * spinner (root ticks mean the spinner runs as a root task) */
	if (bpf_get_smp_processor_id() == 0)
		bump(K2_ROOTTICKS);
	setcnt(K2_DSQDEPTH, scx_bpf_dsq_nr_queued(SHARED_DSQ));
	arena_touch();
	k2_ticks++;
	if (revoke_at && k2_ticks == revoke_at) {
		u32 nr = scx_bpf_nr_cids();

		cmask_init(&cap_mask, 0, nr);
		cmask_set_range(&cap_mask, 0, nr);
		scx_bpf_sub_revoke(grant_cgid, SCX_CAP_PERF, &cap_mask);
		bump(K2_REVOKE_DONE);
	}
}

void BPF_STRUCT_OPS(k2_cid_online, s32 cpu)
{
	arena_touch();
}

void BPF_STRUCT_OPS(k2_cid_offline, s32 cpu)
{
	arena_touch();
}

s32 BPF_STRUCT_OPS(k2_init_task, struct task_struct *p,
		   struct scx_init_task_args *args)
{
	arena_touch();
	return 0;
}

s32 BPF_STRUCT_OPS_SLEEPABLE(k2_init)
{
	asm volatile("" :: "r"(&arena));	/* make the verifier claim the arena map */
	return scx_bpf_create_dsq(SHARED_DSQ, -1);
}

void BPF_STRUCT_OPS(k2_exit, struct scx_exit_info *ei)
{
	UEI_RECORD(uei, ei);
}

SCX_OPS_CID_DEFINE(k2_ops,
		   .select_cid		= (void *)k2_select_cid,
		   .enqueue		= (void *)k2_enqueue,
		   .dispatch		= (void *)k2_dispatch,
		   .sub_attach		= (void *)k2_sub_attach,
		   .sub_detach		= (void *)k2_sub_detach,
		   .tick		= (void *)k2_tick,
		   .cid_online		= (void *)k2_cid_online,
		   .cid_offline		= (void *)k2_cid_offline,
		   .init_task		= (void *)k2_init_task,
		   .init		= (void *)k2_init,
		   .exit		= (void *)k2_exit,
		   .name		= "k2");