[RFC PATCH 14/16] sched/core: Introduce chain-wakeup to activate blocked donors

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

Activating blocked task with sched_proxy_exec() needed grabbing
p->blocked_lock for all wakeups to prevent missing additions to
blocked_head during wakeup.

With proxy_enqueue_on_owner() converted to plug the race between enqueue
on owner and owner's wakeup, simplify the wakeup path with the following
three step process:

activate_blocked_task(owner)
  /* Part 1: Enqueue race */		  enqueue_on_owner(owner, p)
  owner->on_rq = TASK_ON_RQ_MIGRATING;	    list_add(&p->blocked_node, &owner->blocked_head)

  smp_mb();				    smp_mb();

  if (list_empty(&owner->blocked_head))	    if (owner->on_rq) /* Raced */
    /*					      list_del_init(&p->blocked_node)
     * Fast-path: Future additions will	      return;
     * see owner->on_rq != 0 and dequeue    /*
     * themself.			     * Owner has to observe !list_empty() after
     */					     * this point so it is safe to block.
     __activate_task(p)			     */
					    block_task(p);

  /* END: Part 1 */

enqueue_on_owner() holds mutex->wait_lock and "owner" needs to grab the
same wait_lock during unlock so "owner" cannot disappear during
enqueue_on_owner().

  /* Part 2: Prepare migration chain */

  rq_unlock(rq);

  LIST_HEAD(migration_list)
  LIST_HEAD(wakeup_list)
							try_to_wake_up()
 							  /* Sees proxy_task_is_linked(p) */
							  ttwu_runnable(donor)
  rq_lock(rq_of(owner->blocked_cpu))			    task_rq_lock(donor) /* Same rq */

  do {

    list_for_each_entry(donor, &owner->blocked_head)	  if (task_is_linked(donor)
      donor->on_rq = TASK_ON_RQ_MIGRATING /* Part 1 */	    list_del_init(&p->blocked_node)

      if (list_empty(&owner->blocked_head)		activate_blocked_task(donor)
        list_add(&doner, &migration_list)		  /* Same as left */
      else
        list_add(&donor, &wakeup_list)

      list_add(&owner, &wakeup_list)
      owner = list_next_entry(&migration_list)

  } while (!list_empty(&migration_list))

  rq_unlock(rq_of(owner->blocked_cpu))

  /* END: Part 2 */

Since owner->blocked_cpu matches with task_cpu() of the entire chain
(including tasks queued on the blocked donors, and task queued on them,
and so on), grabbing "rq_of(owner->blocked_cpu)" is enough to stall any
concurrent wakeups for any of the linked tasks.

If try_to_wake_up() wins, donor will remove itself from the chain and
continue its wakeup as if it is a blocked task waking up.

If activate_blocked_task() wins, it will set all the donors to
TASK_ON_RQ_MIGRATING, repeating Part 1, to prevent other tasks from
queueing on it form then on and iterates its "blocked_head". This is
done in a breadth-first fashion until full chain is set to
TASK_ON_RQ_MIGRATING.

"owner->blocked_head" traversal and manipulations are still guarded by
"owner->blocked_lock" which has not been depicted here.

  /* Part 3: Wakeup chain */

  rq_lock(rq);

  __activate_task(p); /* Activate top owner first */

  /* Activate the entire chain next */
  list_for_each_entry(donor, &owner->blocked_head)
    list_del_init(&donor->blocked_node)
    __activate_task(donor)

  /* END: Part 3 */

  The wakeup steps skipped for donors is performormed during the chain
  wakeup.

Advantages of this method:

o Does not involve juggling p->pi_lock, rq_lock, owner->blocked_lock
  in addition to  checking task state. Task on chain is fully blocked
  and only one rq_lock + onwer_blocked lock can hold off
  ttwu_runnable() from waking them.

o No need to borrow reference to "owner" before queuing. The mechanism
  ensures p->blocked_owner remains valid as long as is task is on chain.

Disadvantages:

o Since chain wakeups funnel through a single CPU, it is possible for
  increased lock contention however large chain wakeups and concurrent
  wakeups of donors for reasons other than an unlock is rare.

Co-developed-by: John Stultz <jstultz@google.com>
Signed-off-by: John Stultz <jstultz@google.com>
Signed-off-by: K Prateek Nayak <kprateek.nayak@amd.com>
---
 kernel/sched/core.c | 183 +++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 181 insertions(+), 2 deletions(-)

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index e6abcea97683..3153ad007447 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -2247,6 +2247,8 @@ void activate_task(struct rq *rq, struct task_struct *p, int en_flags)
 
 #ifdef CONFIG_SCHED_PROXY_EXEC
 
+static inline void proxy_set_task_cpu(struct task_struct *p, int cpu);
+
 static void __proxy_dequeue_from_owner(struct task_struct *p)
 {
 	list_del_init(&p->blocked_node);
@@ -2337,6 +2339,183 @@ static bool proxy_try_dequeue_from_owner(struct task_struct *p)
 	return true;
 }
 
+static void
+proxy_activate_blocked_task(struct rq *rq, struct task_struct *p, int en_flags)
+{
+	int iowait_count = 0, load_contrib_count = 0;
+	struct task_struct *donor, *owner, *tmp;
+	int blocked_cpu, this_cpu = cpu_of(rq);
+	LIST_HEAD(migration_head);
+	LIST_HEAD(wakeup_head);
+	struct rq *blocked_rq;
+	bool needs_migration;
+
+	WRITE_ONCE(p->on_rq, TASK_ON_RQ_MIGRATING);
+	ASSERT_EXCLUSIVE_WRITER(p->on_rq);
+
+	/*
+	 * Pairs against smp_mb() in proxy_enqueue_on_owner() which
+	 * orders ownwer->on_rq state against the blocked_head addition.
+	 */
+	smp_mb();
+
+	/*
+	 * Fast-path: Tasks blocking on us will see p->on_rq updated and
+	 * will bail out natturally in proxy_enqueue_on_owner().
+	 *
+	 * proxy_enqueue_on_owner() holds the wait_lock to prevent owner
+	 * from running and disappearing before the transient task can
+	 * observe p->on_rq change and dequeue itself.
+	 */
+	if (list_empty(&p->blocked_head)) {
+		__activate_task(rq, p, en_flags);
+		return;
+	}
+
+	/*
+	 * Slow-path: Since we need to wakeup tasks (and task queued on
+	 * them, and task queued on them, and ... you get the gist) we
+	 * need to grab some locks.
+	 *
+	 * Lucky for us, we have simplified this via p->blocked_cpu
+	 * which is the task_cpu() for the entire chain. Wakeup is
+	 * broken into two parts.
+	 *
+	 * Part1: Under the rq_lock + blocked lock, mark all tasks on
+	 * the chain as TASK_ON_RQ_MIGRATING in breadth-first manner.
+	 *
+	 * Part2: Active all the tasks in bulk.
+	 */
+
+	/* Part 1: Prepare the blocked donor chain. */
+	blocked_cpu = p->blocked_cpu;
+	blocked_rq = cpu_rq(blocked_cpu);
+
+	/*
+	 * Blocked chain is linked to a different rq. Drop the rq_lock
+	 * and migrate the chain over before wkaing it up here.
+	 */
+	needs_migration = blocked_rq != rq;
+	if (needs_migration) {
+		raw_spin_rq_unlock(rq);
+		raw_spin_rq_lock(blocked_rq);
+	}
+
+	owner = p;
+
+	do {
+		guard(raw_spinlock)(&owner->blocked_lock);
+
+		/* Iterate the tasks blocked on this owner. */
+		list_for_each_entry_safe(donor, tmp, &owner->blocked_head, blocked_node) {
+			/* Everything hinges on this assumption. */
+			WARN_ON_ONCE(task_cpu(donor) != blocked_cpu);
+			WARN_ON_ONCE(donor->se.sched_delayed);
+
+			WRITE_ONCE(donor->on_rq, TASK_ON_RQ_MIGRATING);
+			ASSERT_EXCLUSIVE_WRITER(donor->on_rq);
+
+			if (donor->in_iowait) {
+				delayacct_blkio_end(donor);
+				iowait_count++;
+			}
+
+			if (donor->sched_contributes_to_load)
+				load_contrib_count++;
+
+			/*
+			 * Pairs against the smp_mb() in proxy_enqueue_on_owner().
+			 * See the comment at the beginning of the function.
+			 */
+			smp_mb();
+
+			/*
+			 * If the donor has more task queued on it,
+			 * add it to the migration list else directly
+			 * queue it onto the wakeup list.
+			 */
+			if (!list_empty(&donor->blocked_head)) {
+				list_move_tail(&donor->blocked_node, &migration_head);
+			} else {
+				list_move_tail(&donor->blocked_node, &wakeup_head);
+			}
+		}
+
+		owner = list_first_entry_or_null(&migration_head,
+						 typeof(*owner),
+						 blocked_node);
+
+		/*
+		 * Move the owner to the wakeup list before preparing
+		 * the sub-chain queued on it.
+		 */
+		if (owner)
+			list_move_tail(&owner->blocked_node, &wakeup_head);
+
+	} while (owner);
+
+	/* Part 2: Bulk wakeup */
+
+	atomic_sub(iowait_count, &blocked_rq->nr_iowait);
+
+	if (needs_migration) {
+		raw_spin_rq_unlock(blocked_rq);
+		raw_spin_rq_lock(rq);
+
+		update_rq_clock(rq);
+		en_flags |= ENQUEUE_NOCLOCK;
+	}
+
+	rq->nr_uninterruptible -= load_contrib_count;
+
+	if (!(en_flags & ENQUEUE_NOCLOCK)) {
+		update_rq_clock(rq);
+		en_flags |= ENQUEUE_NOCLOCK;
+	}
+
+	/* Activate the top owner first. */
+	__activate_task(rq, p, en_flags);
+
+	/*
+	 * Add ENQUEUE_MIGRATED for all tasks on the chain
+	 * since they are all sched_migrated_on_blocking
+	 * and have skipped ->migrate_task_rq() callback.
+	 */
+	en_flags |= ENQUEUE_MIGRATING;
+
+	list_for_each_entry_safe(donor, tmp, &wakeup_head, blocked_node) {
+		__proxy_dequeue_from_owner(donor);
+
+		/*
+		 * Correct the task_cpu() before activating.
+		 *
+		 * Since chain preparation already adds enough barrier
+		 * after storing p->on_rq = MIGRATING, no additional
+		 * barriers are required after modifying p->is_linked
+		 * via __proxy_dequeue_from_owner() above.
+		 */
+		if (needs_migration)
+			proxy_set_task_cpu(donor, this_cpu);
+
+		__activate_task(rq, donor, en_flags);
+		wakeup_preempt(rq, donor, en_flags);
+		/*
+		 * Blocked tasks do not add push callbacks and it is
+		 * safe to skip calling ->task_woken() here.
+		 */
+	}
+}
+
+static void activate_blocked_task(struct rq *rq, struct task_struct *p, int en_flags)
+{
+	if (!sched_proxy_exec()) {
+		__activate_task(rq, p, en_flags);
+		return;
+	}
+
+	proxy_activate_blocked_task(rq, p, en_flags);
+}
+
 #else /* !CONFIG_SCHED_PROXY_EXEC */
 
 static bool proxy_try_dequeue_from_owner(struct task_struct *p)
@@ -2344,13 +2523,13 @@ static bool proxy_try_dequeue_from_owner(struct task_struct *p)
 	return false;
 }
 
-#endif /* CONFIG_SCHED_PROXY_EXEC */
-
 static void activate_blocked_task(struct rq *rq, struct task_struct *p, int en_flags)
 {
 	__activate_task(rq, p, en_flags);
 }
 
+#endif /* CONFIG_SCHED_PROXY_EXEC */
+
 void deactivate_task(struct rq *rq, struct task_struct *p, int flags)
 {
 	WARN_ON_ONCE(flags & DEQUEUE_SLEEP);
-- 
2.34.1