[PATCH v3] sched/proxy_exec: Detect cycles in proxy walks

["zhidao su (Xiaomi)" <soolaugust@gmail.com>]

find_proxy_task() can keep walking the same blocked_on chain if the chain
contains a cycle. This happens with a simple A->B->A mutex deadlock under
proxy execution, and can leave the CPU spinning in __schedule() with the rq
lock held.

Mark each task visited during one proxy walk with the rq's pick sequence.
If the walk sees the same task again in the same pick, break the cycle by
clearing the blocked_on state at the detection point and deactivating that
task.

Clear the marker when proxy execution migrates a blocked task, so a later
walk on another CPU cannot match stale state.

Tested with a PE cycle reproducer in virtme-ng:

  buggy kernel: vng timed out without returning
  fixed kernel: WARN_ONCE "sched/pe: deadlock cycle detected"

Signed-off-by: zhidao su (Xiaomi) <soolaugust@gmail.com>
---
Changes since v2:
- Drop the proxy-chain depth limit.
- Detect cycles anywhere in the proxy walk, not only cycles back to the
  original donor.
- Use a separate per-pick marker instead of overloading blocked_donor.
- Clear the marker on proxy-exec blocked-task migration.

 include/linux/sched.h |  4 ++++
 kernel/fork.c         |  4 ++++
 kernel/sched/core.c   | 34 ++++++++++++++++++++++++++++++++--
 kernel/sched/sched.h  |  1 +
 4 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/include/linux/sched.h b/include/linux/sched.h
index 968b18a7f4702..b83a4e775d736 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -1256,6 +1256,10 @@ struct task_struct {
 	 * under preempt_disable().
 	 */
 	struct task_struct		*blocked_donor;
+#ifdef CONFIG_SCHED_PROXY_EXEC
+	int				proxy_pick_cpu;
+	u64				proxy_pick_seq; /* find_proxy_task() seq */
+#endif
 
 #ifdef CONFIG_DETECT_HUNG_TASK_BLOCKER
 	/*
diff --git a/kernel/fork.c b/kernel/fork.c
index 13e38e89a1f30..2219cce602c62 100644
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -2242,6 +2242,10 @@ __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
+	p->proxy_pick_cpu = -1;
+	p->proxy_pick_seq = 0;
+#endif
 
 #ifdef CONFIG_BCACHE
 	p->sequential_io	= 0;
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 2e7cde033a319..3db4e6c2b786a 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -6724,6 +6724,7 @@ static bool try_to_block_task(struct rq *rq, struct task_struct *p,
 }
 
 #ifdef CONFIG_SCHED_PROXY_EXEC
+
 static inline void proxy_set_task_cpu(struct task_struct *p, int cpu)
 {
 	unsigned int wake_cpu;
@@ -6830,6 +6831,10 @@ static void proxy_migrate_task(struct rq *rq, struct rq_flags *rf,
 
 	deactivate_task(rq, p, DEQUEUE_NOCLOCK);
 	proxy_set_task_cpu(p, target_cpu);
+#ifdef CONFIG_SCHED_PROXY_EXEC
+	WRITE_ONCE(p->proxy_pick_cpu, -1);
+	WRITE_ONCE(p->proxy_pick_seq, 0);
+#endif
 
 	proxy_release_rq_lock(rq, rf);
 
@@ -6839,14 +6844,14 @@ static void proxy_migrate_task(struct rq *rq, struct rq_flags *rf,
 }
 
 /*
- * Find runnable lock owner to proxy for mutex blocked donor
+ * Find runnable lock owner to proxy for a blocked donor
  *
  * Follow the blocked-on relation:
  *
  *                ,-> task
  *                |     | blocked-on
  *                |     v
- *  blocked_donor |   mutex
+ *  blocked_donor | blocking primitive
  *                |     | owner
  *                |     v
  *                `-- task
@@ -6874,6 +6879,10 @@ find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf)
 	struct task_struct *p;
 	int owner_cpu;
 
+	rq->proxy_pick_seq++;
+	if (unlikely(!rq->proxy_pick_seq))
+		rq->proxy_pick_seq++;
+
 	/* Follow blocked_on chain. */
 	for (p = donor; p->is_blocked; p = owner) {
 		/* if its PROXY_WAKING, do return migration or run if current */
@@ -6905,6 +6914,16 @@ find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf)
 			return NULL;
 		}
 
+		if (READ_ONCE(p->proxy_pick_cpu) == this_cpu &&
+		    READ_ONCE(p->proxy_pick_seq) == rq->proxy_pick_seq) {
+			WARN_ONCE(1, "sched/pe: deadlock cycle detected, pid %d\n",
+				  p->pid);
+			__clear_task_blocked_on(p, NULL);
+			goto deactivate;
+		}
+		WRITE_ONCE(p->proxy_pick_cpu, this_cpu);
+		WRITE_ONCE(p->proxy_pick_seq, rq->proxy_pick_seq);
+
 		if (task_current(rq, p))
 			curr_in_chain = true;
 
@@ -6990,6 +7009,14 @@ find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf)
 			 */
 			return proxy_resched_idle(rq);
 		}
+
+		if (READ_ONCE(owner->proxy_pick_cpu) == this_cpu &&
+		    READ_ONCE(owner->proxy_pick_seq) == rq->proxy_pick_seq) {
+			WARN_ONCE(1, "sched/pe: deadlock cycle detected, pid %d\n",
+				  p->pid);
+			__clear_task_blocked_on(p, NULL);
+			goto deactivate;
+		}
 		/*
 		 * OK, now we're absolutely sure @owner is on this
 		 * rq, therefore holding @rq->lock is sufficient to
@@ -9057,6 +9084,9 @@ void __init sched_init(void)
 		raw_spin_lock_init(&rq->cpu_epoch_lock);
 		rq->cpu_epoch_next = jiffies;
 #endif
+#ifdef CONFIG_SCHED_PROXY_EXEC
+		rq->proxy_pick_seq = 0;
+#endif
 
 		zalloc_cpumask_var_node(&rq->scratch_mask, GFP_KERNEL, cpu_to_node(i));
 	}
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 26ae13c86b699..f63c024d160e8 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -1154,6 +1154,7 @@ struct rq {
 #ifdef CONFIG_SCHED_PROXY_EXEC
 	struct task_struct __rcu	*donor;  /* Scheduling context */
 	struct task_struct __rcu	*curr;   /* Execution context */
+	u64				proxy_pick_seq;  /* find_proxy_task() seq */
 #else
 	union {
 		struct task_struct __rcu *donor; /* Scheduler context */

base-commit: 04998aa54848f15332202d0bea008d2ca1ed1713
-- 
2.43.0