[RFC PATCH 2/3] sched/core: Fix forceidle when lock owner has mismatchd cookie

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

Move Vasily's fix from 7cd4088f70e0 ("sched/core: Don't proxy-exec
unmatched cookie lock owners") into a separate helper and fix the
forceidle stats, balance callback queuing when rq is not the one
dictating the core cookie.

The helper will be extended in subsequent commits to allow for re-pick
with cookie from the lock owner if rq is the pick leader.

Signed-off-by: K Prateek Nayak <kprateek.nayak@amd.com>
---
 kernel/sched/core.c | 43 ++++++++++++++++++++++++++++++++++++-------
 1 file changed, 36 insertions(+), 7 deletions(-)

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 31fa6bcee616..ed11410f1416 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -6598,6 +6598,8 @@ __pick_next_task(struct rq *rq, struct rq_flags *rf)
 	BUG(); /* The idle class should always have a runnable task. */
 }
 
+static void proxy_deactivate(struct rq *rq, struct task_struct *donor);
+
 #ifdef CONFIG_SCHED_CORE
 static inline bool is_task_rq_idle(struct task_struct *t)
 {
@@ -7100,6 +7102,23 @@ static inline void sched_core_cpu_dying(unsigned int cpu)
 		rq->core = rq;
 }
 
+static struct task_struct *
+sched_core_swap_pick(struct rq *rq, struct task_struct *next)
+{
+	if (!rq->core_pick_leader) {
+		if (!rq->core->core_forceidle_count++)
+			rq->core->core_forceidle_seq++;
+
+		proxy_resched_idle(rq);
+		queue_core_balance(rq);
+		return rq->idle;
+	}
+
+	clear_task_blocked_on(rq->donor, NULL);
+	proxy_deactivate(rq, rq->donor);
+	return RETRY_TASK;
+}
+
 #else /* !CONFIG_SCHED_CORE: */
 
 static inline void sched_core_cpu_starting(unsigned int cpu) {}
@@ -7113,6 +7132,18 @@ pick_next_task(struct rq *rq, struct rq_flags *rf)
 	return __pick_next_task(rq, rf);
 }
 
+static struct task_struct *
+sched_core_swap_pick(struct rq *rq, struct task_struct *next)
+{
+	/*
+	 * Caller should always check !sched_cpu_cookie_match(rq, next)
+	 * which is always false for !CONFIG_SCHED_CORE and execution
+	 * should never reach here.
+	 */
+	BUG_ON();
+	return next;
+}
+
 #endif /* !CONFIG_SCHED_CORE */
 
 /*
@@ -7512,13 +7543,6 @@ find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf)
 	}
 	WARN_ON_ONCE(owner && !owner->on_rq);
 
-	if (owner && !sched_cpu_cookie_match(rq, owner)) {
-		if (curr_in_chain)
-			return proxy_resched_idle(rq);
-		p = donor; /* Deactivate the donor, not the runnable owner */
-		clear_task_blocked_on(p, NULL);
-		goto deactivate;
-	}
 	return owner;
 
 deactivate:
@@ -7686,6 +7710,11 @@ static void __sched notrace __schedule(int sched_mode)
 				zap_balance_callbacks(rq);
 				goto keep_resched;
 			}
+			if (!sched_cpu_cookie_match(rq, next)) {
+				next = sched_core_swap_pick(rq, next);
+				if (next == RETRY_TASK)
+					goto pick_again;
+			}
 			trace_sched_found_proxy_task(rq->donor, next, cpu);
 		}
 		if (rq->donor == prev_donor && prev != next) {
-- 
2.43.0