[PATCH v2 1/2] sched_ext: Don't run ops.dequeue() with a DSQ lock held

[Qiurong Fang]

From: fangqiurong <fangqiurong@xxxxxxxxxx>

ops.dequeue() is called with the source user DSQ's lock still held on
the consume and move paths (scx_consume_dispatch_q(),
move_task_between_dsqs()) and with the terminal global/bypass DSQ's
lock still held in scx_dispatch_enqueue(). A BPF scheduler that locks
the same DSQ from ops.dequeue() - e.g. by iterating it with
bpf_iter_scx_dsq, which takes the DSQ lock on every step -
self-deadlocks.

Move the invocation after the DSQ unlock on all three paths.
SCX_TASK_IN_CUSTODY is cleared under the lock so that the callback is
invoked exactly once; it is not ordered against consumption of the
task and may run after the task has been moved to, or consumed from,
a terminal DSQ.

Fixes: ebf1ccff79c4 ("sched_ext: Fix ops.dequeue() semantics")
Acked-by: Andrea Righi <arighi@xxxxxxxxxx>
Signed-off-by: fangqiurong <fangqiurong@xxxxxxxxxx>
---
 Documentation/scheduler/sched-ext.rst |  4 +-
 kernel/sched/ext/ext.c                | 67 +++++++++++++++++++++------
 2 files changed, 56 insertions(+), 15 deletions(-)

diff --git a/Documentation/scheduler/sched-ext.rst b/Documentation/scheduler/sched-ext.rst
index 794ae80b3ba3..6e65d3b38951 100644
--- a/Documentation/scheduler/sched-ext.rst
+++ b/Documentation/scheduler/sched-ext.rst
@@ -361,7 +361,9 @@ The following briefly shows how a waking task is scheduled and executed.
    ``scx_bpf_dsq_reenq()``. The task stays in BPF custody the entire time.
 
    When a task leaves BPF scheduler custody, ``ops.dequeue()`` is invoked.
-   The dequeue can happen for different reasons, distinguished by flags:
+   The callback is not ordered against consumption of the task and may run
+   after the task has been moved to, or consumed from, a terminal DSQ. The
+   dequeue can happen for different reasons, distinguished by flags:
 
    1. **Regular dispatch**: when a task in BPF custody is dispatched to a
       terminal DSQ from ``ops.dispatch()`` (leaving BPF custody for
diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 40fa1697bdb7..13efddb729e3 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -1499,20 +1499,28 @@ static inline bool task_scx_migrating(struct task_struct *p)
 	return p->scx.sticky_cpu >= 0;
 }
 
