/* SPDX-License-Identifier: GPL-2.0 */
/*
 * Copyright (c) 2024 Meta Platforms, Inc. and affiliates.
 * Copyright (c) 2024 Tejun Heo <tj@kernel.org>
 * Copyright (c) 2024 David Vernet <dvernet@meta.com>
 */
#ifndef __SCX_COMPAT_H
#define __SCX_COMPAT_H

#include <bpf/btf.h>
#include <bpf/libbpf.h>
#include <fcntl.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "enums_abi.autogen.h"

struct btf *__COMPAT_vmlinux_btf __attribute__((weak));

static inline void __COMPAT_load_vmlinux_btf(void)
{
	if (!__COMPAT_vmlinux_btf) {
		__COMPAT_vmlinux_btf = btf__load_vmlinux_btf();
		SCX_BUG_ON(!__COMPAT_vmlinux_btf, "btf__load_vmlinux_btf()");
	}
}

/*
 * Recover the true value of a 64-bit enum enumerator whose kernel BTF entry
 * was truncated to its low 32 bits.
 *
 * Kernels whose BTF was generated without BTF_KIND_ENUM64 support encode
 * 64-bit enums as 8-byte BTF_KIND_ENUM entries whose enumerator values only
 * carry the low 32 bits. This happens with pahole < 1.24, which predates
 * ENUM64, and with pahole passing --skip_encoding_btf_enum64 (e.g. Google's
 * Container-Optimized OS / GKE kernels deliberately pass it for backward
 * compatibility with older BTF consumers). The high bits
 * can't be recovered from kernel BTF, so substitute the value from the
 * vmlinux.h this tree was built against, cross-checked against the low 32
 * bits the kernel did provide.
 *
 * Note that this is a best-effort recovery, not a ground truth. The
 * substitution assumes the running kernel agrees with this tree's vmlinux.h
 * on the high 32 bits, but only the low 32 bits can actually be verified.
 * The cross-check is vacuous for enumerators whose value has no low bits
 * set (e.g. SCX_DSQ_FLAG_BUILTIN, __SCX_ENQ_INTERNAL_MASK,
 * SCX_ENQ_CLEAR_OPSS, SCX_ECODE_*): their lo32 is 0 and matches anything,
 * so those substitutions rest entirely on the high bits never moving. An
 * enumerator missing from the table (a kernel newer than this tree's
 * vmlinux.h, or a stale autogen table) can't be recovered at all. If a
 * substitution is ever wrong, the scheduler operates on bogus values (e.g.
 * dispatching to nonexistent DSQ ids or silently dropping flags) and can
 * wildly malfunction, which is why the mismatch and table-miss paths refuse
 * instead of guessing.
 */
static inline bool __COMPAT_recover_truncated_enum64(const char *type,
						     const char *name,
						     u32 lo32, u64 *v)
{
	static bool warned;
	size_t i;

	for (i = 0; i < sizeof(__scx_enum_abi_vals) / sizeof(__scx_enum_abi_vals[0]); i++) {
		const struct __scx_enum_abi_val *e = &__scx_enum_abi_vals[i];

		if (strcmp(e->type, type) || strcmp(e->name, name))
			continue;

		if (e->val <= (u64)UINT32_MAX) {
			*v = lo32;
			return true;
		}

		if ((u32)e->val != lo32) {
			fprintf(stderr, "ERROR: kernel BTF value of %s::%s (0x%x) doesn't match the low 32 bits of the vmlinux.h value (0x%llx); refusing to substitute\n",
				type, name, lo32, (unsigned long long)e->val);
			return false;
		}

		if (!warned) {
			fprintf(stderr,
				"WARNING: kernel BTF lacks BTF_KIND_ENUM64 encoding (generated by\n"
				"WARNING: pahole < 1.24 or with --skip_encoding_btf_enum64), so 64-bit\n"
				"WARNING: scx enum values are truncated to their low 32 bits in kernel\n"
				"WARNING: BTF. Substituting the full 64-bit values from the vmlinux.h\n"
				"WARNING: this binary was built against, cross-checked against the low\n"
				"WARNING: 32 bits the kernel does provide. The high 32 bits cannot be\n"
				"WARNING: verified: if the running kernel's actual values differ from\n"
				"WARNING: the build-time vmlinux.h (e.g. an enum that moved in a newer\n"
				"WARNING: kernel), the scheduler will operate on bogus values, such as\n"
				"WARNING: dispatching to nonexistent DSQ ids, and can wildly malfunction.\n");
			warned = true;
		}
		*v = e->val;
		return true;
	}

	/*
	 * Unknown enumerator (likely a stale autogen table). Fail
	 * pessimistically to avoid returning an invalid value.
	 */
	fprintf(stderr, "ERROR: kernel BTF truncates 64-bit enum %s::%s to 0x%x; 64-bit variant not found in vmlinux.h\n",
		type, name, lo32);
	return false;
}

