[PATCH 07/10] sched_ext: Add proxy destination query kfuncs

[Andrea Righi <arighi@nvidia.com>]

BPF schedulers admitting blocked proxy donors may want to know the CPU
or cid where the mutex owner will execute.

Introduce scx_bpf_task_proxy_cpu() and scx_bpf_task_proxy_cid() to
return the CPU or cid of the next mutex owner in the proxy chain.

The owner relationship may change immediately after the query, so expose
the result only as a scheduling hint. Return a negative errno when no
valid proxy destination or cid mapping is available.

Provide compatibility wrappers that return -EOPNOTSUPP when the kfuncs
are unavailable.

Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 kernel/sched/core.c                      | 28 ++++++++++++++++
 kernel/sched/ext/ext.c                   | 42 ++++++++++++++++++++++++
 kernel/sched/ext/internal.h              |  6 ++--
 kernel/sched/sched.h                     |  2 ++
 tools/sched_ext/include/scx/common.bpf.h |  2 ++
 tools/sched_ext/include/scx/compat.bpf.h | 18 ++++++++++
 6 files changed, 96 insertions(+), 2 deletions(-)

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 3d72f64ffe627..39e2689ea6c3b 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -7046,6 +7046,34 @@ find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf)
 	return NULL;
 }
 
+int task_proxy_cpu(struct task_struct *p)
+{
+	struct task_struct *owner;
+	struct mutex *mutex;
+
+	if (!sched_proxy_exec() || !READ_ONCE(p->is_blocked))
+		return -ENOENT;
+
+	guard(raw_spinlock_irqsave)(&p->blocked_lock);
+
+	mutex = __get_task_blocked_on(p);
+	if (!mutex)
+		return -ENOENT;
+
+	/*
+	 * @blocked_lock stabilizes @blocked_on and thus the mutex lifetime.
+	 * The owner is an atomic snapshot used only as a scheduling hint and
+	 * may change as soon as this function returns, so wait_lock is not
+	 * needed here.
+	 */
+	owner = __mutex_owner(mutex);
+	if (!owner)
+		return -ENOENT;
+	if (!READ_ONCE(owner->on_rq) || owner->se.sched_delayed)
+		return -ENOENT;
+
+	return task_cpu(owner);
+}
 #else /* SCHED_PROXY_EXEC */
 static struct task_struct *
 find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf)
diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 60154b25ce975..5d755d586e1cc 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -9432,6 +9432,45 @@ __bpf_kfunc bool scx_bpf_task_running(const struct task_struct *p)
 	return task_rq(p)->curr == p;
 }
 
+/**
+ * scx_bpf_task_proxy_cpu - Return the next proxy execution CPU
+ * @p: task of interest
+ *
+ * Return the CPU of the mutex owner toward which @p's scheduling context
+ * would next be migrated for proxy execution. The owner relationship can
+ * change after this function returns, so the result is only a scheduling
+ * hint and the returned CPU may not be an allowed BPF dispatch destination
+ * for @p. Returns a negative errno if no valid proxy destination is available.
+ */
+__bpf_kfunc s32 scx_bpf_task_proxy_cpu(struct task_struct *p)
+{
+	return task_proxy_cpu(p);
+}
+
+/**
+ * scx_bpf_task_proxy_cid - Return the next proxy execution cid
+ * @p: task of interest
+ *
+ * cid-addressed equivalent of scx_bpf_task_proxy_cpu(). Return the cid of the
+ * mutex owner toward which @p's scheduling context would next be migrated for
+ * proxy execution. The owner relationship can change after this function
+ * returns, so the result is only a scheduling hint. Returns a negative errno
+ * if no valid proxy destination or cid mapping is available.
+ */
+__bpf_kfunc s32 scx_bpf_task_proxy_cid(struct task_struct *p)
+{
+	s16 *tbl = READ_ONCE(scx_cpu_to_cid_tbl);
+	s32 cpu;
+
+	cpu = task_proxy_cpu(p);
+	if (cpu < 0)
+		return cpu;
+	if (!tbl)
+		return -EINVAL;
+
+	return READ_ONCE(tbl[cpu]);
+}
+
 /**
  * scx_bpf_task_cpu - CPU a task is currently associated with
  * @p: task of interest
@@ -9735,6 +9774,8 @@ BTF_ID_FLAGS(func, scx_bpf_get_possible_cpumask, KF_ACQUIRE)
 BTF_ID_FLAGS(func, scx_bpf_get_online_cpumask, KF_ACQUIRE)
 BTF_ID_FLAGS(func, scx_bpf_put_cpumask, KF_RELEASE)
 BTF_ID_FLAGS(func, scx_bpf_task_running, KF_RCU)
+BTF_ID_FLAGS(func, scx_bpf_task_proxy_cpu, KF_RCU)
+BTF_ID_FLAGS(func, scx_bpf_task_proxy_cid, KF_RCU)
 BTF_ID_FLAGS(func, scx_bpf_task_cpu, KF_RCU)
 BTF_ID_FLAGS(func, scx_bpf_task_cid, KF_RCU)
 BTF_ID_FLAGS(func, scx_bpf_locked_rq, KF_IMPLICIT_ARGS | KF_RET_NULL)
@@ -9770,6 +9811,7 @@ static const struct btf_kfunc_id_set scx_kfunc_set_any = {
  */
 BTF_KFUNCS_START(scx_kfunc_ids_cpu_only)
 BTF_ID_FLAGS(func, scx_bpf_kick_cpu, KF_IMPLICIT_ARGS)
