scx_layered: Replace task hinting with task local data

scheds/rust/scx_layered/src/bpf/main.bpf.c

@@ -22,6 +22,7 @@
 #include "intf.h"
 #include "timer.bpf.h"
 #include "util.bpf.h"
+#include "task_local_data.bpf.h"
 
 char _license[] SEC("license") = "GPL";
 
@@ -94,18 +95,10 @@ u32 nr_empty_layer_ids;
 
 UEI_DEFINE(uei);
 
-struct task_hint {
-	u64 hint;
-	u64 __reserved[3];
+struct tld_keys {
+	tld_key_t priority_key;
 };
 
-struct {
-	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
-	__uint(map_flags, BPF_F_NO_PREALLOC);
-	__type(key, int);
-	__type(value, struct task_hint);
-} scx_layered_task_hint_map SEC(".maps");
-
 struct {
 	__uint(type, BPF_MAP_TYPE_ARRAY);
 	__type(key, u32);
@@ -616,29 +609,37 @@ static struct task_ctx *lookup_task_ctx(struct task_struct *p)
 	return taskc;
 }
 
-static struct task_hint *lookup_task_hint(struct task_struct *p)
+static u64 *lookup_task_hint_priority(struct task_struct *p)
 {
-	struct task_hint *hint;
+	struct tld_object tld_obj;
+	u64 *priority;
+	int err;
 
 	if (!task_hint_map_enabled)
 		return NULL;
-	hint = bpf_task_storage_get(&scx_layered_task_hint_map, p, NULL, 0);
+
+	err = tld_object_init(p, &tld_obj);
+	if (err)
+		return NULL;
+
 	/* Only values in the range [0, 1024] are valid hints. */
-	if (hint && hint->hint > 1024)
-		hint = NULL;
-	return hint;
+	priority = tld_get_data(&tld_obj, priority_key, "thread_hint_priority", sizeof(*priority));
+	if (!priority || *priority > 1024)
+		return NULL;
+
+	return priority;
 }
 
 static int lookup_task_hint_layer_id(struct task_struct *p) {
-	struct task_hint *hint;
 	u32 *layer_idp;
+	u64 *priority;
 	u32 hint_val;
 
-	hint = lookup_task_hint(p);
-	if (!hint)
+	priority = lookup_task_hint_priority(p);
+	if (!priority)
 		return -ENOENT;
 
-	hint_val = hint->hint;
+	hint_val = *priority;
 	layer_idp = bpf_map_lookup_elem(&hint_to_layer_id_map, &hint_val);
 	if (!layer_idp)
 		return -EFAULT;
@@ -2490,11 +2491,11 @@ static __noinline bool match_one(struct layer_match *match,
 				avg_runtime_us < match->max_avg_runtime_us;
 	}
 	case MATCH_HINT_EQUALS: {
-		struct task_hint *hint = lookup_task_hint(p);
+		u64 *priority = lookup_task_hint_priority(p);
 
-		if (!hint)
+		if (!priority)
 			return false;
-		return match->hint == hint->hint;
+		return match->hint == *priority;
 	}
 
 	default:

scheds/rust/scx_layered/src/bpf/task_local_data.bpf.h

@@ -0,0 +1,237 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+#ifndef __TASK_LOCAL_DATA_BPF_H
+#define __TASK_LOCAL_DATA_BPF_H
+
+/*
+ * Task local data is a library that facilitates sharing per-task data
+ * between user space and bpf programs.
+ *
+ *
+ * USAGE
+ *
+ * A TLD, an entry of data in task local data, first needs to be created by the
+ * user space. This is done by calling user space API, TLD_DEFINE_KEY() or
+ * tld_create_key(), with the name of the TLD and the size.
+ *
+ * TLD_DEFINE_KEY(prio, "priority", sizeof(int));
+ *
+ * or
+ *
+ * void func_call(...) {
+ *     tld_key_t prio, in_cs;
+ *
+ *     prio = tld_create_key("priority", sizeof(int));
+ *     in_cs = tld_create_key("in_critical_section", sizeof(bool));
+ *     ...
+ *
+ * A key associated with the TLD, which has an opaque type tld_key_t, will be
+ * initialized or returned. It can be used to get a pointer to the TLD in the
+ * user space by calling tld_get_data().
+ *
+ * In a bpf program, tld_object_init() first needs to be called to initialized a
+ * tld_object on the stack. Then, TLDs can be accessed by calling tld_get_data().
+ * The API will try to fetch the key by the name and use it to locate the data.
+ * A pointer to the TLD will be returned. It also caches the key in a task local
+ * storage map, tld_key_map, whose value type, struct tld_keys, must be defined
+ * by the developer.
+ *
+ * struct tld_keys {
+ *     tld_key_t prio;
+ *     tld_key_t in_cs;
+ * };
+ *
+ * SEC("struct_ops")
+ * void prog(struct task_struct task, ...)
+ * {
+ *     struct tld_object tld_obj;
+ *     int err, *p;
+ *
+ *     err = tld_object_init(task, &tld_obj);
+ *     if (err)
+ *         return;
+ *
+ *     p = tld_get_data(&tld_obj, prio, "priority", sizeof(int));
+ *     if (p)
+ *         // do something depending on *p
+ */
+#include <errno.h>
+#include <bpf/bpf_helpers.h>
+
+#define TLD_ROUND_MASK(x, y) ((__typeof__(x))((y) - 1))
+#define TLD_ROUND_UP(x, y) ((((x) - 1) | TLD_ROUND_MASK(x, y)) + 1)
+
+#define TLD_MAX_DATA_CNT (__PAGE_SIZE / sizeof(struct tld_metadata) - 1)
+
+#ifndef TLD_NAME_LEN
+#define TLD_NAME_LEN 62
+#endif
+
+#ifndef TLD_KEY_MAP_CREATE_RETRY
+#define TLD_KEY_MAP_CREATE_RETRY 10
+#endif
+
+typedef struct {
+	__s16 off;
+} tld_key_t;
+
+struct tld_metadata {
+	char name[TLD_NAME_LEN];
+	__u16 size;
+};
+
+struct tld_meta_u {
+	__u8 cnt;
+	__u16 size;
+	struct tld_metadata metadata[TLD_MAX_DATA_CNT];
+};
+
+struct tld_data_u {
+	__u64 start; /* offset of tld_data_u->data in a page */
+	char data[__PAGE_SIZE - sizeof(__u64)];
+};
+
+struct tld_map_value {
+	struct tld_data_u __uptr *data;
+	struct tld_meta_u __uptr *meta;
+};
+
+typedef struct tld_uptr_dummy {
+	struct tld_data_u data[0];
+	struct tld_meta_u meta[0];
+} *tld_uptr_dummy_t;
+
+struct tld_object {
+	struct tld_map_value *data_map;
+	struct tld_keys *key_map;
+	/*
+	 * Force the compiler to generate the actual definition of tld_meta_u
+	 * and tld_data_u in BTF. Without it, tld_meta_u and u_tld_data will
+	 * be BTF_KIND_FWD.
+	 */
+	tld_uptr_dummy_t dummy[0];
+};
+
+/*
+ * Map value of tld_key_map for caching keys. Must be defined by the developer.
+ * Members should be tld_key_t and passed to the 3rd argument of tld_fetch_key().
+ */
+struct tld_keys;
+
+struct {
+	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
+	__uint(map_flags, BPF_F_NO_PREALLOC);
+	__type(key, int);
+	__type(value, struct tld_map_value);
+} tld_data_map SEC(".maps");
+
+struct {
+	__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
+	__uint(map_flags, BPF_F_NO_PREALLOC);
+	__type(key, int);
+	__type(value, struct tld_keys);
+} tld_key_map SEC(".maps");
+
+/**
+ * tld_object_init() - Initialize a tld_object.
+ *
+ * @task: The task_struct of the target task
+ * @tld_obj: A pointer to a tld_object to be initialized
+ *
+ * Return 0 on success; -ENODATA if the user space did not initialize task local data
+ * for the current task through tld_get_data(); -ENOMEM if the creation of tld_key_map
+ * fails
+ */
+__attribute__((unused))
+static int tld_object_init(struct task_struct *task, struct tld_object *tld_obj)
+{
+	int i;
+
+	tld_obj->data_map = bpf_task_storage_get(&tld_data_map, task, 0, 0);
+	if (!tld_obj->data_map)
+		return -ENODATA;
+
+	bpf_for(i, 0, TLD_KEY_MAP_CREATE_RETRY) {
+		tld_obj->key_map = bpf_task_storage_get(&tld_key_map, task, 0,
+							BPF_LOCAL_STORAGE_GET_F_CREATE);
+		if (likely(tld_obj->key_map))
+			break;
+	}
+	if (!tld_obj->key_map)
+		return -ENOMEM;
+
+	return 0;
+}
+
+/*
+ * Return the offset of TLD if @name is found. Otherwise, return the current TLD count
+ * using the nonpositive range so that the next tld_get_data() can skip fetching key if
+ * no new TLD is added or start comparing name from the first newly added TLD.
+ */
+__attribute__((unused))
+static int __tld_fetch_key(struct tld_object *tld_obj, const char *name, int i_start)
+{
+	struct tld_metadata *metadata;
+	int i, cnt, start, off = 0;
+
+	if (!tld_obj->data_map || !tld_obj->data_map->data || !tld_obj->data_map->meta)
+		return 0;
+
+	start = tld_obj->data_map->data->start;
+	cnt = tld_obj->data_map->meta->cnt;
+	metadata = tld_obj->data_map->meta->metadata;
+
+	bpf_for(i, 0, cnt) {
+		if (i >= TLD_MAX_DATA_CNT)
+			break;
+
+		if (i >= i_start && !bpf_strncmp(metadata[i].name, TLD_NAME_LEN, name))
+			return start + off;
+
+		off += TLD_ROUND_UP(metadata[i].size, 8);
+	}
+
+	return -cnt;
+}
+
+/**
+ * tld_get_data() - Retrieve a pointer to the TLD associated with the name.
+ *
+ * @tld_obj: A pointer to a valid tld_object initialized by tld_object_init()
+ * @key: The cached key of the TLD in tld_key_map
+ * @name: The name of the key associated with a TLD
+ * @size: The size of the TLD. Must be a known constant value
+ *
+ * Return a pointer to the TLD associated with @name; NULL if not found or @size is too
+ * big. @key is used to cache the key if the TLD is found to speed up subsequent calls.
+ * It should be defined as an member of tld_keys of tld_key_t type by the developer.
+ */
+#define tld_get_data(tld_obj, key, name, size)						\
+	({										\
+		void *data = NULL, *_data = (tld_obj)->data_map->data;			\
+		long off = (tld_obj)->key_map->key.off;					\
+		int cnt;								\
+											\
+		if (likely(_data)) {							\
+			if (likely(off > 0)) {						\
+				barrier_var(off);					\
+				if (likely(off < __PAGE_SIZE - size))			\
+					data = _data + off;				\
+			} else {							\
+				cnt = -(off);						\
+				if (likely((tld_obj)->data_map->meta) &&		\
+				    cnt < (tld_obj)->data_map->meta->cnt) {		\
+					off = __tld_fetch_key(tld_obj, name, cnt);	\
+					(tld_obj)->key_map->key.off = off;		\
+											\
+					if (likely(off < __PAGE_SIZE - size)) {		\
+						barrier_var(off);			\
+						if (off > 0)				\
+							data = _data + off;		\
+					}						\
+				}							\
+			}								\
+		}									\
+		data;									\
+	})
+
+#endif

scheds/rust/scx_layered/src/main.rs

@@ -2395,7 +2395,7 @@ impl<'a> Scheduler<'a> {
         // so that we can keep reusing the older map already pinned on scheduler
         // restarts.
         let layered_task_hint_map_path = &opts.task_hint_map;
-        let hint_map = &mut skel.maps.scx_layered_task_hint_map;
+        let hint_map = &mut skel.maps.tld_data_map;
         // Only set pin path if a path is provided.
         if layered_task_hint_map_path.is_empty() == false {
             hint_map.set_pin_path(layered_task_hint_map_path).unwrap();