static inline bool __COMPAT_read_enum(const char *type, const char *name, u64 *v)
{
	const struct btf_type *t;
	const char *n;
	s32 tid;
	__u32 i;

	__COMPAT_load_vmlinux_btf();

	tid = btf__find_by_name(__COMPAT_vmlinux_btf, type);
	if (tid < 0)
		return false;

	t = btf__type_by_id(__COMPAT_vmlinux_btf, tid);
	SCX_BUG_ON(!t, "btf__type_by_id(%d)", tid);

	if (btf_is_enum(t)) {
		struct btf_enum *e = btf_enum(t);

		for (i = 0; i < btf_vlen(t); i++) {
			n = btf__name_by_offset(__COMPAT_vmlinux_btf, e[i].name_off);
			SCX_BUG_ON(!n, "btf__name_by_offset()");
			if (!strcmp(n, name)) {
				/*
				 * Try to recover a 64-bit enum from an 8-byte
				 * BTF_KIND_ENUM that was encoded without ENUM64
				 * support (old pahole or
				 * --skip_encoding_btf_enum64). Only scx_*
				 * types are covered by the substitution table;
				 * non-scx types fall through to the raw value
				 * so this generic utility keeps working for
				 * them.
				 */
				if (t->size == 8 && !strncmp(type, "scx_", 4))
					return __COMPAT_recover_truncated_enum64(type, name,
										 (u32)e[i].val, v);
				*v = e[i].val;
				return true;
			}
		}
	} else if (btf_is_enum64(t)) {
		struct btf_enum64 *e = btf_enum64(t);

		for (i = 0; i < btf_vlen(t); i++) {
			n = btf__name_by_offset(__COMPAT_vmlinux_btf, e[i].name_off);
			SCX_BUG_ON(!n, "btf__name_by_offset()");
			if (!strcmp(n, name)) {
				*v = btf_enum64_value(&e[i]);
				return true;
			}
		}
	}

	return false;
}

#define __COMPAT_ENUM_OR_ZERO(__type, __ent)					\
({										\
	u64 __val = 0;								\
	__COMPAT_read_enum(__type, __ent, &__val);				\
	__val;									\
})

static inline bool __COMPAT_has_ksym(const char *ksym)
{
	__COMPAT_load_vmlinux_btf();
	return btf__find_by_name(__COMPAT_vmlinux_btf, ksym) >= 0;
}

static inline bool __COMPAT_struct_has_field(const char *type, const char *field)
{
	const struct btf_type *t;
	const struct btf_member *m;
	const char *n;
	s32 tid;
	__u32 i;

	__COMPAT_load_vmlinux_btf();
	tid = btf__find_by_name_kind(__COMPAT_vmlinux_btf, type, BTF_KIND_STRUCT);
	if (tid < 0)
		return false;

	t = btf__type_by_id(__COMPAT_vmlinux_btf, tid);
	SCX_BUG_ON(!t, "btf__type_by_id(%d)", tid);

	m = btf_members(t);

	for (i = 0; i < btf_vlen(t); i++) {
		n = btf__name_by_offset(__COMPAT_vmlinux_btf, m[i].name_off);
		SCX_BUG_ON(!n, "btf__name_by_offset()");
			if (!strcmp(n, field))
				return true;
	}

	return false;
}

