[RFC PATCH 11/16] sched/core: Prepare to inspect ->is_linked alongside ->on_rq during wakeup

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

To handle wakeup of blocked donors queued on owner's blocked_head via
ttwu_runnable(), prepare to inspect p->is_blocked in addition to
p->on_rq on paths that demand grabbing rq_lock to prevent "->on_rq"
transitions.

Unionize p->on_rq with p->is_linked to derive a 16-bit p->needs_rq_sync
state that can be atomically inspected on all architectures.
p->needs_rq_sync check follows the same ordering requirements as
p->on_rq check.

Subsequent tasks will guarntee ordering of p->is_linked against p->on_rq
transitions ensuring that loading p->needs_rq_sync state is enough to
cover both and the readers don't falsely observe them to be cleared if
read individually during their transition.

get_wchan() has been reflowed using guards as a pert of the change to
prevent awkward indentation.

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

diff --git a/include/linux/sched.h b/include/linux/sched.h
index ec65abf090a4..76de348fb88b 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -852,9 +852,23 @@ struct task_struct {
 	struct alloc_tag		*alloc_tag;
 #endif
 
+	union {
+		struct {
+			u8		on_rq;
+			u8		is_linked;
+		};
+		/*
+		 * Allows inspecting on_rq and is_linked
+		 * atomically with a single read outside
+		 * the task_rq_lock().
+		 *
+		 * See the comment above the load in
+		 * try_to_wake_up() for ordering
+		 * guarantees.
+		 */
+		u16			needs_rq_sync;
+	};
 	u8				on_cpu;
-	u8				on_rq;
-	u8				is_linked;
 	u8				is_blocked;
 
 	struct __call_single_node	wake_entry;
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index d3cdfcb400d1..957f30632b82 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -2171,12 +2171,16 @@ unsigned long get_wchan(struct task_struct *p)
 		return 0;
 
 	/* Only get wchan if task is blocked and we can keep it that way. */
-	raw_spin_lock_irq(&p->pi_lock);
+	guard(raw_spinlock_irq)(&p->pi_lock);
+
 	state = READ_ONCE(p->__state);
+	if (state == TASK_RUNNING || state == TASK_WAKING)
+		return 0;
+
 	smp_rmb(); /* see try_to_wake_up() */
-	if (state != TASK_RUNNING && state != TASK_WAKING && !p->on_rq)
+
+	if (!READ_ONCE(p->needs_rq_sync))
 		ip = __get_wchan(p);
-	raw_spin_unlock_irq(&p->pi_lock);
 
 	return ip;
 }
@@ -4365,10 +4369,49 @@ int try_to_wake_up(struct task_struct *p, unsigned int state, int wake_flags)
 		 * Pairs with the LOCK+smp_mb__after_spinlock() on rq->lock in
 		 * __schedule().  See the comment for smp_mb__after_spinlock().
 		 *
+		 * Additionally, this also guards against activation of blocked
+		 * donors queued on a sleeping owner when a wakeup races for
+		 * the same task with sched_proxy_exec().
+		 *
+		 * proxy_enqueue_on_owner()
+		 *   STORE p->is_linked = 1
+		 *
+		 *   block_task()			try_to_wake_up()
+		 *     smp_mb()
+		 *     STRORE p->on_rq = 0		  LOAD p->needs_rq_sync
+		 *
+		 * The read below will atomically observe either p->on_rq or
+		 * p->is_linked being set and put the task on ttwu_runnable()
+		 * path.
+		 *
+		 * On the chain-wakeup path, p->on_rq is first transitioned to
+		 * TASK_ON_RQ_MIGRATING before p->is_linked is cleared and the
+		 * task is transitioned to TASK_ON_RQ_QUEUED
+		 *
+		 * activate_blocked_task()
+		 *   task_rq_lock()			ttwu_runnable()
+		 *   STORE donor->on_rq = MIGRATING	  task_rq_lock()
+		 *					    # STALL
+		 *   smp_wmb();
+		 *
+		 *   proxy_dequeue_from_owner()
+		 *     STORE donor->is_linked = 0
+		 *
+		 *   activate_task()
+		 *     STORE donor->on_rq = QUEUED
+		 *
+		 *   rq_unlock()			    # ACQUIRED
+		 *					  LOAD donor->needs_rq_sync
+		 *
+		 * In this case too try_to_wake_up() will correctly observe
+		 * either p->on_rq != 0 or p->is_linked != 0 by atomically
+		 * reading p->needs_rq_sync when the task is being activated as
+		 * a part of chain wakeup.
+		 *
 		 * A similar smp_rmb() lives in __task_needs_rq_lock().
 		 */
 		smp_rmb();
-		if (READ_ONCE(p->on_rq) && ttwu_runnable(p, wake_flags))
+		if (READ_ONCE(p->needs_rq_sync) && ttwu_runnable(p, wake_flags))
 			break;
 
 		/*
@@ -4488,7 +4531,7 @@ static bool __task_needs_rq_lock(struct task_struct *p)
 	 * See try_to_wake_up() for a longer comment.
 	 */
 	smp_rmb();
-	if (p->on_rq)
+	if (READ_ONCE(p->needs_rq_sync))
 		return true;
 
 	/*
-- 
2.34.1