sched_ext: Track the cpu a task is runnable on

Add p->scx.runnable_cpu, the cpu @p is runnable on, or -1 when it is not.
It is stamped as @p joins the runnable_list (set_task_runnable()) and
cleared as it leaves (clr_task_runnable()), both under the rq lock.

task_cpu() can't answer "is @p on this rq" reliably: a remote wakeup changes
it under @p's pi_lock alone, without the source rq lock, so it can read as
the locked rq while @p is really elsewhere. runnable_cpu changes only under
the rq lock, so a caller holding an rq lock can compare against it to know
whether that is @p's current rq.

Signed-off-by: Tejun Heo <tj@kernel.org>
Reviewed-by: Andrea Righi <arighi@nvidia.com>

1 files changed, 9 insertions, 0 deletions

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 35aca50796fd..9ae8d78738b4 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -1865,11 +1865,19 @@ static void set_task_runnable(struct rq *rq, struct task_struct *p)
 	 * appended to the runnable_list.
 	 */
 	list_add_tail(&p->scx.runnable_node, &rq->scx.runnable_list);
+
+	/*
+	 * Record the rq @p is runnable on, maintained under the rq lock so it
+	 * stays valid unlike task_cpu(), which a remote wakeup can move under
+	 * pi_lock alone.
+	 */
+	WRITE_ONCE(p->scx.runnable_cpu, cpu_of(rq));
 }
 
 static void clr_task_runnable(struct task_struct *p, bool reset_runnable_at)
 {
 	list_del_init(&p->scx.runnable_node);
+	WRITE_ONCE(p->scx.runnable_cpu, -1);
 	if (reset_runnable_at)
 		p->scx.flags |= SCX_TASK_RESET_RUNNABLE_AT;
 }
@@ -3531,6 +3539,7 @@ void init_scx_entity(struct sched_ext_entity *scx)
 	RB_CLEAR_NODE(&scx->dsq_priq);
 	scx->sticky_cpu = -1;
 	scx->holding_cpu = -1;
+	scx->runnable_cpu = -1;
 	INIT_LIST_HEAD(&scx->runnable_node);
 	scx->runnable_at = jiffies;
 	scx->ddsp_dsq_id = SCX_DSQ_INVALID;