#define SCX_OPS_FLAG(name) __COMPAT_ENUM_OR_ZERO("scx_ops_flags", #name)

#define SCX_OPS_KEEP_BUILTIN_IDLE SCX_OPS_FLAG(SCX_OPS_KEEP_BUILTIN_IDLE)
#define SCX_OPS_ENQ_LAST SCX_OPS_FLAG(SCX_OPS_ENQ_LAST)
#define SCX_OPS_ENQ_EXITING  SCX_OPS_FLAG(SCX_OPS_ENQ_EXITING)
#define SCX_OPS_SWITCH_PARTIAL SCX_OPS_FLAG(SCX_OPS_SWITCH_PARTIAL)
#define SCX_OPS_ENQ_MIGRATION_DISABLED SCX_OPS_FLAG(SCX_OPS_ENQ_MIGRATION_DISABLED)
#define SCX_OPS_ALLOW_QUEUED_WAKEUP SCX_OPS_FLAG(SCX_OPS_ALLOW_QUEUED_WAKEUP)
#define SCX_OPS_BUILTIN_IDLE_PER_NODE SCX_OPS_FLAG(SCX_OPS_BUILTIN_IDLE_PER_NODE)
#define SCX_OPS_ALWAYS_ENQ_IMMED SCX_OPS_FLAG(SCX_OPS_ALWAYS_ENQ_IMMED)

#define SCX_PICK_IDLE_FLAG(name) __COMPAT_ENUM_OR_ZERO("scx_pick_idle_cpu_flags", #name)

#define SCX_PICK_IDLE_CORE SCX_PICK_IDLE_FLAG(SCX_PICK_IDLE_CORE)
#define SCX_PICK_IDLE_IN_NODE SCX_PICK_IDLE_FLAG(SCX_PICK_IDLE_IN_NODE)

static inline long scx_hotplug_seq(void)
{
	int fd;
	char buf[32];
	char *endptr;
	ssize_t len;
	long val;

	fd = open("/sys/kernel/sched_ext/hotplug_seq", O_RDONLY);
	if (fd < 0)
		return -ENOENT;

	len = read(fd, buf, sizeof(buf) - 1);
	SCX_BUG_ON(len <= 0, "read failed (%ld)", len);
	buf[len] = 0;
	close(fd);

	errno = 0;
	val = strtoul(buf, &endptr, 10);
	SCX_BUG_ON(errno == ERANGE || endptr == buf ||
		   (*endptr != '\n' && *endptr != '\0'), "invalid num hotplug events: %ld", val);

	return val;
}

/*
 * Open the sched_ext_ops skeleton.
 *
 * struct sched_ext_ops can change over time. Two complementary mechanisms
 * keep BPF schedulers built against newer headers running on older kernels:
 *
 * 1. Load-time fix-up (SCX_OPS_OPEN()). For each optional ops callback or field
 *    added to struct sched_ext_ops, an explicit stanza below probes the
 *    running kernel's BTF via __COMPAT_struct_has_field() and, if the field
 *    is missing, clears it in the in-memory struct_ops (with a warning to
 *    stderr) before load. Handles additive changes - a new stanza must be
 *    added here for each new optional field.
 *
 * 2. Multi-variant struct_ops via compat.bpf.h::SCX_OPS_DEFINE(). That
 *    macro can be expanded to emit several variants of struct sched_ext_ops,
 *    and SCX_OPS_LOAD()/ATTACH() can pick the right one based on what the
 *    kernel supports. Needed when an existing operation has to change
 *    incompatibly (e.g. a callback signature changes); the load-time
 *    fix-up above only handles purely additive changes.
 *
 * ec7e3b0463e1 ("implement-ops") in https://github.com/sched-ext/sched_ext is
 * the current minimum required kernel version.
 *
 * COMPAT:
 * - v6.17: ops.cgroup_set_bandwidth()
 * - v6.19: ops.cgroup_set_idle()
 * - v7.1:  ops.sub_attach(), ops.sub_detach(), ops.sub_cgroup_id
 * - v7.3:  ops.rescue_bandwidth_ppt, ops.rescue_quantum_us
 */