+BTF_ID_FLAGS(func, scx_bpf_task_proxy_cpu, KF_RCU)
 BTF_ID_FLAGS(func, scx_bpf_task_cpu, KF_RCU)
 BTF_ID_FLAGS(func, scx_bpf_cpu_curr, KF_IMPLICIT_ARGS | KF_RET_NULL | KF_RCU_PROTECTED)
 BTF_ID_FLAGS(func, scx_bpf_cpu_node, KF_IMPLICIT_ARGS)
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index 8ac195934f63b..df7420b816a04 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -214,8 +214,10 @@ enum scx_ops_flags {
 
 	/*
 	 * If set, mutex-blocked tasks remain runnable as proxy donors and are
-	 * passed to ops.enqueue() with %SCX_ENQ_BLOCKED. The BPF scheduler controls
-	 * when donors are dispatched and whether they should preempt other work.
+	 * passed to ops.enqueue() with %SCX_ENQ_BLOCKED. The BPF scheduler can
+	 * query the next mutex owner's CPU or cid with scx_bpf_task_proxy_cpu()
+	 * or scx_bpf_task_proxy_cid(). It controls when donors are dispatched and
+	 * whether they should preempt work on the owner's CPU.
 	 *
 	 * If clear, mutex-blocked tasks are removed from the runqueue normally
 	 * and cannot donate their scheduling context through proxy execution.
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index dfa0cb722c00c..8e55c922e16ca 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -2472,8 +2472,10 @@ static inline bool task_is_blocked(struct task_struct *p)
 }
 
 #ifdef CONFIG_SCHED_PROXY_EXEC
+int task_proxy_cpu(struct task_struct *p);
 void sched_proxy_block_task(struct rq *rq, struct task_struct *p);
 #else
+static inline int task_proxy_cpu(struct task_struct *p) { return -EOPNOTSUPP; }
 static inline void sched_proxy_block_task(struct rq *rq, struct task_struct *p) {}
 #endif
 
diff --git a/tools/sched_ext/include/scx/common.bpf.h b/tools/sched_ext/include/scx/common.bpf.h
index e7b3ba491c5e8..98a69be2dd502 100644
--- a/tools/sched_ext/include/scx/common.bpf.h
+++ b/tools/sched_ext/include/scx/common.bpf.h
@@ -95,6 +95,8 @@ s32 scx_bpf_pick_idle_cpu(const cpumask_t *cpus_allowed, u64 flags) __ksym;
 s32 scx_bpf_pick_any_cpu_node(const cpumask_t *cpus_allowed, int node, u64 flags) __ksym __weak;
 s32 scx_bpf_pick_any_cpu(const cpumask_t *cpus_allowed, u64 flags) __ksym;
 bool scx_bpf_task_running(const struct task_struct *p) __ksym;
+s32 scx_bpf_task_proxy_cpu(struct task_struct *p) __ksym __weak;
+s32 scx_bpf_task_proxy_cid(struct task_struct *p) __ksym __weak;
 s32 scx_bpf_task_cpu(const struct task_struct *p) __ksym;
 struct rq *scx_bpf_locked_rq(void) __ksym;
 struct task_struct *scx_bpf_cpu_curr(s32 cpu) __ksym __weak;
diff --git a/tools/sched_ext/include/scx/compat.bpf.h b/tools/sched_ext/include/scx/compat.bpf.h
index 133058578668d..d25996163206a 100644
--- a/tools/sched_ext/include/scx/compat.bpf.h
+++ b/tools/sched_ext/include/scx/compat.bpf.h
@@ -133,6 +133,24 @@ static inline void scx_bpf_cid_override(const s32 *cpu_to_cid, u32 cpu_to_cid__s
 		return scx_bpf_cid_override___compat(cpu_to_cid, cpu_to_cid__sz);
 }
 
+/*
+ * v7.3: scx_bpf_task_proxy_cpu() and scx_bpf_task_proxy_cid() for querying
+ * the next proxy execution destination. Return -EOPNOTSUPP if unavailable.
+ */
+static inline s32 __COMPAT_scx_bpf_task_proxy_cpu(struct task_struct *p)
+{
+	if (bpf_ksym_exists(scx_bpf_task_proxy_cpu))
+		return scx_bpf_task_proxy_cpu(p);
+	return -EOPNOTSUPP;
+}
+
+static inline s32 __COMPAT_scx_bpf_task_proxy_cid(struct task_struct *p)
+{
+	if (bpf_ksym_exists(scx_bpf_task_proxy_cid))
+		return scx_bpf_task_proxy_cid(p);
+	return -EOPNOTSUPP;
+}
+
 /**
  * __COMPAT_is_enq_cpu_selected - Test if SCX_ENQ_CPU_SELECTED is on
  * in a compatible way. We will preserve this __COMPAT helper until v6.16.
-- 
2.55.0