[RFC PATCH 09/16] sched/core: Track CPU where the task was blocked on

[K Prateek Nayak]

Track the CPU where task was blocked on when running with
sched_proxy_exec().

This is used to re-direct the activation of blocked donors queued on the
blocked_head via the said CPU in the activation slowpath that will be
added in the subsequent patches.

p->wake_cpu or task_cpu() is not sufficient for this purpose since
p->wake_cpu is not stable when !task_on_rq_queued() and there is a
window between set_task_cpu() and activet_blocked_task() in the wakeup
path that needs to be plugged in.

Signed-off-by: K Prateek Nayak <kprateek.nayak@xxxxxxx>
---
 include/linux/sched.h |  1 +
 kernel/sched/core.c   | 16 ++++++++++++++++
 kernel/sched/sched.h  |  6 ++++++
 3 files changed, 23 insertions(+)

diff --git a/include/linux/sched.h b/include/linux/sched.h
index bf0f4b6be7c4..2360fdb1ce3f 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -1255,6 +1255,7 @@ struct task_struct {
 	struct list_head		blocked_head;  /* tasks blocked on this task */
 	struct list_head		blocked_node;  /* our entry on someone elses blocked_head */
 	struct task_struct		*sleeping_owner; /* task our blocked_node is enqueued on */
+	int				blocked_cpu;	 /* CPU where task was blocked. */
 #endif
 
 	/*
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index d635f5452a1b..b5cc5612a12c 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -3773,6 +3773,22 @@ static inline void proxy_reset_donor(struct rq *rq)
 	resched_curr(rq);
 }
 
+void __proxy_block_task(struct task_struct *p)
+{
+	if (!sched_proxy_exec())
+		return;
+
+	/*
+	 * This is the CPU where the __active_blocked_donor() slowpath will
+	 * resume activation of blocked donors queued on this task.
+	 *
+	 * This exists to cover the race between set_task_cpu() changing the
+	 * task_cpu() before task is transitioned to TASK_ON_RQ_QUEUED in
+	 * activate_blocked_task().
+	 */
+	WRITE_ONCE(p->blocked_cpu, task_cpu(p));
+}
+
 /*
  * Checks to see if task p has been proxy-migrated to another rq
  * and needs to be returned. If so, we deactivate the task here
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 56d9c09c485e..0d90506c34f1 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -1445,11 +1445,15 @@ static inline void rq_set_donor(struct rq *rq, struct task_struct *t)
 {
 	rcu_assign_pointer(rq->donor, t);
 }
+
+void __proxy_block_task(struct task_struct *p);
 #else
 static inline void rq_set_donor(struct rq *rq, struct task_struct *t)
 {
 	/* Do nothing */
 }
+
+static inline void __proxy_block_task(struct task_struct *p) {}
 #endif
 
 #ifdef CONFIG_SCHED_CORE
@@ -3039,6 +3043,8 @@ static inline void __block_task(struct rq *rq, struct task_struct *p)
 		delayacct_blkio_start();
 	}
 
+	__proxy_block_task(p);
+
 	ASSERT_EXCLUSIVE_WRITER(p->on_rq);
 
 	/*
-- 
2.34.1