#define __SCX_OPS_OPEN(__ops_name, __scx_name, __ops_struct) ({			\
	struct __scx_name *__oskel;						\
										\
	SCX_BUG_ON(!__COMPAT_struct_has_field(__ops_struct, "dump"),		\
		   __ops_struct ".dump() missing, kernel too old?");		\
										\
	__oskel = __scx_name##__open();						\
	SCX_BUG_ON(!__oskel, "Could not open " #__scx_name);			\
	__oskel->struct_ops.__ops_name->hotplug_seq = scx_hotplug_seq();	\
	SCX_ENUM_INIT(__oskel);							\
	__oskel;								\
})

#define SCX_OPS_OPEN(__ops_name, __scx_name) ({					\
	struct __scx_name *__skel;						\
										\
	__skel = __SCX_OPS_OPEN(__ops_name, __scx_name, "sched_ext_ops");	\
	if (__skel->struct_ops.__ops_name->cgroup_set_bandwidth &&		\
	    !__COMPAT_struct_has_field("sched_ext_ops", "cgroup_set_bandwidth")) { \
		fprintf(stderr, "WARNING: kernel doesn't support ops.cgroup_set_bandwidth()\n"); \
		__skel->struct_ops.__ops_name->cgroup_set_bandwidth = NULL;	\
	}									\
	if (__skel->struct_ops.__ops_name->cgroup_set_idle &&			\
	    !__COMPAT_struct_has_field("sched_ext_ops", "cgroup_set_idle")) { \
		fprintf(stderr, "WARNING: kernel doesn't support ops.cgroup_set_idle()\n"); \
		__skel->struct_ops.__ops_name->cgroup_set_idle = NULL;	\
	}									\
	if (__skel->struct_ops.__ops_name->sub_attach &&			\
	    !__COMPAT_struct_has_field("sched_ext_ops", "sub_attach")) {	\
		fprintf(stderr, "WARNING: kernel doesn't support ops.sub_attach()\n"); \
		__skel->struct_ops.__ops_name->sub_attach = NULL;		\
	}									\
	if (__skel->struct_ops.__ops_name->sub_detach &&			\
	    !__COMPAT_struct_has_field("sched_ext_ops", "sub_detach")) {	\
		fprintf(stderr, "WARNING: kernel doesn't support ops.sub_detach()\n"); \
		__skel->struct_ops.__ops_name->sub_detach = NULL;		\
	}									\
	if (__skel->struct_ops.__ops_name->sub_cgroup_id > 0 &&		\
	    !__COMPAT_struct_has_field("sched_ext_ops", "sub_cgroup_id")) { \
		fprintf(stderr, "WARNING: kernel doesn't support ops.sub_cgroup_id\n"); \
		__skel->struct_ops.__ops_name->sub_cgroup_id = 0;		\
	}									\
	if (__skel->struct_ops.__ops_name->rescue_bandwidth_ppt > 0 &&		\
	    !__COMPAT_struct_has_field("sched_ext_ops", "rescue_bandwidth_ppt")) { \
		fprintf(stderr, "WARNING: kernel doesn't support ops.rescue_bandwidth_ppt\n"); \
		__skel->struct_ops.__ops_name->rescue_bandwidth_ppt = 0;	\
	}									\
	if (__skel->struct_ops.__ops_name->rescue_quantum_us > 0 &&		\
	    !__COMPAT_struct_has_field("sched_ext_ops", "rescue_quantum_us")) { \
		fprintf(stderr, "WARNING: kernel doesn't support ops.rescue_quantum_us\n"); \
		__skel->struct_ops.__ops_name->rescue_quantum_us = 0;		\
	}									\
	__skel; 								\
})

