[RFC PATCH 12/16] sched:core: Add MIGRATING flags when blocking and activating linked donors

[K Prateek Nayak <kprateek.nayak@amd.com>]

Linked donor can go through a proxy_set_task_cpu() on the way to
blocking. Add DEQUEUE_MIGRATING with DEQUEUE_SLEEP when blocking proxy
donors and add ENQUEUE_MIGRATING on the activation path if the task woke
up for reason other than proxy chain-wakeup.

Signed-off-by: K Prateek Nayak <kprateek.nayak@amd.com>
---
 include/linux/sched.h | 1 +
 kernel/sched/core.c   | 9 ++++++++-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/include/linux/sched.h b/include/linux/sched.h
index 76de348fb88b..a05240b52806 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -1003,6 +1003,7 @@ struct task_struct {
 	unsigned			sched_contributes_to_load:1;
 	unsigned			sched_migrated:1;
 	unsigned			sched_task_hot:1;
+	unsigned			sched_migrated_on_blocking:1;
 
 	/* Force alignment to the next boundary: */
 	unsigned			:0;
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 957f30632b82..ff499d4532ad 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -2285,6 +2285,10 @@ static void block_task(struct rq *rq, struct task_struct *p, unsigned long task_
 		!(task_state & TASK_NOLOAD) &&
 		!(task_state & TASK_FROZEN);
 
+	p->sched_migrated_on_blocking = !!p->is_linked;
+	if (p->is_linked)
+		flags |= DEQUEUE_MIGRATING;
+
 	if (unlikely(is_special_task_state(task_state)))
 		flags |= DEQUEUE_SPECIAL;
 
@@ -3415,7 +3419,7 @@ void set_task_cpu(struct task_struct *p, unsigned int new_cpu)
 	trace_sched_migrate_task(p, new_cpu);
 
 	if (task_cpu(p) != new_cpu) {
-		if (p->sched_class->migrate_task_rq)
+		if (!p->sched_migrated_on_blocking && p->sched_class->migrate_task_rq)
 			p->sched_class->migrate_task_rq(p, new_cpu);
 		p->se.nr_migrations++;
 		perf_event_task_migrate(p);
@@ -3863,6 +3867,9 @@ ttwu_do_activate(struct rq *rq, struct task_struct *p, int wake_flags,
 	if (p->sched_contributes_to_load)
 		rq->nr_uninterruptible--;
 
+	if (p->sched_migrated_on_blocking)
+		en_flags |= ENQUEUE_MIGRATING;
+
 	if (wake_flags & WF_RQ_SELECTED)
 		en_flags |= ENQUEUE_RQ_SELECTED;
 	if (wake_flags & WF_MIGRATED)
-- 
2.34.1