[RFC PATCH 3/3] sched/core: Swap to lock owner for core-wide pick on core-cookie mismatch

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

Proxy donor can have a different core cookie compared to the lock owner.
Simply running a lock owner with a mismatched core-cookie in not an
option as this breaks the core scheduling guarantees.

To prevent a cookie mismatch, use the "rq->core_pick_leader" indicator
to know if the CPU with the blocked donor is the one with the highest
priority task. If the rq contains the highest priority task, it is
that rq's pick which dictates the core-cookie.

Once the find_proxy_task() chain stabilizes on the rq (find_proxy_task()
returns the valid lock owner queued on the same rq), set a second
"rq->core_pick_blocked_donor" and retry the pick.

The flattened flow looks as follows:

  donor = pick_next_task()
    rq_set_donor(rq, donor)

  if (donor->is_blocked)
     next = find_proxy_task()

  /*
   * next has a different cookie compared to doner and
   * rq->core_pick_leader is set.
   */
  sched_core_swap_task()
    rq->core_pick_blocked_donor = 1; /* rq->donor was the last pick */
    return RETRY_TASK;

  pick_next_task()
    donor = rq->donor;
    owner = NULL;

    if (rq->core_pick_blocked_donor)
      owner = find_proxy_task(donor);
      goto restart_multi;

    /*
     * Finish the core-pick and validate that this rq
     * is still the pick leaader.
     */
    if (rq->core_pick_blocked_donor &&
        rq->core_pick_leader &&
        rq->core_pick == donor)
      rq->core_pick = owner; /* Swapped donor with owner. */

    rq->core->coore_cookie = rq->core_pick->cookie; /* owner cookie */

    /* Continue with the rest */

    return rq->core_pick; /* Owner */

    if (sched_core_retain_donor(rq))
       rq->sched_core_proxy_pick = 0 /* Consume */

Since the owner is not blocked, it continues to run and rq->donor is
preserved while the "rq->core_pick_blocked_donor" gets consumed.

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

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index ed11410f1416..ab3b3104cd86 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -6598,7 +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);
+static struct task_struct *
+find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf);
 
 #ifdef CONFIG_SCHED_CORE
 static inline bool is_task_rq_idle(struct task_struct *t)
@@ -6643,10 +6644,16 @@ extern void task_vruntime_update(struct rq *rq, struct task_struct *p, bool in_f
 
 static void queue_core_balance(struct rq *rq);
 
+static __always_inline bool sched_core_proxy_pick(struct rq *rq)
+{
+	return sched_proxy_exec() && unlikely(rq->core_pick_blocked_donor);
+}
+
 static struct task_struct *
 pick_next_task(struct rq *rq, struct rq_flags *rf)
 	__must_hold(__rq_lockp(rq))
 {
+	struct task_struct *owner = NULL, *donor = rq->donor;
 	struct task_struct *next, *p, *max;
 	const struct cpumask *smt_mask;
 	bool fi_before = false;
@@ -6728,6 +6735,24 @@ pick_next_task(struct rq *rq, struct rq_flags *rf)
 	 */
 	rq->core->core_task_seq++;
 
+	/* Last core_pick resolved to a blocked_donor! */
+	if (sched_core_proxy_pick(rq)) {
+		WARN_ON_ONCE(!donor->is_blocked);
+
+		donor->blocked_donor = NULL;
+		owner = find_proxy_task(rq, donor, rf);
+		if (owner && owner != rq->idle)
+			goto restart_multi;
+		/*
+		 * Something changed in the proxy chain!
+		 * Retry pick like normal. Increment core_task_seq
+		 * since the core-wide lock might have been dropped
+		 * during proxy-migration in find_proxy_task().
+		 */
+		rq->core_pick_blocked_donor = false;
+		rq->core->core_task_seq++;
+	}
+
 	/*
 	 * Optimize for common case where this CPU has no cookies
 	 * and there are no cookied tasks running on siblings.
@@ -6786,6 +6811,21 @@ pick_next_task(struct rq *rq, struct rq_flags *rf)
 	}
 
 	rq_max->core_pick_leader = true;
+
+	if (sched_core_proxy_pick(rq)) {
+		WARN_ON_ONCE(!owner);
+
+		/*
+		 * Core-wide pick resolved to the same state as last time.
+		 * Swap the donor with the lock owner and continue the
+		 * rest of the pick sequence.
+		 */
+		if (rq->core_pick_leader && rq->core_pick == donor) {
+			rq_max->core_pick = owner;
+			max = owner;
+		}
+	}
+
 	cookie = rq->core->core_cookie = max->core_cookie;
 
 	/*
@@ -6883,9 +6923,16 @@ pick_next_task(struct rq *rq, struct rq_flags *rf)
 	}
 
 out_set_next:
-	put_prev_set_next_task(rq, rq->donor, next);
-	if (rq->core->core_forceidle_count && next == rq->idle)
-		queue_core_balance(rq);
+	/*
+	 * If this is a redo for setting the owner's core-cookie,
+	 * put_prev_set_next_task() was already done during the
+	 * last pick and there is nothing to do now.
+	 */
+	if (!sched_core_proxy_pick(rq)) {
+		put_prev_set_next_task(rq, rq->donor, next);
+		if (rq->core->core_forceidle_count && next == rq->idle)
+			queue_core_balance(rq);
+	}
 
 	return next;
 }
@@ -7114,11 +7161,21 @@ sched_core_swap_pick(struct rq *rq, struct task_struct *next)
 		return rq->idle;
 	}
 
-	clear_task_blocked_on(rq->donor, NULL);
-	proxy_deactivate(rq, rq->donor);
+	rq->core_pick_blocked_donor = true;
 	return RETRY_TASK;
 }
 
+static bool sched_core_retain_donor(struct rq *rq)
+{
+	bool retain = rq->core_pick_blocked_donor;
+
+	if (!sched_core_enabled(rq))
+		return false;
+
+	rq->core_pick_blocked_donor = false;
+	return retain;
+}
+
 #else /* !CONFIG_SCHED_CORE: */
 
 static inline void sched_core_cpu_starting(unsigned int cpu) {}
@@ -7144,6 +7201,11 @@ sched_core_swap_pick(struct rq *rq, struct task_struct *next)
 	return next;
 }
 
+static bool sched_core_retain_donor(struct rq *rq)
+{
+	return false;
+}
+
 #endif /* !CONFIG_SCHED_CORE */
 
 /*
@@ -7694,8 +7756,11 @@ static void __sched notrace __schedule(int sched_mode)
 	if (sched_proxy_exec()) {
 		struct task_struct *prev_donor = rq->donor;
 
-		rq_set_donor(rq, next);
-		next->blocked_donor = NULL;
+		if (!sched_core_retain_donor(rq)) {
+			rq_set_donor(rq, next);
+			next->blocked_donor = NULL;
+		}
+
 		if (unlikely(next->is_blocked)) {
 			next = find_proxy_task(rq, next, &rf);
 			if (!next) {
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index b80a5eb384f0..541b2ae673d8 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -1352,6 +1352,7 @@ struct rq {
 
 	/* private state for sched_proxy_exec()  */
 	bool			core_pick_leader;
+	bool			core_pick_blocked_donor;
 
 	/* shared state -- careful with sched_core_cpu_deactivate() */
 	unsigned int		core_task_seq;
-- 
2.43.0