/*
 * Open a cid-form (struct sched_ext_ops_cid) skeleton. The cid form postdates
 * every op the load-time fix-ups above handle, so none of them apply.
 */
#define SCX_OPS_CID_OPEN(__ops_name, __scx_name)				\
	__SCX_OPS_OPEN(__ops_name, __scx_name, "sched_ext_ops_cid")

/*
 * Associate non-struct_ops BPF programs with the scheduler's struct_ops map so
 * that scx_prog_sched() can determine which scheduler a BPF program belongs
 * to. Requires libbpf >= 1.7.
 */
#if LIBBPF_MAJOR_VERSION > 1 ||							\
	(LIBBPF_MAJOR_VERSION == 1 && LIBBPF_MINOR_VERSION >= 7)
static inline void __scx_ops_assoc_prog(struct bpf_program *prog,
					struct bpf_map *map,
					const char *ops_name)
{
	s32 err = bpf_program__assoc_struct_ops(prog, map, NULL);
	if (err)
		fprintf(stderr,
			"ERROR: Failed to associate %s with %s: %d\n",
			bpf_program__name(prog), ops_name, err);
}
#else
static inline void __scx_ops_assoc_prog(struct bpf_program *prog,
					struct bpf_map *map,
					const char *ops_name)
{
}
#endif

/* See SCX_OPS_OPEN() above for backward-compatibility handling. */
#define SCX_OPS_LOAD(__skel, __ops_name, __scx_name, __uei_name) ({		\
	struct bpf_program *__prog;						\
	UEI_SET_SIZE(__skel, __ops_name, __uei_name);				\
	SCX_BUG_ON(__scx_name##__load((__skel)), "Failed to load skel");	\
	bpf_object__for_each_program(__prog, (__skel)->obj) {			\
		if (bpf_program__type(__prog) == BPF_PROG_TYPE_STRUCT_OPS)	\
			continue;						\
		__scx_ops_assoc_prog(__prog, (__skel)->maps.__ops_name,		\
				     #__ops_name);				\
	}									\
})

/*
 * New versions of bpftool now emit additional link placeholders for BPF maps,
 * and set up BPF skeleton in such a way that libbpf will auto-attach BPF maps
 * automatically, assuming libbpf is recent enough (v1.5+). Old libbpf will do
 * nothing with those links and won't attempt to auto-attach maps.
 *
 * To maintain compatibility with older libbpf while avoiding trying to attach
 * twice, disable the autoattach feature on newer libbpf.
 */
#if LIBBPF_MAJOR_VERSION > 1 ||							\
	(LIBBPF_MAJOR_VERSION == 1 && LIBBPF_MINOR_VERSION >= 5)
#define __SCX_OPS_DISABLE_AUTOATTACH(__skel, __ops_name)			\
	bpf_map__set_autoattach((__skel)->maps.__ops_name, false)
#else
#define __SCX_OPS_DISABLE_AUTOATTACH(__skel, __ops_name) do {} while (0)
#endif

#define SCX_OPS_ATTACH(__skel, __ops_name, __scx_name) ({			\
	struct bpf_link *__link;						\
	__SCX_OPS_DISABLE_AUTOATTACH(__skel, __ops_name);			\
	SCX_BUG_ON(__scx_name##__attach((__skel)), "Failed to attach skel");	\
	__link = bpf_map__attach_struct_ops((__skel)->maps.__ops_name);		\
	SCX_BUG_ON(!__link, "Failed to attach struct_ops");			\
	__link;									\
})

#endif	/* __SCX_COMPAT_H */