[PATCH sched_ext/for-6.18] sched_ext: Acquire task reference in scx_bpf_cpu_curr()

[Andrea Righi <arighi@nvidia.com>]

scx_bpf_cpu_curr() has been introduced to retrieve the current task of a
given runqueue, allowing schedulers to interact with that task.

The kfunc assumes that it is always called in an RCU context, but this
is not always guaranteed and some BPF schedulers can trigger the
following warning:

  WARNING: suspicious RCU usage
  sched_ext: BPF scheduler "cosmos_1.0.2_gd0e71ca_x86_64_unknown_linux_gnu_debug" enabled
  6.17.0-rc1 #1-NixOS Not tainted
  -----------------------------
  kernel/sched/ext.c:6415 suspicious rcu_dereference_check() usage!

The correct behavior is to acquire a reference to the returned task, so
the scheduler can safely access it and then release it with
bpf_task_release().

Update the kfunc and the corresponding compatibility helper to implement
reference acquisition and prevent potential RCU warnings.

Fixes: 20b158094a1ad ("sched_ext: Introduce scx_bpf_cpu_curr()")
Cc: Christian Loehle <christian.loehle@arm.com>
Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 kernel/sched/ext.c                       | 19 +++++++++++++++----
 tools/sched_ext/include/scx/compat.bpf.h |  8 +++++++-
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index 477eccf023388..760f49a9c846c 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -6392,16 +6392,27 @@ __bpf_kfunc struct rq *scx_bpf_locked_rq(void)
 }
 
 /**
- * scx_bpf_cpu_curr - Return remote CPU's curr task
+ * scx_bpf_cpu_curr - Acquire a reference and return the current task
+ * running on a target CPU
  * @cpu: CPU of interest
  *
- * Callers must hold RCU read lock (KF_RCU).
+ * A task acquired by this kfunc must be released by calling
+ * bpf_task_release().
  */
 __bpf_kfunc struct task_struct *scx_bpf_cpu_curr(s32 cpu)
 {
+	struct task_struct *p;
+
 	if (!kf_cpu_valid(cpu, NULL))
 		return NULL;
-	return rcu_dereference(cpu_rq(cpu)->curr);
+
+	rcu_read_lock();
+	p = rcu_dereference(cpu_rq(cpu)->curr);
+	if (p && !refcount_inc_not_zero(&p->rcu_users))
+		p = NULL;
+	rcu_read_unlock();
+
+	return p;
 }
 
 /**
@@ -6569,7 +6580,7 @@ BTF_ID_FLAGS(func, scx_bpf_task_running, KF_RCU)
 BTF_ID_FLAGS(func, scx_bpf_task_cpu, KF_RCU)
 BTF_ID_FLAGS(func, scx_bpf_cpu_rq)
 BTF_ID_FLAGS(func, scx_bpf_locked_rq, KF_RET_NULL)
-BTF_ID_FLAGS(func, scx_bpf_cpu_curr, KF_RET_NULL | KF_RCU)
+BTF_ID_FLAGS(func, scx_bpf_cpu_curr, KF_RET_NULL | KF_ACQUIRE)
 #ifdef CONFIG_CGROUP_SCHED
 BTF_ID_FLAGS(func, scx_bpf_task_cgroup, KF_RCU | KF_ACQUIRE)
 #endif
diff --git a/tools/sched_ext/include/scx/compat.bpf.h b/tools/sched_ext/include/scx/compat.bpf.h
index dd9144624dc99..551635746f25c 100644
--- a/tools/sched_ext/include/scx/compat.bpf.h
+++ b/tools/sched_ext/include/scx/compat.bpf.h
@@ -238,13 +238,19 @@ static inline bool __COMPAT_is_enq_cpu_selected(u64 enq_flags)
 static inline struct task_struct *__COMPAT_scx_bpf_cpu_curr(int cpu)
 {
 	struct rq *rq;
+	struct task_struct *p;
 
 	if (bpf_ksym_exists(scx_bpf_cpu_curr))
 		return scx_bpf_cpu_curr(cpu);
 
+	bpf_rcu_read_lock();
 	rq = scx_bpf_cpu_rq(cpu);
+	p = rq ? rq->curr : NULL;
+	if (p)
+		p = bpf_task_from_pid(p->pid);
+	bpf_rcu_read_unlock();
 
-	return rq ? rq->curr : NULL;
+	return p;
 }
 
 /*
-- 
2.51.0