+/* Must be called under the lock serializing @p's custody transfers. */
+static bool task_leave_custody(struct task_struct *p)
+{
+	if (!(p->scx.flags & SCX_TASK_IN_CUSTODY) || task_scx_migrating(p))
+		return false;
+
+	p->scx.flags &= ~SCX_TASK_IN_CUSTODY;
+	return true;
+}
+
 /*
  * Call ops.dequeue() if the task is in BPF custody and not migrating.
- * Clears %SCX_TASK_IN_CUSTODY when the callback is invoked.
+ * Clears %SCX_TASK_IN_CUSTODY before the callback is invoked.
  */
 static void call_task_dequeue(struct scx_sched *sch, struct rq *rq,
 			      struct task_struct *p, u64 deq_flags)
 {
-	if (!(p->scx.flags & SCX_TASK_IN_CUSTODY) || task_scx_migrating(p))
+	if (!task_leave_custody(p))
 		return;
 
 	if (SCX_HAS_OP(sch, dequeue))
 		SCX_CALL_OP_TASK(sch, dequeue, rq, p, deq_flags);
-
-	p->scx.flags &= ~SCX_TASK_IN_CUSTODY;
 }
 
 static void rq_owned_post_enq(struct scx_sched *sch, struct rq *rq,
@@ -1705,20 +1713,28 @@ static void scx_dispatch_enqueue(struct scx_sched *sch, struct rq *rq,
 	if (is_rq_owned) {
 		rq_owned_post_enq(sch, rq, dsq, p, enq_flags);
 	} else {
+		bool call_dequeue = false;
+
 		/*
 		 * Global and bypass DSQs are terminal - the task leaves the
-		 * scheduler's custody, so ops.dequeue() fires here. It can run
+		 * scheduler's custody, so ops.dequeue() fires. It can run
 		 * without @p's rq lock (finish_dispatch() passes the dispatch
 		 * rq); that's safe because dequeue_task_scx() waits on
 		 * SCX_OPSS_DISPATCHING (see the ops_state note above) and so
 		 * can't race it. A non-terminal DSQ keeps the task in custody.
+		 * The custody transfer happens under @dsq->lock so that
+		 * later consumers see the flag clear; the callback runs
+		 * unlocked - it must not run with a DSQ lock held.
 		 */
 		if (dsq->id == SCX_DSQ_GLOBAL || dsq->id == SCX_DSQ_BYPASS)
-			call_task_dequeue(sch, rq, p, 0);
+			call_dequeue = task_leave_custody(p);
 		else
 			p->scx.flags |= SCX_TASK_IN_CUSTODY;
 
 		raw_spin_unlock(&dsq->lock);
+
+		if (call_dequeue && SCX_HAS_OP(sch, dequeue))
+			SCX_CALL_OP_TASK(sch, dequeue, rq, p, 0);
 	}
 
 	/*
@@ -2373,11 +2389,13 @@ static void wakeup_preempt_scx(struct rq *rq, struct task_struct *p, int wake_fl
 		scx_schedule_reenq_local(rq, 0);
 }
 
-void scx_move_local_task_to_local_dsq(struct scx_sched *sch, struct task_struct *p,
-				      u64 enq_flags, struct scx_dispatch_q *src_dsq,
-				      struct rq *dst_rq)
+static struct scx_dispatch_q *
+__scx_move_local_task_to_local_dsq(struct scx_sched *sch,
+				   struct task_struct *p, u64 *enq_flags,
+				   struct scx_dispatch_q *src_dsq,
+				   struct rq *dst_rq)
 {
-	struct scx_dispatch_q *dst_dsq = scx_resolve_local_dsq(sch, dst_rq, p, &enq_flags);
+	struct scx_dispatch_q *dst_dsq = scx_resolve_local_dsq(sch, dst_rq, p, enq_flags);
 
 	/* @p is on @dst_rq, an rq-owned @src_dsq is covered by the rq lock */
 	if (!dsq_is_rq_owned(src_dsq))
@@ -2386,14 +2404,25 @@ void scx_move_local_task_to_local_dsq(struct scx_sched *sch, struct task_struct
 
 	WARN_ON_ONCE(p->scx.holding_cpu >= 0);
 
-	if (enq_flags & (SCX_ENQ_HEAD | SCX_ENQ_PREEMPT))
+	if (*enq_flags & (SCX_ENQ_HEAD | SCX_ENQ_PREEMPT))
 		dsq_insert_head(dst_dsq, p);
 	else
 		list_add_tail(&p->scx.dsq_list.node, &dst_dsq->list);
 
-	dsq_inc_nr(dst_dsq, p, enq_flags);
+	dsq_inc_nr(dst_dsq, p, *enq_flags);
 	p->scx.dsq = dst_dsq;
 
+	return dst_dsq;
+}
+
+void scx_move_local_task_to_local_dsq(struct scx_sched *sch, struct task_struct *p,
+				      u64 enq_flags, struct scx_dispatch_q *src_dsq,
+				      struct rq *dst_rq)
+{
+	struct scx_dispatch_q *dst_dsq;
+
+	dst_dsq = __scx_move_local_task_to_local_dsq(sch, p, &enq_flags,
+						     src_dsq, dst_rq);
 	rq_owned_post_enq(sch, dst_rq, dst_dsq, p, enq_flags);
 }
 
@@ -2628,9 +2657,14 @@ static struct rq *move_task_between_dsqs(struct scx_sched *sch,
 	if (dst_dsq->id == SCX_DSQ_LOCAL) {
 		/* @p is going from a non-local DSQ to a local DSQ */
 		if (src_rq == dst_rq) {
+			struct scx_dispatch_q *ldsq;
+
 			scx_task_unlink_from_dsq(p, src_dsq);
-			scx_move_local_task_to_local_dsq(sch, p, enq_flags, src_dsq, dst_rq);
+			ldsq = __scx_move_local_task_to_local_dsq(sch, p,
+								  &enq_flags,
+								  src_dsq, dst_rq);
 			raw_spin_unlock(&src_dsq->lock);
+			rq_owned_post_enq(sch, dst_rq, ldsq, p, enq_flags);
 		} else {
 			raw_spin_unlock(&src_dsq->lock);
 			move_remote_task_to_local_dsq(sch, p, enq_flags, src_rq, dst_rq);
@@ -2679,9 +2713,14 @@ bool scx_consume_dispatch_q(struct scx_sched *sch, struct rq *rq,
 			break;
 
 		if (rq == task_rq) {
+			struct scx_dispatch_q *ldsq;
+
 			scx_task_unlink_from_dsq(p, dsq);
-			scx_move_local_task_to_local_dsq(sch, p, enq_flags, dsq, rq);
+			ldsq = __scx_move_local_task_to_local_dsq(sch, p,
+								  &enq_flags,
+								  dsq, rq);
 			raw_spin_unlock(&dsq->lock);
+			rq_owned_post_enq(sch, rq, ldsq, p, enq_flags);
 			return true;
 		}
 
-- 
2.43.0