[RFC PATCH 06/16] sched/core: Queue blocked donor onto sleeping owner for chain-wakeup

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

Add a blocked_donor to sleeping owner's "blocked_head" when
find_proxy_task() resolves to a blocked task. Once added to the list,
double check if the owner has woken up by checking "owner->on_rq".

If the owner has woken up, remove the task from "blocked_head" and
continue try find_proxy_task() agiain to re-evaluate the state of owner
and take the correct steps.

Since find_proxy_task() is called with wait_lock held, which the owner
needs during unlock, it is guaranteed that owner cannot disappear under
us in the process.

The added data memebers in task_struct serve the following pusrpose:

- blocked_head: Contains the blocked donors queued on us
- blocked_node: The list head used to queue onto blocked_head of a
  sleeping owner
- sleeping_owner: Sleeping owner on which the task is queued.

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>
---
 include/linux/sched.h |  5 ++++
 init/init_task.c      |  5 ++++
 kernel/fork.c         |  5 ++++
 kernel/sched/core.c   | 68 +++++++++++++++++++++++++++++++++++++++++--
 4 files changed, 81 insertions(+), 2 deletions(-)

diff --git a/include/linux/sched.h b/include/linux/sched.h
index 373bcc0598d1..bf0f4b6be7c4 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -1251,6 +1251,11 @@ struct task_struct {
 
 	struct mutex			*blocked_on;	/* lock we're blocked on */
 	raw_spinlock_t			blocked_lock;
+#ifdef CONFIG_SCHED_PROXY_EXEC
+	struct list_head		blocked_head;  /* tasks blocked on this task */
+	struct list_head		blocked_node;  /* our entry on someone elses blocked_head */
+	struct task_struct		*sleeping_owner; /* task our blocked_node is enqueued on */
+#endif
 
 	/*
 	 * The task that is boosting this task; a back link for the current
diff --git a/init/init_task.c b/init/init_task.c
index b67ef6040a65..a097c0def4c4 100644
--- a/init/init_task.c
+++ b/init/init_task.c
@@ -211,6 +211,11 @@ struct task_struct init_task __aligned(L1_CACHE_BYTES) = {
 						 &init_task.alloc_lock),
 #endif
 	.blocked_donor = NULL,
+#ifdef CONFIG_SCHED_PROXY_EXEC
+	.blocked_head = LIST_HEAD_INIT(init_task.blocked_head),
+	.blocked_node = LIST_HEAD_INIT(init_task.blocked_node),
+	.sleeping_owner = NULL,
+#endif
 #ifdef CONFIG_RT_MUTEXES
 	.pi_waiters	= RB_ROOT_CACHED,
 	.pi_top_task	= NULL,
diff --git a/kernel/fork.c b/kernel/fork.c
index f0e2e131a9a5..88f2b6e08c46 100644
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -2247,6 +2247,11 @@ __latent_entropy struct task_struct *copy_process(
 
 	p->blocked_on = NULL; /* not blocked yet */
 	p->blocked_donor = NULL; /* nobody is boosting p yet */
+#ifdef CONFIG_SCHED_PROXY_EXEC
+	INIT_LIST_HEAD(&p->blocked_head);
+	INIT_LIST_HEAD(&p->blocked_node);
+	p->sleeping_owner = NULL;
+#endif
 
 #ifdef CONFIG_BCACHE
 	p->sequential_io	= 0;
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index b4a2771eb290..e53967a126a9 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -2229,6 +2229,16 @@ void activate_task(struct rq *rq, struct task_struct *p, int en_flags)
 	__activate_task(rq, p, en_flags | ENQUEUE_MIGRATING);
 }
 
+#ifdef CONFIG_SCHED_PROXY_EXEC
+
+static void __proxy_dequeue_from_owner(struct task_struct *p)
+{
+	list_del_init(&p->blocked_node);
+	WRITE_ONCE(p->sleeping_owner, NULL);
+}
+
+#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);
@@ -6872,6 +6882,43 @@ static void proxy_migrate_task(struct rq *rq, struct rq_flags *rf,
 	proxy_reacquire_rq_lock(rq, rf);
 }
 
+static void proxy_enqueue_on_owner(struct rq *rq, struct task_struct *owner,
+				   struct task_struct *p)
+{
+	lockdep_assert_rq_held(rq);
+	lockdep_assert_held(&owner->blocked_lock);
+
+	WARN_ON(!p->on_rq);
+	WARN_ON(p->sleeping_owner);
+
+	WRITE_ONCE(p->sleeping_owner, owner);
+	list_add(&p->blocked_node, &owner->blocked_head);
+	proxy_resched_idle(rq);
+
+	/*
+	 * Order against __activate_blocked_task_slowpath() checking
+	 * owner->blocked_list after setting owner->on_rq.
+	 */
+	smp_mb();
+
+	if (READ_ONCE(owner->on_rq)) {
+		/*
+		 * owner has woken up and may miss activating us.
+		 * Remove ourself from "owner->blocked_head" and try
+		 * find_proxy_task() again considering the owner's
+		 * new state.
+		 */
+		__proxy_dequeue_from_owner(p);
+		return;
+	}
+
+	/*
+	 * Owner is fully blocked. __activate_blocked_task_slowpath()
+	 * will see us on the list during wakeup and DTRT.
+	 */
+	block_task(rq, p, READ_ONCE(p->__state));
+}
+
 /*
  * Find runnable lock owner to proxy for mutex blocked donor
  *
@@ -6962,8 +7009,25 @@ find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf)
 			/* XXX Don't handle blocked owners yet */
 			if (curr_in_chain)
 				return proxy_resched_idle(rq);
-			__clear_task_blocked_on(p, NULL);
-			goto deactivate;
+			/*
+			 * If !@owner->on_rq, holding @rq->lock will not pin the task,
+			 * so we cannot drop @mutex->wait_lock until we're sure its a blocked
+			 * task on this rq.
+			 *
+			 * We use @owner->blocked_lock to serialize against ttwu_activate().
+			 * Either we see its new owner->on_rq or it will see our list_add().
+			 */
+			WARN_ON(owner == p);
+
+			raw_spin_unlock(&p->blocked_lock);
+			raw_spin_lock(&owner->blocked_lock);
+
+			proxy_enqueue_on_owner(rq, owner, p);
+
+			raw_spin_unlock(&owner->blocked_lock);
+			raw_spin_lock(&p->blocked_lock);
+
+			return NULL; /* retry task selection */
 		}
 
 		owner_cpu = task_cpu(owner);
-- 
2.34.1