sched_ext: Fix the bypass dispatch claim for schedulers bypassed by an ancestor

scx_bypass() claims bypass dispatch when the scheduler's depth goes from
zero to one and releases it when the depth returns to zero, but propagates
only the depth to the descendants, not the claim. A descendant that starts
bypassing with its ancestor never claims, and one whose own bypass ends
under the ancestor's keeps its claim, as its depth does not return to zero.

A sub-scheduler enable whose bypass ends under a PM suspend thus leaks its
claim: the parent keeps hosting bypass dispatch and the sub-scheduler's next
bypass trips the WARN_ON_ONCE() on the claim. A sub-scheduler disable that
starts under a PM suspend has no claim, so once the resume lifts the root's
bypass its tasks sit on the parent's bypass DSQs with nobody consuming them.

Claim for every descendant whose depth goes from zero to one in the
propagation and release every scheduler in the subtree whose depth returned
to zero once the tasks have moved out of the bypass DSQs. A parent that is
bypassing itself when a descendant claims consumes only its own bypass DSQs,
empty until it leaves bypass and becomes the host.

Fixes: aa2a0a1968 ("sched_ext: Separate bypass dispatch enabling from bypass depth tracking")
Cc: stable@vger.kernel.org # v7.1+
Signed-off-by: Tejun Heo <tj@kernel.org>
Reviewed-by: Andrea Righi <arighi@nvidia.com>

kernel/sched/ext/ext.c

@@ -6099,15 +6099,12 @@ static void enable_bypass_dsp(struct scx_sched *sch)
 
 	/*
 	 * When a sub-sched bypasses, its tasks are queued on the bypass DSQs of
-	 * the nearest non-bypassing ancestor or root. As enable_bypass_dsp() is
-	 * called iff @sch is not already bypassed due to an ancestor bypassing,
-	 * we can assume that the parent is not bypassing and thus will be the
-	 * host of the bypass DSQs.
-	 *
-	 * While the situation may change in the future, the following
-	 * guarantees that the nearest non-bypassing ancestor or root has bypass
-	 * dispatch enabled while a descendant is bypassing, which is all that's
-	 * required.
+	 * the nearest non-bypassing ancestor or root. The parent hosts them
+	 * unless it is bypassing too. Then they sit further up and the parent
+	 * consumes only its own bypass DSQs, empty until it leaves bypass and
+	 * becomes the host. Either way the nearest non-bypassing ancestor or
+	 * root has bypass dispatch enabled while a descendant is bypassing,
+	 * which is all that's required.
 	 *
 	 * scx_bypass_dsp_enabled() test is used to determine whether to enter
 	 * the bypass dispatch handling path from both bypassing and hosting
@@ -6232,16 +6229,19 @@ void scx_bypass(struct scx_sched *sch, bool bypass)
 	 * Bypass state is propagated to all descendants - an scx_sched bypasses
 	 * if itself or any of its ancestors are in bypass mode. A descendant's
 	 * depth is the sum of its own bypasses and its ancestors', so every
-	 * call propagates whether or not @sch changed state.
+	 * call propagates whether or not @sch changed state. The descendants'
+	 * claims are released below with @sch's, after their tasks have moved.
 	 */
 	raw_spin_lock(&scx_sched_lock);
 	scx_for_each_descendant_pre(pos, sch) {
 		if (pos == sch)
 			continue;
-		if (bypass)
-			inc_bypass_depth(pos);
-		else
+		if (bypass) {
+			if (inc_bypass_depth(pos))
+				enable_bypass_dsp(pos);
+		} else {
 			dec_bypass_depth(pos);
+		}
 	}
 	raw_spin_unlock(&scx_sched_lock);
 
@@ -6330,8 +6330,12 @@ void scx_bypass(struct scx_sched *sch, bool bypass)
 	}
 
 	/* disarming must come after moving all tasks out of the bypass DSQs */
-	if (!bypass)
-		scx_disable_bypass_dsp(sch);
+	if (!bypass) {
+		guard(raw_spinlock)(&scx_sched_lock);
+		scx_for_each_descendant_pre(pos, sch)
+			if (!pos->bypass_depth)
+				scx_disable_bypass_dsp(pos);
+	}
 unlock:
 	raw_spin_unlock_irqrestore(&scx_bypass_lock, flags);
 }