// SPDX-License-Identifier: GPL-2.0
/* scx_k2 userspace. Usage:
 *   root mode: scx_k2 [-g <sub-cgroup-dir>] [-c <caps>] [-R <ticks>] [-p <perf>]
 *   sub  mode: scx_k2 -s <my-cgroup-dir> [-p <perf>]
 * Prints all counters every 2s; SIGTERM exits (triggers sub disable =>
 * sweep path).
 */
#include <stdio.h>
#include <unistd.h>
#include <signal.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/stat.h>
#include <bpf/bpf.h>
#include <scx/common.h>

/* Mirror of kernel/sched/ext/types.h's arena cap mask: the arena global
 * cap_mask is mirrored to userspace via the skeleton (must be defined
 * before the skel include). The bits[] flexible array is omitted and a
 * __pad field pads the struct to the 16 bytes BTF reports (FAM alignment
 * padding) - userspace never accesses the object, and a FAM would
 * trigger -Wgnu-variable-sized-type-not-at-end. */
struct scx_cmask {
	u32 base;
	u32 nr_cids;
	u32 alloc_words;
	u32 __pad;
};

#include "scx_k2.bpf.skel.h"

static bool verbose;
static volatile int exit_req;

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 sig_handler(int sig)
{
	exit_req = 1;
}

static const char *cnt_name(u32 i)
{
	static const char *names[] = { "grant_ret", "revoke_done", "setcnt", "set_err", "rootticks", "detach", "detach_ok", "dsqdepth", "disp0", "moved0", "subdisp0", "submoved0" };

	return names[i];
}

static u64 read_cnt(struct scx_k2 *skel, u32 idx)
{
	u64 v = 0;

	bpf_map_lookup_elem(bpf_map__fd(skel->maps.counters), &idx, &v);
	return v;
}

int main(int argc, char **argv)
{
	struct scx_k2 *skel;
	struct bpf_link *link;
	__s32 opt;
	u64 sub_cgid = 0, grant_cgid = 0, grant_caps = 0, revoke_at = 0;
	u32 restore_on_detach = 0;
	u32 perf_target = 1;

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

	skel = SCX_OPS_CID_OPEN(k2_ops, scx_k2);
	while ((opt = getopt(argc, argv, "s:g:c:R:e:p:v")) != -1) {
		struct stat st;

		switch (opt) {
		case 's':	/* sub mode, arg = own cgroup */
			if (stat(optarg, &st)) { perror(optarg); return 1; }
			sub_cgid = st.st_ino;
			break;
		case 'g':	/* root mode, arg = sub's cgroup */
			if (stat(optarg, &st)) { perror(optarg); return 1; }
			grant_cgid = st.st_ino;
			break;
		case 'c':
			grant_caps = strtoull(optarg, NULL, 0);
			break;
		case 'R':
			revoke_at = strtoull(optarg, NULL, 0);
			break;
		case 'e':	/* on sub_detach parent restores neutral values (1=cpu0 only, 2=all cids) */
			restore_on_detach = (u32)strtoul(optarg, NULL, 0);
			break;
		case 'p':
			perf_target = (u32)strtoul(optarg, NULL, 0);
			break;
		case 'v':
			verbose = true;
			break;
		default:
			fprintf(stderr, "scx_k2: -s <cg>|-g <cg> [-c caps] [-R ticks] [-e 1|2] [-p perf]\n");
			return opt != 'h';
		}
	}

	if (sub_cgid) {
		skel->struct_ops.k2_ops->sub_cgroup_id = sub_cgid;
		skel->rodata->sub_cgroup_id = sub_cgid;
		skel->rodata->is_sub = 1;
	}
	skel->rodata->grant_cgid = grant_cgid;
	/* no SCX_CAP_* enum in the userspace headers (kernel internal.h):
	 * ENQ_IMMED=BIT(0), PERF=BIT(3) */
	skel->rodata->grant_caps = grant_caps ? grant_caps : 0x9;
	skel->rodata->revoke_at = revoke_at;
	skel->rodata->restore_on_detach = restore_on_detach;
	skel->rodata->perf_target = perf_target;

	SCX_OPS_LOAD(skel, k2_ops, scx_k2, uei);
	link = SCX_OPS_ATTACH(skel, k2_ops, scx_k2);
	fprintf(stderr, "[k2] up: %s grant_cgid=%llu caps=0x%llx revoke_at=%llu perf=%u\n",
		sub_cgid ? "sub" : "root",
		(unsigned long long)grant_cgid,
		(unsigned long long)skel->rodata->grant_caps,
		(unsigned long long)revoke_at, perf_target);

	while (!exit_req) {
		u32 i;

		sleep(2);
		for (i = 0; i < 12; i++)
			fprintf(stderr, "[k2] %s=%llu%s", cnt_name(i),
				(unsigned long long)read_cnt(skel, i),
				i == 3 ? "\n" : " ");
	}

	bpf_link__destroy(link);
	scx_k2__destroy(skel);
	fprintf(stderr, "[k2] exited cleanly\n");
	return 0;
}