[PATCH v5] 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. A simple A->B->A deadlock can leave the CPU spinning in
__schedule() with rq->lock held.

Use the rq pick sequence as a per-walk marker. Mark each task visited by
the current walk. If the walk sees the same marker again, break the cycle
by clearing blocked_on at the detection point and deactivating that task.

The marker is only consumed while holding rq->lock. Clear it when a task is
activated, so stale state from an earlier pick or another rq is not carried
into the next queued lifetime.

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"

Reviewed-by: K Prateek Nayak <kprateek.nayak@amd.com>
Signed-off-by: zhidao su (Xiaomi) <soolaugust@gmail.com>
---
Changes since v4:
- Use sched_proxy_enqueue_task() to clear stale proxy-walk state when a
  task is activated.
- Stop marking the donor early; only mark the finalized owner.
- Use pr_warn_once() for the cycle warning.

 include/linux/sched.h |  3 +++
 kernel/sched/core.c   | 19 +++++++++++++++++--
 kernel/sched/sched.h  | 11 +++++++++++
 3 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/include/linux/sched.h b/include/linux/sched.h
index 968b18a7f4702..b1c46e21bc268 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -1256,6 +1256,9 @@ struct task_struct {
 	 * under preempt_disable().
 	 */
 	struct task_struct		*blocked_donor;
+#ifdef CONFIG_SCHED_PROXY_EXEC
+	u64				proxy_pick_seq;
+#endif
 
 #ifdef CONFIG_DETECT_HUNG_TASK_BLOCKER
 	/*
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 2e7cde033a319..c1bb4f98dd64a 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -2221,6 +2221,7 @@ void activate_task(struct rq *rq, struct task_struct *p, int flags)
 	if (task_on_rq_migrating(p))
 		flags |= ENQUEUE_MIGRATED;
 
+	sched_proxy_enqueue_task(p);
 	enqueue_task(rq, p, flags);
 
 	WRITE_ONCE(p->on_rq, TASK_ON_RQ_QUEUED);
@@ -6724,6 +6725,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;
@@ -6839,14 +6841,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 +6876,8 @@ 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++;
+
 	/* Follow blocked_on chain. */
 	for (p = donor; p->is_blocked; p = owner) {
 		/* if its PROXY_WAKING, do return migration or run if current */
@@ -6990,6 +6994,14 @@ find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf)
 			 */
 			return proxy_resched_idle(rq);
 		}
+
+		if (owner->proxy_pick_seq == rq->proxy_pick_seq) {
+			pr_warn_once("sched/pe: deadlock cycle detected, pid %d\n",
+				     p->pid);
+			__clear_task_blocked_on(p, NULL);
+			goto deactivate;
+		}
+		owner->proxy_pick_seq = rq->proxy_pick_seq;
 		/*
 		 * OK, now we're absolutely sure @owner is on this
 		 * rq, therefore holding @rq->lock is sufficient to
@@ -9057,6 +9069,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 = 1;
+#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..d5d9de509dfdb 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;
 #else
 	union {
 		struct task_struct __rcu *donor; /* Scheduler context */
@@ -3092,6 +3093,16 @@ static inline void __block_task(struct rq *rq, struct task_struct *p)
 extern void activate_task(struct rq *rq, struct task_struct *p, int flags);
 extern void deactivate_task(struct rq *rq, struct task_struct *p, int flags);
 
+static inline void sched_proxy_enqueue_task(struct task_struct *p)
+{
+#ifdef CONFIG_SCHED_PROXY_EXEC
+	if (!sched_proxy_exec())
+		return;
+
+	p->proxy_pick_seq = 0;
+#endif
+}
+
 extern void wakeup_preempt(struct rq *rq, struct task_struct *p, int flags);
 
 /*

base-commit: 04998aa54848f15332202d0bea008d2ca1ed1713
-- 
2.43.0