[PATCHSET v4 sched_ext/for-7.4-fixes] sched_ext: Fix bypassed sub-scheduler tasks stalling when the host lacks caps

[PATCHSET v4 sched_ext/for-7.4-fixes] sched_ext: Fix bypassed sub-scheduler tasks stalling when the host lacks caps

[Tejun Heo]

Hello,

v4 (rebased onto for-7.4-fixes, the stall watchdog fix posted separately):
- Rebase onto sched_ext/for-7.4-fixes, as for-7.4 is in mainline.
- Depends on "sched_ext: Fix the stall watchdog testing and blaming the
  wrong sched under bypass", posted separately for sched_ext/for-7.4-fixes:
  https://lore.kernel.org/r/20261026204300.3618150-1-tj@kernel.org
  It replaces v3's 19/17 (the stall watchdog blaming a bypassing owner),
  which is dropped: the watchdog now blames the sched running the task and
  tests that sched's timeout (Andrea, sashiko).
- Fold v3's 18/17, a NONE rescuee leaving sched_ext with rescue.curr
  dangling, into #13, which introduces the bug (sashiko).
- Add v3's 20/17, a lone IMMED task waiting behind a kept prev, as #18
  (sashiko).
- Skip a task the donee would refuse on a balancing donor instead of
  reenqueueing it at home, where it went to the tail of its own DSQ on every
  pass (#17, Andrea). Correct its kerneldoc and condense its description.
- Update descriptions and comments (#3, #13, #18).

v3 (uncounted reenqueues under bypass, served time from the charge):
- Add a patch testing bypass before the reenqueue count in
  scx_do_enqueue_task(). A bypassing scheduler's reenqueues are the kernel's
  placements, and counting them could strand a bypassed task and deadlock
  the disable, which relies on bypass keeping every task running (#3,
  sashiko).
- Add a patch counting rescue served time from the charged runtime. The
  rescuee's sum_exec_runtime follows the mutex owner under proxy execution
  while the slice follows the donor, so the two disagreed in both
  directions. scx_rescue_keep() now WARNs and ends the rescue if its restore
  is refused (#6, sashiko).
- With rescue disabled, admit a migration-disabled or migration-pending task
  to the local DSQ on a cpu where the host lacks BASE instead of leaving it
  on the unconsumed bypass DSQ (#13).
- Update descriptions and comments, among them #17's rescue-off and
  zero-interval cases (#4, #7, #9, #10, #11, #16, #17).

v2 (keep a bypassed rescuee in dispatch_one()):
- Add a patch expiring a bypassing prev's slice in dispatch_one() before
  its keep test and dropping the bypass test from it. A bypassed rescuee
  now keeps the cpu on its kernel slice and scx_dispatch_sched()'s
  post-dispatch test no longer keeps a prev that dispatch_one() declined
  (#10, sashiko).
- Cover a rescuee that passes the keep-last test in the dispatch_one()
  comment and description (#6, sashiko).
- Update the descriptions and correct the refusal causes listed in #15
  (#11, #15).

v3: https://lore.kernel.org/r/20261023175608.4120910-1-tj@kernel.org
v2: https://lore.kernel.org/r/20261023052452.2889061-1-tj@kernel.org
v1: https://lore.kernel.org/r/20261022213836.2259069-1-tj@kernel.org

A bypassing sub-scheduler's tasks run through the per-cpu bypass DSQs of
its host, the nearest ancestor that is not bypassing. Caps are per cpu, and
a cpu's dispatch descends into a scheduler only where it holds
SCX_CAP_BASE, so a bypass DSQ on a cpu where the host lacks BASE is never
consumed. Bypass placement and the bypass load balancer ignored that. A
bypassed task could be queued or moved onto such a cpu and sit there until
the host regained the cap or the bypass window ended, and a host holding
BASE but not ENQ looped a bypassed task between its bypass DSQ and the
reject drain instead of running it.

The fixes make SCX_CAP_BASE the one cap a bypassed task needs and keep
bypassed tasks off the cpus where the host lacks it. A bypass consume puts
the task on the cpu being dispatched for, to run right away, which is what
BASE grants. Where the host lacks BASE on a task's cpu, the task runs under
rescue at bounded bandwidth instead of waiting for a dispatch that can't
reach it. The bypass load balancer spreads tasks over the cpus where the
host holds BASE only, and when the host loses BASE on a cpu it drains that
cpu's bypass DSQ to the cpus that can still run the tasks, or sends them
home to be placed again.

Running bypassed tasks under rescue leans on the rescue slice harder than
before and exposed gaps there: a cpu could go idle with a runnable rescuee
queued on it, the kernel's own slice writes under bypass served the grant a
tick at a time, and a scheduler could extend the grant past the rescue's
bandwidth. The grant is now a kernel-owned slice that the kernel's slice
housekeeping leaves alone and schedulers can shorten but not extend, and a
cpu no longer idles over a runnable rescuee.

Tested with a three-level hierarchy cycling bypass under pinned spinners,
fork and context-switch storms, with caps revoked under queued tasks. The
balancer's remote move and the drain were exercised, and the holding_cpu
handshake, the cap-revoke placement and the rescue keep were model checked
for small cpu counts.

Based on sched_ext/for-7.4-fixes (651694bc47a1) with the stall watchdog fix
above applied.

This patchset contains the following 18 patches.

 0001 sched_ext: Pack scx_ent_flags' multi-bit fields at the top
 0002 sched_ext: Clear holding_cpu where the claim is confirmed
 0003 sched_ext: Fix bypassed tasks stranded by the reenqueue limit
 0004 sched_ext: Fix bypassed tasks looping on a host without SCX_CAP_ENQ
 0005 sched_ext: Move refill_task_slice_dfl() below apply_slice_vtime()
 0006 sched_ext: Fix rescue served time disagreeing with the slice under proxy execution
 0007 sched_ext: Restore a rescuee's slice after ops.stopping()
 0008 sched_ext: Fix a cpu idling with a runnable rescuee
 0009 sched_ext: Mark the rescue grant as a kernel slice with SCX_TASK_KERNEL_SLICE
 0010 sched_ext: Refuse extending a kernel slice
 0011 sched_ext: Fix a cpu idling after put_prev_task_scx() queues its task on the local DSQ
 0012 sched_ext: Expire prev's slice in dispatch_one() under bypass
 0013 sched_ext: Fix bypassed tasks stalling when the host lacks SCX_CAP_BASE
 0014 sched_ext: Fix the bypass load balancer stranding tasks on cpus where the host lacks SCX_CAP_BASE
 0015 sched_ext: Factor out bypass_lb_move() and insert with the donor DSQ unlocked
 0016 sched_ext: Let the bypass load balancer move tasks homed on other cpus
 0017 sched_ext: Fix queued bypassed tasks stalling when the host loses SCX_CAP_BASE
 0018 sched_ext: Fix a lone IMMED task waiting behind a kept prev

The patches are organized as follows:

- 0001-0002 prepare: pack scx_ent_flags' multi-bit fields at the top so that
  the single-bit flags stay contiguous, and clear holding_cpu where the
  claim is confirmed.
- 0003 keeps a bypassing scheduler's reenqueues out of the reenqueue limit,
  as the kernel places its tasks and a strand under bypass can deadlock the
  disable.
- 0004 consumes bypass DSQs with SCX_ENQ_IMMED, so that a host with BASE but
  not ENQ runs its bypassed tasks instead of looping them.
- 0005-0010 fix the rescue slice handling that bypassed tasks under rescue
  depend on: count served time from the charge, restore the remainder after
  ops.stopping(), keep a cpu from idling with a runnable rescuee, mark the
  grant as a kernel slice that the kernel's slice writes leave alone, and
  refuse extending it.
- 0011 reschedules a cpu after put_prev_task_scx() queues its task on the
  local DSQ through the plain insert too.
- 0012 expires a bypassing prev's slice in dispatch_one() before its keep
  test instead of refusing to keep it, so a bypassed rescuee keeps the cpu
  on its kernel slice and the post-dispatch test can't keep a declined prev.
- 0013 sends a bypassed task to rescue when the host lacks BASE on its cpu.
- 0014-0017 make the bypass load balancer cap-aware: capacity and donees
  from BASE cpus only, a factored-out move that inserts with the donor DSQ
  unlocked, moves of tasks homed on other cpus, and the drain of a cpu's
  bypass DSQ when the host loses BASE there.
- 0018 schedules the IMMED reenqueue in dispatch_one() when dispatch keeps
  prev, so that a lone IMMED task queued during the dispatch doesn't wait
  behind prev.

The patchset is also available in the following git branch:

 git://git.kernel.org/pub/scm/linux/kernel/git/tj/sched_ext.git bypass-caps-v4

diffstat follows. Thanks.

 include/linux/sched/ext.h                       |  22 +-
 kernel/sched/ext/ext.c                          | 603 +++++++++++++++++-------
 kernel/sched/ext/inlines.h                      |  15 +-
 kernel/sched/ext/internal.h                     |  22 +-
 kernel/sched/ext/sub.c                          |  80 ++--
 kernel/sched/ext/sub.h                          |  16 +
 kernel/sched/sched.h                            |   2 +-
 tools/sched_ext/include/scx/enum_defs.autogen.h |   1 +
 tools/sched_ext/include/scx/enums_abi.autogen.h |  29 +-
 9 files changed, 569 insertions(+), 221 deletions(-)

--
tejun

[PATCH 01/18] sched_ext: Pack scx_ent_flags' multi-bit fields at the top

[Tejun Heo]

The following patches add a flag for a kernel-owned slice to scx_ent_flags.
The two 3-bit fields, the task state at bits 8 to 10 and the reenqueue
reason at bits 12 to 14, sit right above the single-bit flags, so the new
flag would have to take bit 11 or 15, out of sequence with the rest.

Pack the fields at the top of the word, below SCX_TASK_CURSOR, the state at
bits 25 to 27 and the reason at bits 28 to 30, so that the single-bit flags
continue from bit 8.

The renumbering changes the values a BPF scheduler reads from p->scx.flags.
Nothing in the kernel or the in-tree schedulers hardcodes the numbers. The
enum ABI table in tools/sched_ext/include/scx/enums_abi.autogen.h, kept for
repairing enum values that a kernel's BTF truncates, is regenerated for the
two fields. In the scx repository, scx_eevdf tests SCX_TASK_REENQ_IMMED and
SCX_TASK_REENQ_PREEMPTED against SCX_TASK_REENQ_REASON_MASK. Those three
names belong in its load-time enum table, so that they follow the running
kernel. No behavior change.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 include/linux/sched/ext.h                     | 11 +++++---
 .../sched_ext/include/scx/enums_abi.autogen.h | 28 +++++++++----------
 2 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/include/linux/sched/ext.h b/include/linux/sched/ext.h
index efd635a99309..b341c328948e 100644
--- a/include/linux/sched/ext.h
+++ b/include/linux/sched/ext.h
@@ -107,7 +107,10 @@ enum scx_ent_flags {
 	SCX_TASK_RUN_TRACKED	= 1 << 7, /* task is in an ops.running()/stopping() session */
 
 	/*
-	 * Bits 8 to 10 are used to carry task state:
+	 * The multi-bit fields are packed at the top, below SCX_TASK_CURSOR, so
+	 * that the single-bit flags stay contiguous from bit 0.
+	 *
+	 * Bits 25 to 27 are used to carry task state:
 	 *
 	 * NONE		ops.init_task() not called yet
 	 * INIT_BEGIN	ops.init_task() in flight; see sched_ext_dead()
@@ -116,7 +119,7 @@ enum scx_ent_flags {
 	 * ENABLED	fully initialized and in sched_ext
 	 * DEAD		terminal state set by sched_ext_dead()
 	 */
-	SCX_TASK_STATE_SHIFT	= 8,
+	SCX_TASK_STATE_SHIFT	= 25,
 	SCX_TASK_STATE_BITS	= 3,
 	SCX_TASK_STATE_MASK	= ((1 << SCX_TASK_STATE_BITS) - 1) << SCX_TASK_STATE_SHIFT,
 
@@ -128,7 +131,7 @@ enum scx_ent_flags {
 	SCX_TASK_DEAD		= 5 << SCX_TASK_STATE_SHIFT,
 
 	/*
-	 * Bits 12 to 14 are used to carry reenqueue reason. In addition to
+	 * Bits 28 to 30 are used to carry reenqueue reason. In addition to
 	 * %SCX_ENQ_REENQ flag, ops.enqueue() can also test for
 	 * %SCX_TASK_REENQ_REASON_NONE to distinguish reenqueues.
 	 *
@@ -139,7 +142,7 @@ enum scx_ent_flags {
 	 * CAP		sub-sched cap miss, see p->scx.reenq_reason_*
 	 * PROXY	proxy state prevented a remote DSQ transfer
 	 */
-	SCX_TASK_REENQ_REASON_SHIFT = 12,
+	SCX_TASK_REENQ_REASON_SHIFT = 28,
 	SCX_TASK_REENQ_REASON_BITS = 3,
 	SCX_TASK_REENQ_REASON_MASK = ((1 << SCX_TASK_REENQ_REASON_BITS) - 1) << SCX_TASK_REENQ_REASON_SHIFT,
 
diff --git a/tools/sched_ext/include/scx/enums_abi.autogen.h b/tools/sched_ext/include/scx/enums_abi.autogen.h
index 1186daa9cb14..05a460611c4d 100644
--- a/tools/sched_ext/include/scx/enums_abi.autogen.h
+++ b/tools/sched_ext/include/scx/enums_abi.autogen.h
@@ -120,24 +120,24 @@ static const struct __scx_enum_abi_val __scx_enum_abi_vals[]
 	{ "scx_ent_flags", "SCX_TASK_IMMED", 0x20LLU },
 	{ "scx_ent_flags", "SCX_TASK_PROTECTED", 0x40LLU },
 	{ "scx_ent_flags", "SCX_TASK_RUN_TRACKED", 0x80LLU },
-	{ "scx_ent_flags", "SCX_TASK_STATE_SHIFT", 0x8LLU },
+	{ "scx_ent_flags", "SCX_TASK_STATE_SHIFT", 0x19LLU },
 	{ "scx_ent_flags", "SCX_TASK_STATE_BITS", 0x3LLU },
-	{ "scx_ent_flags", "SCX_TASK_STATE_MASK", 0x700LLU },
+	{ "scx_ent_flags", "SCX_TASK_STATE_MASK", 0xe000000LLU },
 	{ "scx_ent_flags", "SCX_TASK_NONE", 0x0LLU },
-	{ "scx_ent_flags", "SCX_TASK_INIT_BEGIN", 0x100LLU },
-	{ "scx_ent_flags", "SCX_TASK_INIT", 0x200LLU },
-	{ "scx_ent_flags", "SCX_TASK_READY", 0x300LLU },
-	{ "scx_ent_flags", "SCX_TASK_ENABLED", 0x400LLU },
-	{ "scx_ent_flags", "SCX_TASK_DEAD", 0x500LLU },
-	{ "scx_ent_flags", "SCX_TASK_REENQ_REASON_SHIFT", 0xcLLU },
+	{ "scx_ent_flags", "SCX_TASK_INIT_BEGIN", 0x2000000LLU },
+	{ "scx_ent_flags", "SCX_TASK_INIT", 0x4000000LLU },
+	{ "scx_ent_flags", "SCX_TASK_READY", 0x6000000LLU },
+	{ "scx_ent_flags", "SCX_TASK_ENABLED", 0x8000000LLU },
+	{ "scx_ent_flags", "SCX_TASK_DEAD", 0xa000000LLU },
+	{ "scx_ent_flags", "SCX_TASK_REENQ_REASON_SHIFT", 0x1cLLU },
 	{ "scx_ent_flags", "SCX_TASK_REENQ_REASON_BITS", 0x3LLU },
-	{ "scx_ent_flags", "SCX_TASK_REENQ_REASON_MASK", 0x7000LLU },
+	{ "scx_ent_flags", "SCX_TASK_REENQ_REASON_MASK", 0x70000000LLU },
 	{ "scx_ent_flags", "SCX_TASK_REENQ_NONE", 0x0LLU },
-	{ "scx_ent_flags", "SCX_TASK_REENQ_KFUNC", 0x1000LLU },
-	{ "scx_ent_flags", "SCX_TASK_REENQ_IMMED", 0x2000LLU },
-	{ "scx_ent_flags", "SCX_TASK_REENQ_PREEMPTED", 0x3000LLU },
-	{ "scx_ent_flags", "SCX_TASK_REENQ_CAP", 0x4000LLU },
-	{ "scx_ent_flags", "SCX_TASK_REENQ_PROXY", 0x5000LLU },
+	{ "scx_ent_flags", "SCX_TASK_REENQ_KFUNC", 0x10000000LLU },
+	{ "scx_ent_flags", "SCX_TASK_REENQ_IMMED", 0x20000000LLU },
+	{ "scx_ent_flags", "SCX_TASK_REENQ_PREEMPTED", 0x30000000LLU },
+	{ "scx_ent_flags", "SCX_TASK_REENQ_CAP", 0x40000000LLU },
+	{ "scx_ent_flags", "SCX_TASK_REENQ_PROXY", 0x50000000LLU },
 	{ "scx_ent_flags", "SCX_TASK_CURSOR", 0xffffffff80000000LLU },
 	{ "scx_exit_code", "SCX_ECODE_RSN_HOTPLUG", 0x100000000LLU },
 	{ "scx_exit_code", "SCX_ECODE_RSN_CGROUP_OFFLINE", 0x200000000LLU },
-- 
2.56.0

[PATCH 02/18] sched_ext: Clear holding_cpu where the claim is confirmed

[Tejun Heo]

The bypass load balancer will take tasks off a remote cpu's bypass DSQ with
unlink_dsq_and_switch_rq_lock() in the following patches and insert them
into another bypass DSQ. Clearing holding_cpu after the claim is up to each
claimant, so the balancer would need a clear of its own.

unlink_dsq_and_switch_rq_lock() marks a task with holding_cpu for the gap
between dropping the DSQ lock and taking the task's rq lock, so that a
dequeue in that gap can tell the claimant it lost. Once the claimant holds
the rq lock and finds the marker intact, nothing can take the task away
under it and the marker has no reader left.

Each claimant clears it at its own site after that point:

 - dispatch_to_local_dsq(), which runs the handshake inline, in its same-rq
   and fallback branches
 - scx_proxy_reject_task() for the proxy reject
 - the dequeue inside deactivate_task() for both remote moves

Clear it where the claim is confirmed instead and drop the explicit clears.
No behavior change.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c      | 19 +++++++++++--------
 kernel/sched/ext/internal.h | 11 ++++++-----
 2 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index a9621bbdb183..0c0046766542 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -2865,7 +2865,6 @@ static void scx_proxy_reject_task(struct scx_sched *sch, struct rq *rq,
 	lockdep_assert_rq_held(rq);
 	WARN_ON_ONCE(p->scx.flags & SCX_TASK_REENQ_REASON_MASK);
 
-	p->scx.holding_cpu = -1;
 	p->scx.flags |= SCX_TASK_REENQ_PROXY;
 	scx_divert_strip_flags(p, &enq_flags);
 
@@ -2899,8 +2898,9 @@ static void scx_proxy_reject_task(struct scx_sched *sch, struct rq *rq,
  * tell whether we lost to dequeue by testing whether the holding_cpu still
  * points to this CPU. See scx_dispatch_dequeue() for the counterpart.
  *
- * On return, @dsq is unlocked and @src_rq is locked. Returns %true if @p is
- * still valid and detached from @dsq. %false if lost to dequeue.
+ * On return, @dsq is unlocked and @src_rq is locked. Return %true if @p is
+ * still valid and detached from @dsq, with holding_cpu cleared. Return %false
+ * if @p was lost to a dequeue.
  */
 static bool unlink_dsq_and_switch_rq_lock(struct task_struct *p,
 					  struct scx_dispatch_q *dsq,
@@ -2923,11 +2923,13 @@ static bool unlink_dsq_and_switch_rq_lock(struct task_struct *p,
 	if (likely(p->scx.holding_cpu == cpu) && !WARN_ON_ONCE(src_rq != task_rq(p))) {
 		/*
 		 * Keep ->dsq set until we own @src_rq so that a racing dequeue
-		 * can find @dsq and clear holding_cpu. Once ownership is
-		 * confirmed, clear it under @src_rq so that deactivate_task()
-		 * takes the !dsq path instead of retaking @dsq->lock.
+		 * can find @dsq and clear holding_cpu. Once we own @src_rq, the
+		 * handshake is over. Clear ->dsq so that deactivate_task()
+		 * takes the !dsq path instead of retaking @dsq->lock. Clear
+		 * holding_cpu with it so that no caller has to.
 		 */
 		p->scx.dsq = NULL;
+		p->scx.holding_cpu = -1;
 		scx_wait_dispatching(p);
 		return true;
 	}
@@ -3163,20 +3165,21 @@ static void dispatch_to_local_dsq(struct scx_sched *sch, struct rq *rq,
 	    !WARN_ON_ONCE(src_rq != task_rq(p))) {
 		bool fallback = false;
 
+		/* the handshake is over, see unlink_dsq_and_switch_rq_lock() */
+		p->scx.holding_cpu = -1;
+
 		/*
 		 * If @p is staying on the same rq, there's no need to go
 		 * through the full deactivate/activate cycle. Optimize by
 		 * abbreviating move_remote_task_to_local_dsq().
 		 */
 		if (src_rq == dst_rq) {
-			p->scx.holding_cpu = -1;
 			scx_dispatch_enqueue(sch, dst_rq, &dst_rq->scx.local_dsq, p,
 					     slice, vtime, enq_flags | SCX_ENQ_APPLY_SLICE);
 		} else if (unlikely(task_proxy_unsafe_to_move(p))) {
 			fallback = true;
 			scx_proxy_reject_task(sch, src_rq, p, enq_flags);
 		} else if (unlikely(!task_can_run_on_remote_rq(sch, p, dst_rq, true))) {
-			p->scx.holding_cpu = -1;
 			fallback = true;
 			scx_dispatch_enqueue(sch, src_rq, find_global_dsq(sch, task_cpu(p)),
 					     p, slice, vtime,
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index 80c746054ce7..752a9a2a94b0 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -2157,12 +2157,13 @@ static const char *scx_enable_state_str[] = {
  *   4. Lock src_rq (where the task currently lives)
  *   5. Verify p->scx.holding_cpu == current CPU, if not, dequeue won the
  *      race (dequeue clears holding_cpu to -1 when it takes the task), in
- *      this case migration is aborted
- *   6. If src_rq == dst_rq: clear holding_cpu and enqueue directly
- *      into dst_rq's local DSQ (no lock swap needed)
+ *      this case migration is aborted. Otherwise the claim is confirmed and
+ *      holding_cpu is cleared here, see unlink_dsq_and_switch_rq_lock()
+ *   6. If src_rq == dst_rq: enqueue directly into dst_rq's local DSQ (no lock
+ *      swap needed)
  *   7. Otherwise, verify under src_rq lock that the task can be moved to dst_rq
- *      (CPU affinity, migration_disabled, etc.). If not, clear holding_cpu,
- *      leave the task on src_rq, and enqueue it on the fallback DSQ.
+ *      (CPU affinity, migration_disabled, etc.). If not, leave the task on
+ *      src_rq, and enqueue it on the fallback DSQ.
  *   8. Otherwise (i.e. if the task can be moved to dst_rq), call
  *      move_remote_task_to_local_dsq(), which releases src_rq, locks dst_rq,
  *      and performs the deactivate/activate migration cycle
-- 
2.56.0

[PATCH 03/18] sched_ext: Fix bypassed tasks stranded by the reenqueue limit

[Tejun Heo]

scx_do_enqueue_task() counts a task's reenqueues until the task runs. Past
SCX_REENQ_MAX_REPEAT it ejects the owning scheduler and leaves the task off
every DSQ for the scheduler's exit to requeue. The count comes before the
bypass test, so it also counts the reenqueues of a task whose scheduler is
bypassing.

Those reenqueues are not the scheduler's placements, and ejecting the
scheduler for them fixes nothing. Under bypass the kernel queues the
scheduler's tasks on the bypass DSQs. A host calling
scx_bpf_reenqueue_local() reenqueues a bypassed child's tasks on the cpu's
local DSQ along with its own, and the reject drain reenqueues a consumed
task that a cap check rejected.

In the following patches, the consume itself reenqueues a task it can't run
right away, and the bypass load balancer reenqueues a task at its home.

The count is reset only when the task runs or leaves sched_ext, so an
ejection can lead to a stranding the exit does not recover:

 1. the task's reenqueues pass the limit, which ejects its scheduler and
    strands the task
 2. the exit's scx_bypass() cycles the task through dequeue and enqueue,
    which requeues it on a bypass DSQ with the count still past the limit
 3. the task's next reenqueue under bypass strands it again
 4. scx_bypass() requeues stranded tasks only when the scheduler wasn't
    bypassing yet, so the task stays off every DSQ until the disable walks
    the tasks
 5. the disable takes scx_enable_mutex, the cgroup lock and scx_fork_rwsem
    before the walk, and a sub-scheduler's disable the cpu hotplug lock too.
    It hangs if the stranded task holds one of them, forking, attaching to a
    cgroup or doing a cpu hotplug

Test bypass before the count. A task reenqueued under bypass stays queued
until a later dispatch or balancer pass, so there is no tight loop to bound,
and the task's first run resets the count. The offline test moves up with it
to stay ahead of the bypass test. A reenqueue on an offline cpu goes to the
local DSQ and isn't counted either.

A host that keeps calling scx_bpf_reenqueue_local() could get a bypassed
child ejected, as the child's tasks on the local DSQ counted against the
child. The reenqueues are the host's, and its own tasks still count against
it.

v2: The description lists the cpu hotplug lock among the locks a
sub-scheduler's disable takes before the walk.

Fixes: 7706d6e4f2e3 ("sched_ext: Bound per-task reenqueues and eject the owning scheduler")
Reported-by: Sashiko <sashiko-bot@kernel.org>
Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c | 39 ++++++++++++++++++++++++---------------
 1 file changed, 24 insertions(+), 15 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 0c0046766542..5d872a155800 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -2216,11 +2216,31 @@ void scx_do_enqueue_task(struct rq *rq, struct task_struct *p, u64 enq_flags,
 	 */
 	p->scx.flags &= ~SCX_TASK_IMMED;
 
+	/*
+	 * If !scx_rq_online(), we already told the BPF scheduler that the CPU
+	 * is offline and are just running the hotplug path. Don't bother the
+	 * BPF scheduler.
+	 */
+	if (!scx_rq_online(rq))
+		goto local;
+
+	if (scx_bypassing(sch, cpu_of(rq))) {
+		__scx_add_event(sch, SCX_EV_BYPASS_DISPATCH, 1);
+		goto bypass;
+	}
+
 	/*
 	 * A task reenqueued too many times without running means the scheduler
 	 * keeps re-deciding a placement it can't honor, e.g. re-inserting to a
 	 * cid it lacks caps on. Eject the owning scheduler and strand the task
 	 * to be picked up during sched exit.
+	 *
+	 * Bypass guarantees that every runnable task keeps running, and the
+	 * disable relies on that to take sleeping locks. The kernel places the
+	 * tasks of a bypassing scheduler and those on an offline rq, and
+	 * counting their reenqueues could strand a task under bypass and
+	 * deadlock the disable on a lock the task holds, so both tests stay
+	 * above the count.
 	 */
 	if (enq_flags & SCX_ENQ_REENQ) {
 		if (++p->scx.reenq_cnt > 1)
@@ -2234,19 +2254,6 @@ void scx_do_enqueue_task(struct rq *rq, struct task_struct *p, u64 enq_flags,
 		}
 	}
 
-	/*
-	 * If !scx_rq_online(), we already told the BPF scheduler that the CPU
-	 * is offline and are just running the hotplug path. Don't bother the
-	 * BPF scheduler.
-	 */
-	if (!scx_rq_online(rq))
-		goto local;
-
-	if (scx_bypassing(sch, cpu_of(rq))) {
-		__scx_add_event(sch, SCX_EV_BYPASS_DISPATCH, 1);
-		goto bypass;
-	}
-
 	if (p->scx.ddsp_dsq_id != SCX_DSQ_INVALID)
 		goto direct;
 
@@ -4993,7 +5000,8 @@ static void process_ddsp_deferred_locals(struct rq *rq)
  * %SCX_TASK_REENQ_IMMED. If the BPF scheduler dispatches back to the same local
  * DSQ with %SCX_ENQ_IMMED while the CPU is still unavailable, this triggers
  * another reenq cycle. Repetitions are bounded by %SCX_REENQ_MAX_REPEAT in
- * scx_do_enqueue_task(), which ejects the task's owning scheduler.
+ * scx_do_enqueue_task(), which ejects the task's owning scheduler. A bypassing
+ * scheduler's tasks are placed by the kernel instead and are not counted.
  */
 static bool local_task_should_reenq(struct rq *rq, struct task_struct *p,
 				    u64 *reenq_flags, u32 *reason)
@@ -5288,7 +5296,8 @@ static void process_deferred_reenq_users(struct rq *rq)
  *
  * A task can be re-rejected repeatedly. Reenqueues are bounded per task by
  * SCX_REENQ_MAX_REPEAT in scx_do_enqueue_task(), which ejects the owning
- * scheduler.
+ * scheduler. A bypassing scheduler's tasks are placed by the kernel instead and
+ * are not counted.
  */
 static void scx_reenq_reject(struct rq *rq)
 {
-- 
2.56.0

[PATCH 04/18] sched_ext: Fix bypassed tasks looping on a host without SCX_CAP_ENQ

[Tejun Heo]

Bypass exists so that runnable tasks make progress without trusting the BPF
scheduler. A host with SCX_CAP_ENQ_IMMED (BASE) but not SCX_CAP_ENQ on a cpu
cannot run a single bypassed task there. The task loops instead:

 1. a bypassed task lands on the host's bypass DSQ for the cpu
 2. the cpu's dispatch descends into the host, which holds BASE there
 3. the host consumes the task into the local DSQ as a plain insert
 4. the insert needs ENQ and the host lacks it, so the task is rejected
 5. the reject drain reenqueues the task and it is back at step 1

A bypass consume puts a task on the cpu being dispatched for, to run right
away. That is the IMMED contract, and BASE is the cap it needs. Consume with
SCX_ENQ_IMMED. BASE is then the one cap a bypassed task ever needs. Dispatch
descends into a host only on a cpu where the host holds BASE, so a consume
never lacks it.

A consumed task is first on the local DSQ when dispatch enters the host with
the DSQ empty. That is always the case for the root and normally for a
parent dispatching into a child.

A consumed task that cannot run right away is reenqueued like any IMMED task
and goes back to the bypass DSQ. Waiting on the local DSQ would need ENQ.
The task cannot run right away when a higher class or a HEAD insert took the
cpu meanwhile, or when the parent queued its own task before dispatching
into the host.

Each such task waits on the bypass DSQ for the next dispatch. The reenqueue
is not counted against the reenqueue limit, as the task's scheduler is
bypassing.

v2: The description says that the reenqueue of a consumed task is not
counted against the reenqueue limit (sashiko).

Fixes: 701b7bcad868 ("sched_ext: Add the SCX_CAP_ENQ cap")
Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/inlines.h | 15 ++++++++++++---
 kernel/sched/ext/sub.c     |  6 ++++++
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/kernel/sched/ext/inlines.h b/kernel/sched/ext/inlines.h
index 8de9d8fb4b25..1e39d12678c4 100644
--- a/kernel/sched/ext/inlines.h
+++ b/kernel/sched/ext/inlines.h
@@ -39,10 +39,18 @@ scx_dispatch_sched(struct scx_sched *sch, struct rq *rq,
 	if (scx_consume_global_dsq(sch, rq))
 		return SCX_DSP_LOCAL;
 
+	/*
+	 * A bypass consume puts a task on the cpu being dispatched for, to run
+	 * right away. That is an IMMED insert and needs only
+	 * %SCX_CAP_ENQ_IMMED, the base cap. Dispatch descends into @sch only on
+	 * a cpu where @sch holds the base cap, so @sch holds it here. The same
+	 * holds for the consume at the end of this function.
+	 */
 	if (scx_bypass_dsp_enabled(sch)) {
 		/* if @sch is bypassing, only the bypass DSQs are active */
 		if (scx_bypassing(sch, cpu)) {
-			if (scx_consume_dispatch_q(sch, rq, scx_bypass_dsq(sch, cpu), 0))
+			if (scx_consume_dispatch_q(sch, rq, scx_bypass_dsq(sch, cpu),
+						   SCX_ENQ_IMMED))
 				return SCX_DSP_LOCAL;
 			return SCX_DSP_NONE;
 		}
@@ -63,7 +71,8 @@ scx_dispatch_sched(struct scx_sched *sch, struct rq *rq,
 		struct scx_sched_pcpu *pcpu = per_cpu_ptr(sch->pcpu, cpu);
 
 		if (!(pcpu->bypass_host_seq++ % SCX_BYPASS_HOST_NTH) &&
-		    scx_consume_dispatch_q(sch, rq, scx_bypass_dsq(sch, cpu), 0)) {
+		    scx_consume_dispatch_q(sch, rq, scx_bypass_dsq(sch, cpu),
+					   SCX_ENQ_IMMED)) {
 			__scx_add_event(sch, SCX_EV_SUB_BYPASS_DISPATCH, 1);
 			return SCX_DSP_LOCAL;
 		}
@@ -135,7 +144,7 @@ scx_dispatch_sched(struct scx_sched *sch, struct rq *rq,
 	 * scheduler's ops.dispatch() doesn't yield any tasks.
 	 */
 	if (scx_bypass_dsp_enabled(sch) &&
-	    scx_consume_dispatch_q(sch, rq, scx_bypass_dsq(sch, cpu), 0)) {
+	    scx_consume_dispatch_q(sch, rq, scx_bypass_dsq(sch, cpu), SCX_ENQ_IMMED)) {
 		__scx_add_event(sch, SCX_EV_SUB_BYPASS_DISPATCH, 1);
 		return SCX_DSP_LOCAL;
 	}
diff --git a/kernel/sched/ext/sub.c b/kernel/sched/ext/sub.c
index f7f33fd0c5e9..a7b14ea16822 100644
--- a/kernel/sched/ext/sub.c
+++ b/kernel/sched/ext/sub.c
@@ -2722,6 +2722,12 @@ __bpf_kfunc_start_defs();
  * child schedulers. The child scheduler runs its dispatch operation to
  * move tasks from dispatch queues to the local runqueue.
  *
+ * Dispatch into a child while the local DSQ is empty, when the cpu needs a
+ * task. A child holding only %SCX_CAP_ENQ_IMMED can insert only IMMED tasks,
+ * and the consume of its bypassing descendants' bypass DSQs is one such insert.
+ * With a task already queued on the local DSQ, such an insert fails and the
+ * task is reenqueued.
+ *
  * Returns: true on success, false if cgroup_id is invalid, not a direct
  * child, or caller lacks dispatch permission.
  */
-- 
2.56.0

[PATCH 05/18] sched_ext: Move refill_task_slice_dfl() below apply_slice_vtime()

[Tejun Heo]

Move refill_task_slice_dfl() below apply_slice_vtime(), with the other slice
writers, to prepare for the kernel-slice handling the following patches add
to them. No functional change.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 5d872a155800..42a349aa2987 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -1569,6 +1569,16 @@ static void apply_slice_vtime(struct task_struct *p, u64 slice, u64 vtime, u64 e
 		p->scx.dsq_vtime = vtime;
 }
 
+static void refill_task_slice_dfl(struct scx_sched *sch, struct task_struct *p)
+{
+	/*
+	 * A default refill is not an explicit request, so it must not drop a
+	 * pending out-of-band one, which is applied when @p next runs.
+	 */
+	set_task_slice_keep_oob(p, READ_ONCE(sch->slice_dfl));
+	__scx_add_event(sch, SCX_EV_REFILL_SLICE_DFL, 1);
+}
+
 static void update_curr_scx(struct rq *rq)
 {
 	struct task_struct *donor;
@@ -1662,16 +1672,6 @@ static void dsq_dec_nr(struct scx_dispatch_q *dsq, struct task_struct *p)
 	}
 }
 
-static void refill_task_slice_dfl(struct scx_sched *sch, struct task_struct *p)
-{
-	/*
-	 * A default refill is not an explicit request, so it must not drop a
-	 * pending out-of-band one, which is applied when @p next runs.
-	 */
-	set_task_slice_keep_oob(p, READ_ONCE(sch->slice_dfl));
-	__scx_add_event(sch, SCX_EV_REFILL_SLICE_DFL, 1);
-}
-
 /*
  * Return true if @p is moving due to an internal SCX migration, false
  * otherwise.
-- 
2.56.0

[PATCH 06/18] sched_ext: Fix rescue served time disagreeing with the slice under proxy execution

[Tejun Heo]

Rescue measures served time as the rescuee's sum_exec_runtime minus a
snapshot taken at admission. Under proxy execution that measure follows the
wrong task. update_se() credits sum_exec_runtime to rq->curr, the mutex
owner that executes, while update_curr_scx() drains the slice of rq->donor
and charges its rescue. A rescuee can be either, so its served time and its
slice disagree in both directions.

As the owner running for another donor, the rescuee accrues served time that
drained no slice and charged no budget, and its rescue ends early. As the
donor, its rescue runs long:

 1. the rescuee blocks on a mutex and the owner runs in its place
 2. the owner's run drains the rescuee's slice and the budget, but not the
    served time
 3. the slice runs out with the rescue still open, and scx_rescue_keep()
    restores a remainder that the owner's run didn't reduce
 4. the owner keeps running on the rescue until it releases the mutex, past
    the admitted slice and overdrawing the budget

Count served time in scx_rescue_charge() from the delta that drains the
slice. Charging the owner's run to the donor's rescue matches proxy
execution, which drains the donor's slice and accounts the cgroup time to
the donor for the owner's run.

scx_rescue_keep() also ignored the return of scx_set_task_slice(). The call
refuses the restore on a rescuee that escalation marked SCX_TASK_PROTECTED,
and the rescuee was kept with a zero slice and its protection set.

That refusal is now unreachable. Escalation sets the slice to the unserved
remainder before it protects the task. From then on every slice write a
scheduler can reach is refused. Only a sub-scheduler's task is rescued, and
sub-schedulers are cid-form, which rejects direct slice writes.

The slice then moves only by what update_curr_scx() drains, and that drain
charges the rescue by the same amount, so the slice runs out exactly when
the charge ends the rescue. WARN and end the rescue if the restore is
refused anyway.

Fixes: 5fd501744b10 ("sched_ext: Add bandwidth-limited rescue execution for stranded tasks")
Reported-by: Sashiko <sashiko-bot@kernel.org>
Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/sub.c | 25 ++++++++++++++++---------
 kernel/sched/sched.h   |  2 +-
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/kernel/sched/ext/sub.c b/kernel/sched/ext/sub.c
index a7b14ea16822..998888ca5ea7 100644
--- a/kernel/sched/ext/sub.c
+++ b/kernel/sched/ext/sub.c
@@ -428,9 +428,7 @@ void scx_init_root_caps(struct scx_sched *sch)
 /* unserved remainder of @rq's rescuee's admitted slice, 0 once fully served */
 static s64 scx_rescue_slice_remaining(struct rq *rq)
 {
-	s64 served = rq->scx.rescue.curr->se.sum_exec_runtime - rq->scx.rescue.exec_snap;
-
-	return max(rq->scx.rescue.slice - served, 0);
+	return max(rq->scx.rescue.slice - rq->scx.rescue.served, 0);
 }
 
 /*
@@ -465,6 +463,14 @@ void scx_rescue_charge(struct rq *rq, s64 delta_exec)
 
 	lockdep_assert_rq_held(rq);
 
+	/*
+	 * Count @delta_exec as served before the clamp below, which bounds only
+	 * the budget charge. update_curr_scx() took the same delta off the
+	 * slice, so a slice set to the unserved remainder, as escalation does,
+	 * reaches zero when the remainder does, see scx_rescue_keep().
+	 */
+	rq->scx.rescue.served += delta_exec;
+
 	/*
 	 * A rescue slice is bounded by one quantum and tick-driven expiry can
 	 * overshoot by up to a tick. Clamp to avoid wild over-charges on VMs.
@@ -525,7 +531,9 @@ bool scx_rescue_keep(struct rq *rq, struct task_struct *p)
 	    scx_bypassing(scx_task_sched(p), cpu_of(rq)))
 		return false;
 
-	scx_set_task_slice(p, remaining);
+	/* never called on a protected rescuee, so the write can't be refused */
+	if (WARN_ON_ONCE(!scx_set_task_slice(p, remaining)))
+		return false;
 	return true;
 }
 
@@ -608,10 +616,9 @@ static void scx_rescue_timer_arm(struct rq *rq)
  * @slice: CPU time to grant
  *
  * The schedulers keep their normal control over @p and may preempt or reslice
- * it. @slice is measured on served CPU time against the snapshot taken here, so
- * neither shortens the rescue, see scx_rescue_charge() and scx_rescue_keep().
- * Prolonged denial escalates into protected execution, see
- * scx_rescue_timerfn().
+ * it. @slice is measured on served CPU time, so neither shortens the rescue,
+ * see scx_rescue_charge() and scx_rescue_keep(). Prolonged denial escalates
+ * into protected execution, see scx_rescue_timerfn().
  */
 static void scx_rescue_admit(struct rq *rq, struct task_struct *p, s64 slice)
 {
@@ -620,7 +627,7 @@ static void scx_rescue_admit(struct rq *rq, struct task_struct *p, s64 slice)
 
 	rq->scx.rescue.curr = p;
 	rq->scx.rescue.slice = slice;
-	rq->scx.rescue.exec_snap = p->se.sum_exec_runtime;
+	rq->scx.rescue.served = 0;
 	scx_set_task_slice(p, slice);
 	scx_rescue_timer_arm(rq);
 }
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 628765910934..dcc639998069 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -816,7 +816,7 @@ struct scx_rq_rescue {
 	u64			clock;			/* last budget accrual timestamp */
 	struct task_struct	*curr;			/* task being rescued, one at a time */
 	s64			slice;			/* curr's admitted slice */
-	u64			exec_snap;		/* sum_exec_runtime at admission */
+	s64			served;			/* curr's served time */
 	struct timer_list	timer;			/* paces admission and escalation */
 	u64			kill_at;		/* last ejection, init before any */
 };
-- 
2.56.0

[PATCH 07/18] sched_ext: Restore a rescuee's slice after ops.stopping()

[Tejun Heo]

put_prev_task_scx() restores a preempted rescuee's unserved slice through
scx_rescue_keep() before it calls ops.stopping(). A rescuee whose slice the
scheduler zeroes there is left to the reinsert instead:

 1. ops.stopping() zeroes the rescuee's slice
 2. the reinsert condition special-cases the rescuee and keeps it on the
    local DSQ with a zero slice
 3. the pick's default refill supplies a slice, or the task runs on a zero
    slice until the tick when its scheduler lacks BASE

Move the keep after the stopping call so that scx_rescue_keep() restores the
remainder whether a preemption or ops.stopping() zeroed the slice, and drop
the special case from the reinsert condition as a rescuee reaching it always
has a slice. A rescuee whose slice ops.stopping() zeroed now goes to the
tail of the local DSQ with its remainder restored, behind what the scheduler
wanted to run first, instead of staying at the head on a zero slice.

The slice end in the same block moves with it, which changes nothing for a
protected task. Its slice is the unserved remainder, and scx_rescue_charge()
ends the protection in the same update_curr_scx() call that zeroes the
slice, so the block never sees a protected task in either position.

v2: The description no longer says that a protected task's protection
outlasts ops.stopping().

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c | 29 +++++++++++++----------------
 1 file changed, 13 insertions(+), 16 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 42a349aa2987..406571afc9d8 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -3617,17 +3617,6 @@ static void put_prev_task_scx(struct rq *rq, struct task_struct *p,
 
 	update_curr_scx(rq);
 
-	/*
-	 * If the slice is consumed, protection ends with it. A rescuee
-	 * preempted beforehand keeps going, see scx_rescue_keep().
-	 */
-	if (!p->scx.slice) {
-		if (unlikely(p == scx_rescuee(rq)))
-			rescue_keep = scx_rescue_keep(rq, p);
-		if (!rescue_keep)
-			scx_task_slice_ended(rq, p);
-	}
-
 	/*
 	 * Preserve the running session when proxy execution refreshes the same
 	 * donor around an execution-context switch on this rq.
@@ -3640,6 +3629,18 @@ static void put_prev_task_scx(struct rq *rq, struct task_struct *p,
 		p->scx.flags &= ~SCX_TASK_RUN_TRACKED;
 	}
 
+	/*
+	 * If the slice is consumed, protection ends with it. A rescuee keeps
+	 * going when a preemption or ops.stopping() zeroed its slice, see
+	 * scx_rescue_keep().
+	 */
+	if (!p->scx.slice) {
+		if (unlikely(p == scx_rescuee(rq)))
+			rescue_keep = scx_rescue_keep(rq, p);
+		if (!rescue_keep)
+			scx_task_slice_ended(rq, p);
+	}
+
 	if (p->scx.flags & SCX_TASK_QUEUED) {
 		if (proxy_block)
 			goto switch_class;
@@ -3664,12 +3665,8 @@ static void put_prev_task_scx(struct rq *rq, struct task_struct *p,
 		 * The exception is an IMMED donor put by proxy_resched_idle():
 		 * the CPU isn't busy, proxy resolution is only retrying through
 		 * idle, so keep the donor where it can be picked again.
-		 *
-		 * An open rescue must keep @p on the local DSQ even if the
-		 * scheduler zeroed the slice in ops.stopping() above.
 		 */
-		if ((p->scx.slice || unlikely(p == scx_rescuee(rq))) &&
-		    !scx_bypassing(sch, cpu_of(rq))) {
+		if (p->scx.slice && !scx_bypassing(sch, cpu_of(rq))) {
 			if ((p->scx.flags & SCX_TASK_IMMED) && !proxy_put) {
 				p->scx.flags |= SCX_TASK_REENQ_PREEMPTED;
 				scx_do_enqueue_task(rq, p, SCX_ENQ_REENQ, -1);
-- 
2.56.0

[PATCH 08/18] sched_ext: Fix a cpu idling with a runnable rescuee

[Tejun Heo]

A scheduler preempts a running rescuee like any other task by zeroing its
slice. The rescue grant survives the preemption. put_prev_task_scx()
restores the unserved remainder through scx_rescue_keep() and queues the
rescuee at the tail of the local DSQ, behind whatever the scheduler wanted
to run first.

That breaks down when the scheduler zeroes the slice but dispatches nothing.
dispatch_one() keeps the previous task running only with slice left or, when
nothing else is runnable, with its scheduler holding SCX_CAP_BASE on the cpu
and not taking the last task through SCX_OPS_ENQ_LAST. A rescuee fails both
once its slice is zeroed if its scheduler lacks BASE there or takes the last
task:

 1. ops.tick() or scx_bpf_task_set_slice() zeroes the rescuee's slice
 2. the pick finds nothing else to run and settles on idle
 3. put_prev_task_scx() restores the remainder and queues the rescuee on the
    local DSQ
 4. nothing reschedules the cpu, as the pick has settled and __schedule()
    clears need_resched right after it

The cpu idles with the rescuee queued until an unrelated wakeup or the
rescue timer's escalation.

Restore the remainder from dispatch_one() too. When nothing else is
runnable, its last keep now keeps the rescuee running on the remainder, so
the pick never settles on idle with a rescuee runnable. put_prev_task_scx()
keeps handling a displaced rescuee, where another task won the pick.

A rescuee that passes the existing keep-last test is kept there instead and
continues as an ordinary task, its rescue ended with the zeroed slice. That
is fine as its scheduler holds SCX_CAP_BASE on the cpu.

The new keep also covers bypassed tasks run under rescue in the following
patches.

v2: The dispatch_one() comment and description now say that a rescuee whose
scheduler holds SCX_CAP_BASE on the cpu passes the keep-last test and
continues as an ordinary task (sashiko).

Fixes: 5fd501744b10 ("sched_ext: Add bandwidth-limited rescue execution for stranded tasks")
Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c | 38 +++++++++++++++++++++++++++++++++++---
 kernel/sched/ext/sub.c | 15 ++++++++-------
 2 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 406571afc9d8..c1361b61f09b 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -3397,6 +3397,34 @@ static enum scx_dsp_verdict dispatch_one(struct rq *rq, struct task_struct *prev
 			verdict = SCX_DSP_PREV;
 			goto has_tasks;
 		}
+
+		/*
+		 * The scheduler preempts a rescuee like any other task by
+		 * zeroing its slice. The grant survives the preemption. The
+		 * rescuee gets the unserved remainder back and runs again
+		 * behind whatever the scheduler wanted to run first.
+		 *
+		 * Two sites restore the remainder. When another task won the
+		 * pick, put_prev_task_scx() restores it and queues the rescuee
+		 * at the tail of the local DSQ.
+		 *
+		 * When nothing else is runnable, the cpu must not go idle with
+		 * the rescuee queued. The keep above takes a rescuee whose
+		 * scheduler holds %SCX_CAP_BASE here and doesn't take the last
+		 * task through %SCX_OPS_ENQ_LAST.
+		 *
+		 * The pick then ends its zeroed slice, and the rescue with it,
+		 * and refills a default slice. The task continues as an
+		 * ordinary one, which is fine as its scheduler holds BASE here.
+		 *
+		 * For any other rescuee, a reinsert from put_prev_task_scx()
+		 * would come after the pick settled on idle. Keep it running on
+		 * the remainder here.
+		 */
+		if (unlikely(prev == scx_rescuee(rq)) && scx_rescue_keep(rq, prev)) {
+			verdict = SCX_DSP_PREV;
+			goto has_tasks;
+		}
 	}
 	rq->scx.flags &= ~SCX_RQ_IN_DISPATCH;
 	return SCX_DSP_NONE;
@@ -3630,9 +3658,13 @@ static void put_prev_task_scx(struct rq *rq, struct task_struct *p,
 	}
 
 	/*
-	 * If the slice is consumed, protection ends with it. A rescuee keeps
-	 * going when a preemption or ops.stopping() zeroed its slice, see
-	 * scx_rescue_keep().
+	 * If the slice is consumed, protection ends with it.
+	 *
+	 * A rescuee keeps going when a preemption or ops.stopping() zeroed its
+	 * slice. Here another task won the pick. scx_rescue_keep() restores the
+	 * unserved remainder and the reinsert below queues the rescuee at the
+	 * tail of the local DSQ. When nothing else is runnable, dispatch_one()
+	 * keeps the rescuee running on the remainder instead.
 	 */
 	if (!p->scx.slice) {
 		if (unlikely(p == scx_rescuee(rq)))
diff --git a/kernel/sched/ext/sub.c b/kernel/sched/ext/sub.c
index 998888ca5ea7..5fa9c660aaa6 100644
--- a/kernel/sched/ext/sub.c
+++ b/kernel/sched/ext/sub.c
@@ -507,16 +507,17 @@ void scx_rescue_end(struct rq *rq)
 }
 
 /**
- * scx_rescue_keep - Keep the rescue going for a preempted-out rescuee
+ * scx_rescue_keep - Keep the rescue going past a zeroed slice
  * @rq: rq @p is running on
- * @p: task under rescue whose slice is exhausted
+ * @p: task under rescue
  *
- * Called from put_prev_task_scx() to decide what an exhausted slice means for
- * the rescuee. scx_rescue_charge() ends the rescue the moment the admitted
+ * Decide what an exhausted slice means for the rescuee. dispatch_one() calls
+ * this when nothing else is runnable and put_prev_task_scx() when another task
+ * won the pick. scx_rescue_charge() ends the rescue the moment the admitted
  * slice is fully served, so arriving here with the rescue still open means @p
- * was preempted. Restore the unserved remainder and return %true - @p stays the
- * rescuee and the caller reinserts it at the tail of the local DSQ, behind
- * whatever preempted the rescuee.
+ * was preempted. Restore the unserved remainder and return %true. @p stays the
+ * rescuee. dispatch_one() keeps it running and put_prev_task_scx() reinserts it
+ * at the tail of the local DSQ, behind whatever preempted the rescuee.
  *
  * Return %false to end the rescue instead - the slice is already fully served,
  * @p is leaving the rq or bypass is dismantling rescues.
-- 
2.56.0

[PATCH 09/18] sched_ext: Mark the rescue grant as a kernel slice with SCX_TASK_KERNEL_SLICE

[Tejun Heo]

In the following patches, bypassed tasks run under rescue when their host
lacks SCX_CAP_BASE on the cpu. Under bypass the kernel zeroes the running
task's slice at every tick, as it does not trust the scheduler's slices.

A rescue grant has to survive that, or it is served a tick at a time with
the rescuee sent to the tail of the local DSQ each time. The generic slice
handling can't tell a kernel grant from a scheduler's slice.

Mark the grant as a kernel slice with SCX_TASK_KERNEL_SLICE, set at rescue
admission and cleared when the slice ends. The generic slice handling leaves
it alone. The default refill does nothing for it, and the three zeroings go
through the new expire_task_slice(), which skips it:

 - the bypass tick's zeroing of the running task
 - the eviction of a running task that can't stay on the cpu when bypass
   lifts
 - the cap-revoke eviction of a running task

Schedulers keep their control. Preemption and scx_bpf_task_set_slice() still
apply, and the rescue stays bounded by served time.

The sites that recognized the rescuee by identity, the reenqueue scan and
put_prev_task_scx()'s reinsert, test the flag instead.
scx_task_has_kernel_slice() hides the test behind the scx_has_subs() static
key, as only a sub-scheduler's task can carry the flag. The flag's entries
join the enum tables in tools/sched_ext.

No kernel slice reaches a zeroing until bypassed tasks run under rescue. The
cap-revoke eviction skipped the rescuee already, and bypass entry ends an
open rescue, so the bypass tick and the eviction at bypass lift see none.

The refill is the one behavior change. The pick refills a zero slice only
for a task that can stay on the cpu, so it concerns a rescuee whose
scheduler holds BASE there. A scheduler may zero such a queued rescuee's
slice, and the next pick used to refill it to the default.

Now the rescuee runs on the zero slice until the next tick. There
put_prev_task_scx() restores the remainder when another task is runnable.

When nothing else is runnable, dispatch_one() keeps the rescuee as the last
task. The pick then ends its zero slice, which ends the rescue, and refills
the default. A scheduler with SCX_OPS_ENQ_LAST skips the keep-last test.
scx_rescue_keep() restores the remainder there instead and the rescue goes
on.

v2:
- scx_rescue_admit()'s kerneldoc no longer mentions a served-time snapshot.
- The description covers a scheduler with SCX_OPS_ENQ_LAST, whose rescuee
  dispatch_one() keeps on the remainder.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 include/linux/sched/ext.h                     | 11 +++++
 kernel/sched/ext/ext.c                        | 45 ++++++++++++++-----
 kernel/sched/ext/sub.c                        | 10 +++--
 kernel/sched/ext/sub.h                        |  7 +++
 .../sched_ext/include/scx/enum_defs.autogen.h |  1 +
 .../sched_ext/include/scx/enums_abi.autogen.h |  1 +
 6 files changed, 60 insertions(+), 15 deletions(-)

diff --git a/include/linux/sched/ext.h b/include/linux/sched/ext.h
index b341c328948e..fcf85b87fdd4 100644
--- a/include/linux/sched/ext.h
+++ b/include/linux/sched/ext.h
@@ -106,6 +106,17 @@ enum scx_ent_flags {
 	SCX_TASK_PROTECTED	= 1 << 6, /* slice and DSQ head position protected */
 	SCX_TASK_RUN_TRACKED	= 1 << 7, /* task is in an ops.running()/stopping() session */
 
+	/*
+	 * The kernel owns the slice, currently the rescue grant. While set:
+	 *
+	 * - refill_task_slice_dfl() and expire_task_slice() leave the slice
+	 *   alone.
+	 * - Reenqueue scans skip the task.
+	 * - put_prev_task_scx() keeps a preempted task on the local DSQ.
+	 * - Scheduler slice writes and preemption still apply.
+	 */
+	SCX_TASK_KERNEL_SLICE	= 1 << 8, /* slice owned by the kernel */
+
 	/*
 	 * The multi-bit fields are packed at the top, below SCX_TASK_CURSOR, so
 	 * that the single-bit flags stay contiguous from bit 0.
diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index c1361b61f09b..27fc3c4c1b3e 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -1472,8 +1472,8 @@ bool scx_set_task_slice(struct task_struct *p, u64 slice)
  * @rq: rq @p is on
  * @p: task of interest
  *
- * End what rides on the slice - the protection, and the rescue if @p is being
- * rescued.
+ * End what rides on the slice - the protection, the kernel's ownership of the
+ * slice, and the rescue if @p is being rescued.
  *
  * A dequeue normally ends the slice too. The exception is a save/restore pair
  * on the running task. Attribute changes like renice cycle the task through
@@ -1485,7 +1485,7 @@ void scx_task_slice_ended(struct rq *rq, struct task_struct *p)
 {
 	lockdep_assert_rq_held(rq);
 
-	p->scx.flags &= ~SCX_TASK_PROTECTED;
+	p->scx.flags &= ~(SCX_TASK_PROTECTED | SCX_TASK_KERNEL_SLICE);
 	if (unlikely(p == scx_rescuee(rq)))
 		scx_rescue_end(rq);
 }
@@ -1569,8 +1569,16 @@ static void apply_slice_vtime(struct task_struct *p, u64 slice, u64 vtime, u64 e
 		p->scx.dsq_vtime = vtime;
 }
 
+/*
+ * refill_task_slice_dfl() supplies a slice when the scheduler left none and
+ * expire_task_slice() takes it away to force the task off the cpu. Neither
+ * touches a kernel slice, see %SCX_TASK_KERNEL_SLICE.
+ */
 static void refill_task_slice_dfl(struct scx_sched *sch, struct task_struct *p)
 {
+	if (scx_task_has_kernel_slice(p))
+		return;
+
 	/*
 	 * A default refill is not an explicit request, so it must not drop a
 	 * pending out-of-band one, which is applied when @p next runs.
@@ -1579,6 +1587,17 @@ static void refill_task_slice_dfl(struct scx_sched *sch, struct task_struct *p)
 	__scx_add_event(sch, SCX_EV_REFILL_SLICE_DFL, 1);
 }
 
+static void expire_task_slice(struct task_struct *p)
+{
+	/*
+	 * Unlike the refill above, an expiry supersedes a pending out-of-band
+	 * request. A kept request would be applied at the next
+	 * update_curr_scx() and undo the expiry.
+	 */
+	if (!scx_task_has_kernel_slice(p))
+		scx_set_task_slice(p, 0);
+}
+
 static void update_curr_scx(struct rq *rq)
 {
 	struct task_struct *donor;
@@ -3706,11 +3725,15 @@ static void put_prev_task_scx(struct rq *rq, struct task_struct *p,
 				u64 enq_flags = 0;
 
 				/*
-				 * Keep a preempted rescue going. If preempted
-				 * by another SCX task, append to the local DSQ,
-				 * see scx_rescue_keep().
+				 * The enqueue caps gate the scheduler's
+				 * placement. The kernel placed a kernel-slice
+				 * task here itself, so reinsert it without a
+				 * cap check, at the head like any preempted
+				 * task. When scx_rescue_keep() restored a slice
+				 * the scheduler zeroed, the scheduler wanted
+				 * something else to run first, so append.
 				 */
-				if (unlikely(p == scx_rescuee(rq))) {
+				if (unlikely(scx_task_has_kernel_slice(p))) {
 					enq_flags |= SCX_ENQ_IGNORE_CAPS;
 					if (!rescue_keep)
 						enq_flags |= SCX_ENQ_HEAD;
@@ -4417,7 +4440,7 @@ static void task_tick_scx(struct rq *rq, struct task_struct *donor, int queued)
 	 * management.
 	 */
 	if (scx_bypassing(sch, cpu_of(rq)))
-		scx_set_task_slice(donor, 0);
+		expire_task_slice(donor);
 	else if (SCX_HAS_OP(sch, tick))
 		SCX_CALL_OP_TASK(sch, tick, rq, donor);
 
@@ -5040,7 +5063,7 @@ static bool local_task_should_reenq(struct rq *rq, struct task_struct *p,
 	first = !(*reenq_flags & SCX_REENQ_TSR_NOT_FIRST);
 	*reenq_flags |= SCX_REENQ_TSR_NOT_FIRST;
 
-	if (unlikely((p->scx.flags & SCX_TASK_PROTECTED) || p == scx_rescuee(rq)))
+	if (unlikely((p->scx.flags & SCX_TASK_PROTECTED) || scx_task_has_kernel_slice(p)))
 		return false;
 
 	*reason = SCX_TASK_REENQ_KFUNC;
@@ -5143,7 +5166,7 @@ static u32 reenq_local(struct scx_sched *sch, struct rq *rq, u64 reenq_flags)
 	if ((reenq_flags & SCX_REENQ_CAP_REVOKE) &&
 	    rq->donor->sched_class == &ext_sched_class &&
 	    scx_task_reenq_on_cap_revoke(rq, rq->donor)) {
-		scx_set_task_slice(rq->donor, 0);
+		expire_task_slice(rq->donor);
 		resched_curr(rq);
 	}
 
@@ -6999,7 +7022,7 @@ void scx_bypass(struct scx_sched *sch, bool bypass)
 			if (bypass)
 				scx_task_slice_ended(rq, p);
 			else if (!scx_task_can_stay_on_cpu(rq, p))
-				scx_set_task_slice(p, 0);
+				expire_task_slice(p);
 		}
 
 		/*
diff --git a/kernel/sched/ext/sub.c b/kernel/sched/ext/sub.c
index 5fa9c660aaa6..139257f3411f 100644
--- a/kernel/sched/ext/sub.c
+++ b/kernel/sched/ext/sub.c
@@ -616,10 +616,11 @@ static void scx_rescue_timer_arm(struct rq *rq)
  * @p: task being admitted, off any DSQ
  * @slice: CPU time to grant
  *
- * The schedulers keep their normal control over @p and may preempt or reslice
- * it. @slice is measured on served CPU time, so neither shortens the rescue,
- * see scx_rescue_charge() and scx_rescue_keep(). Prolonged denial escalates
- * into protected execution, see scx_rescue_timerfn().
+ * The slice is a kernel slice, see %SCX_TASK_KERNEL_SLICE. The schedulers keep
+ * their normal control over @p and may preempt or reslice it. @slice is
+ * measured on served CPU time, so neither shortens the rescue, see
+ * scx_rescue_charge() and scx_rescue_keep(). Prolonged denial escalates into
+ * protected execution, see scx_rescue_timerfn().
  */
 static void scx_rescue_admit(struct rq *rq, struct task_struct *p, s64 slice)
 {
@@ -630,6 +631,7 @@ static void scx_rescue_admit(struct rq *rq, struct task_struct *p, s64 slice)
 	rq->scx.rescue.slice = slice;
 	rq->scx.rescue.served = 0;
 	scx_set_task_slice(p, slice);
+	p->scx.flags |= SCX_TASK_KERNEL_SLICE;
 	scx_rescue_timer_arm(rq);
 }
 
diff --git a/kernel/sched/ext/sub.h b/kernel/sched/ext/sub.h
index eb1dba6d51e2..de80a8af524c 100644
--- a/kernel/sched/ext/sub.h
+++ b/kernel/sched/ext/sub.h
@@ -254,12 +254,19 @@ static inline struct task_struct *scx_rescuee(struct rq *rq)
 	return rq->scx.rescue.curr;
 }
 
+/* does the kernel own @p's slice? see %SCX_TASK_KERNEL_SLICE */
+static inline bool scx_task_has_kernel_slice(const struct task_struct *p)
+{
+	return scx_has_subs() && (p->scx.flags & SCX_TASK_KERNEL_SLICE);
+}
+
 #else	/* CONFIG_EXT_SUB_SCHED */
 
 static inline u64 scx_missing_caps(struct scx_sched *sch, s32 cpu, u64 needed) { return 0; }
 static inline u64 scx_caps_for_preempt(struct scx_sched *sch, struct rq *rq, u64 enq_flags) { return 0; }
 static inline bool scx_task_can_stay_on_cpu(struct rq *rq, struct task_struct *p) { return true; }
 static inline struct task_struct *scx_rescuee(struct rq *rq) { return NULL; }
+static inline bool scx_task_has_kernel_slice(const struct task_struct *p) { return false; }
 
 #endif	/* CONFIG_EXT_SUB_SCHED */
 
diff --git a/tools/sched_ext/include/scx/enum_defs.autogen.h b/tools/sched_ext/include/scx/enum_defs.autogen.h
index 4fa82033b99c..362e8bb34499 100644
--- a/tools/sched_ext/include/scx/enum_defs.autogen.h
+++ b/tools/sched_ext/include/scx/enum_defs.autogen.h
@@ -108,6 +108,7 @@
 #define HAVE_SCX_TASK_IMMED
 #define HAVE_SCX_TASK_PROTECTED
 #define HAVE_SCX_TASK_RUN_TRACKED
+#define HAVE_SCX_TASK_KERNEL_SLICE
 #define HAVE_SCX_TASK_STATE_SHIFT
 #define HAVE_SCX_TASK_STATE_BITS
 #define HAVE_SCX_TASK_STATE_MASK
diff --git a/tools/sched_ext/include/scx/enums_abi.autogen.h b/tools/sched_ext/include/scx/enums_abi.autogen.h
index 05a460611c4d..ad8f56f581ad 100644
--- a/tools/sched_ext/include/scx/enums_abi.autogen.h
+++ b/tools/sched_ext/include/scx/enums_abi.autogen.h
@@ -120,6 +120,7 @@ static const struct __scx_enum_abi_val __scx_enum_abi_vals[]
 	{ "scx_ent_flags", "SCX_TASK_IMMED", 0x20LLU },
 	{ "scx_ent_flags", "SCX_TASK_PROTECTED", 0x40LLU },
 	{ "scx_ent_flags", "SCX_TASK_RUN_TRACKED", 0x80LLU },
+	{ "scx_ent_flags", "SCX_TASK_KERNEL_SLICE", 0x100LLU },
 	{ "scx_ent_flags", "SCX_TASK_STATE_SHIFT", 0x19LLU },
 	{ "scx_ent_flags", "SCX_TASK_STATE_BITS", 0x3LLU },
 	{ "scx_ent_flags", "SCX_TASK_STATE_MASK", 0xe000000LLU },
-- 
2.56.0

[PATCH 10/18] sched_ext: Refuse extending a kernel slice

[Tejun Heo]

A rescuee can keep running after its grant is served, on a cpu its scheduler
has lost. The write rules deny a slice extension only where the scheduler
lacks SCX_CAP_BASE on the cpu, and a scheduler holding BASE but not ENQ can
have a task under rescue there:

 1. the scheduler extends the rescuee's slice
 2. BASE is revoked. The cap-revoke eviction skips the rescuee.
 3. the grant is served and scx_rescue_charge() ends the rescue
 4. the task runs out the extension on the cpu its scheduler lost, outside
    the rescue's bandwidth

A scheduler shortens or zeroes a rescuee's slice to preempt it, and
scx_rescue_keep() resumes the grant afterwards. An extension serves no such
purpose.

Refuse extending a kernel slice in both scheduler write paths, the direct
write in scx_bpf_task_set_slice() and the out-of-band application, counted
in SCX_EV_SLICE_DENIED like an extension without BASE. Shortening stays
allowed.

scx_rescue_keep()'s restore of the remainder raises the slice from zero.
Like the kernel's other slice writes, it goes through scx_set_task_slice()
and not the scheduler write paths.

v2:
- scx_rescue_admit()'s kerneldoc no longer mentions a served-time snapshot.
- SCX_EV_SLICE_DENIED's comment lists every refusal it counts.
- scx_bpf_task_set_slice()'s kerneldoc says that shortening is refused on a
  protected slice.

Fixes: 5fd501744b10 ("sched_ext: Add bandwidth-limited rescue execution for stranded tasks")
Signed-off-by: Tejun Heo <tj@kernel.org>
---
 include/linux/sched/ext.h   |  2 +-
 kernel/sched/ext/ext.c      | 30 +++++++++++++++++++-----------
 kernel/sched/ext/internal.h |  6 ++++--
 kernel/sched/ext/sub.c      | 10 +++++-----
 4 files changed, 29 insertions(+), 19 deletions(-)

diff --git a/include/linux/sched/ext.h b/include/linux/sched/ext.h
index fcf85b87fdd4..f3564581c317 100644
--- a/include/linux/sched/ext.h
+++ b/include/linux/sched/ext.h
@@ -113,7 +113,7 @@ enum scx_ent_flags {
 	 *   alone.
 	 * - Reenqueue scans skip the task.
 	 * - put_prev_task_scx() keeps a preempted task on the local DSQ.
-	 * - Scheduler slice writes and preemption still apply.
+	 * - Schedulers may shorten the slice, never extend it.
 	 */
 	SCX_TASK_KERNEL_SLICE	= 1 << 8, /* slice owned by the kernel */
 
diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 27fc3c4c1b3e..5308c7443b44 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -1385,7 +1385,8 @@ enum scx_slice_oob_consts {
  * p->scx.slice_oob, applied under @p's rq lock at the next slice consideration.
  *
  * While %SCX_TASK_PROTECTED is set, every scheduler-reachable slice update is
- * refused. See set_task_slice_keep_oob().
+ * refused. See set_task_slice_keep_oob(). A kernel slice, see
+ * %SCX_TASK_KERNEL_SLICE, can be shortened but not extended.
  *
  * dsq_vtime orders the next PRIQ insertion and has no running-side consumer, so
  * scx_bpf_task_set_dsq_vtime() writes it directly. Fork-time init and direct
@@ -1511,9 +1512,9 @@ static void set_task_slice_oob(struct scx_sched *sch, struct task_struct *p, u64
 /*
  * Apply a pending out-of-band slice request under @rq's lock. A request whose
  * packed id no longer matches @p's current owner is dropped. An extension needs
- * baseline cpu access on @p's cid, shortening is always allowed, and a
- * protected slice refuses both. %SCX_EV_SLICE_DENIED counts the denials. See
- * the write rules above.
+ * baseline cpu access on @p's cid and is refused on a kernel slice. Shortening
+ * is always allowed. A protected slice refuses both. %SCX_EV_SLICE_DENIED
+ * counts the denials. See the write rules above.
  */
 static void apply_task_slice_oob(struct rq *rq, struct task_struct *p)
 {
@@ -1537,7 +1538,8 @@ static void apply_task_slice_oob(struct rq *rq, struct task_struct *p)
 	slice = dur == SCX_SLICE_OOB_DUR_MASK ? SCX_SLICE_INF : dur;
 
 	if (slice > p->scx.slice &&
-	    unlikely(scx_missing_caps(scx_task_sched(p), cpu_of(rq), SCX_CAP_BASE))) {
+	    unlikely(scx_task_has_kernel_slice(p) ||
+		     scx_missing_caps(scx_task_sched(p), cpu_of(rq), SCX_CAP_BASE))) {
 		__scx_add_event(scx_task_sched(p), SCX_EV_SLICE_DENIED, 1);
 		return;
 	}
@@ -10427,10 +10429,12 @@ __bpf_kfunc_start_defs();
  * rules above for the ownership model.
  *
  * Raising the slice is honored only while the scheduler holds %SCX_CAP_BASE on
- * @p's cpu, otherwise it is counted in %SCX_EV_SLICE_DENIED. Shortening is
- * always allowed. On the stashed path the slice is packed into an atomic64_t
- * with the scheduler id and a flag bit, so a slice too large to fit is clamped
- * and counted in %SCX_EV_SLICE_CLAMPED. %SCX_SLICE_INF is preserved.
+ * @p's cpu, and never on a kernel slice, see %SCX_TASK_KERNEL_SLICE. Shortening
+ * is honored unless @p is protected, see %SCX_TASK_PROTECTED, which refuses
+ * every write. A refused write is counted in %SCX_EV_SLICE_DENIED. On the
+ * stashed path the slice is packed into an atomic64_t with the scheduler id and
+ * a flag bit, so a slice too large to fit is clamped and counted in
+ * %SCX_EV_SLICE_CLAMPED. %SCX_SLICE_INF is preserved.
  *
  * Return %true on success, %false if @p is not on the calling scheduler.
  */
@@ -10468,9 +10472,13 @@ __bpf_kfunc bool scx_bpf_task_set_slice(struct task_struct *p, u64 slice,
 		return true;
 	}
 
-	/* under the rq lock: apply now, extensions gated on baseline access */
+	/*
+	 * @p's rq is locked, so apply now. An extension needs baseline access
+	 * and is refused on a kernel slice.
+	 */
 	if (slice > p->scx.slice &&
-	    unlikely(scx_missing_caps(sch, cpu_of(locked_rq), SCX_CAP_BASE))) {
+	    unlikely(scx_task_has_kernel_slice(p) ||
+		     scx_missing_caps(sch, cpu_of(locked_rq), SCX_CAP_BASE))) {
 		__scx_add_event(sch, SCX_EV_SLICE_DENIED, 1);
 		return true;
 	}
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index 752a9a2a94b0..ed63e29c9442 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -1343,8 +1343,10 @@ struct scx_event_stats {
 	s64		SCX_EV_SLICE_CLAMPED;
 
 	/*
-	 * The number of times a slice extension was denied because the
-	 * scheduler lacked baseline cpu access on the task's cpu.
+	 * The number of times a slice write was refused. A raise is refused
+	 * without baseline cpu access on the task's cpu and on a kernel slice.
+	 * Every write is refused on a protected slice, a preemption's zeroing
+	 * included.
 	 */
 	s64		SCX_EV_SLICE_DENIED;
 
diff --git a/kernel/sched/ext/sub.c b/kernel/sched/ext/sub.c
index 139257f3411f..6fa8cea1ef88 100644
--- a/kernel/sched/ext/sub.c
+++ b/kernel/sched/ext/sub.c
@@ -616,11 +616,11 @@ static void scx_rescue_timer_arm(struct rq *rq)
  * @p: task being admitted, off any DSQ
  * @slice: CPU time to grant
  *
- * The slice is a kernel slice, see %SCX_TASK_KERNEL_SLICE. The schedulers keep
- * their normal control over @p and may preempt or reslice it. @slice is
- * measured on served CPU time, so neither shortens the rescue, see
- * scx_rescue_charge() and scx_rescue_keep(). Prolonged denial escalates into
- * protected execution, see scx_rescue_timerfn().
+ * The slice is a kernel slice, see %SCX_TASK_KERNEL_SLICE. The schedulers may
+ * preempt @p or shorten its slice, never extend it. @slice is measured on
+ * served CPU time, so neither shortens the rescue, see scx_rescue_charge() and
+ * scx_rescue_keep(). Prolonged denial escalates into protected execution, see
+ * scx_rescue_timerfn().
  */
 static void scx_rescue_admit(struct rq *rq, struct task_struct *p, s64 slice)
 {
-- 
2.56.0

[PATCH 11/18] sched_ext: Fix a cpu idling after put_prev_task_scx() queues its task on the local DSQ

[Tejun Heo]

put_prev_task_scx() enqueues a still-runnable task when the pick has settled
on idle. The insert can land the task on the cpu's own local DSQ, and then
the cpu must not go idle over it. 7a9df7667963 ("sched_ext: Fix stall when a
task enqueued with SCX_ENQ_LAST lands back on its CPU's local DSQ")
reschedules the idle task in that case, but only after the SCX_ENQ_LAST
insert.

The plain insert, taken when the task can't stay on the cpu, can land on the
local DSQ too: rescue can admit a task there when its scheduler lacks
SCX_CAP_BASE on the cpu. Without the resched the cpu idles with the task
queued until an unrelated wakeup or the rescue timer's escalation.

Reschedule after either insert. In the following patches, the plain insert
becomes the common path for bypassed tasks under rescue.

v2: The description's forward reference to the plain insert becoming the
common path no longer names the following patch.

Fixes: 5fd501744b10 ("sched_ext: Add bandwidth-limited rescue execution for stranded tasks")
Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c | 31 ++++++++++++++++++-------------
 1 file changed, 18 insertions(+), 13 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 5308c7443b44..9e503510ba1c 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -3658,6 +3658,8 @@ static void put_prev_task_scx(struct rq *rq, struct task_struct *p,
 			 (rq->scx.flags & SCX_RQ_PROXY_PENDING);
 	bool proxy_block = rq->scx.flags & SCX_RQ_PROXY_BLOCKING;
 	bool rescue_keep = false;
+	bool going_idle = !proxy_put && next &&
+			  sched_class_above(&ext_sched_class, next->sched_class);
 
 	rq->scx.flags &= ~SCX_RQ_PROXY_PENDING;
 
@@ -3763,25 +3765,28 @@ static void put_prev_task_scx(struct rq *rq, struct task_struct *p,
 		 * locally runnable and can legitimately go idle with @p still
 		 * runnable (see do_pick_task_scx()).
 		 */
-		if (!proxy_put && next &&
-		    sched_class_above(&ext_sched_class, next->sched_class) &&
-		    scx_task_can_stay_on_cpu(rq, p)) {
+		if (going_idle && scx_task_can_stay_on_cpu(rq, p)) {
 			WARN_ON_ONCE(!sched_core_enabled(rq) &&
 				     !(sch->ops.flags & SCX_OPS_ENQ_LAST));
 			scx_do_enqueue_task(rq, p, SCX_ENQ_LAST, -1);
-
-			/*
-			 * A task put on the local DSQ runs, as anywhere else.
-			 * Here the insert can't reschedule on its own: the pick
-			 * has settled on @next and @p is still curr, so
-			 * resched_curr() would flag @p and __schedule() clears
-			 * that right after the pick. Flag @next instead.
-			 */
-			if (p->scx.dsq == &rq->scx.local_dsq)
-				set_tsk_need_resched(next);
 		} else {
 			scx_do_enqueue_task(rq, p, 0, -1);
 		}
+
+		/*
+		 * @p may be on the local DSQ now. With %SCX_ENQ_LAST the
+		 * scheduler can put the last runnable task there. The plain
+		 * insert runs when @p cannot stay on this cpu, and rescue can
+		 * still admit it here.
+		 *
+		 * The cpu must not go idle with @p queued on its local DSQ. The
+		 * usual resched doesn't work: the pick has settled on @next and
+		 * @p is still curr, so resched_curr() would set need_resched on
+		 * @p, which __schedule() clears right after the pick. Set it on
+		 * @next instead.
+		 */
+		if (going_idle && p->scx.dsq == &rq->scx.local_dsq)
+			set_tsk_need_resched(next);
 	}
 
 switch_class:
-- 
2.56.0

[PATCH 12/18] sched_ext: Expire prev's slice in dispatch_one() under bypass

[Tejun Heo]

The next patch keeps a bypassed rescuee on the local DSQ and needs
dispatch_one() to keep it running too. dispatch_one()'s keep test refuses
any prev whose scheduler is bypassing, so another local task could run ahead
of the rescuee, even an escalated SCX_TASK_PROTECTED one. The test exists
because a cpu-form scheduler can write p->scx.slice directly from BPF, so
the slice can't be trusted under bypass.

The refusal doesn't hold anyway. Since 025b1bd41965 ("sched_ext: Implement
hierarchical bypass mode"), the host's ops.dispatch() runs while prev's
scheduler bypasses, and the post-dispatch test in scx_dispatch_sched() keeps
a queued prev with slice left without checking bypass. A task dispatched in
between waits behind prev until the next bypass tick zeroes the slice, at
most a tick.

Expire the slice before the keep test instead when prev's scheduler is
bypassing, as the bypass tick would. Neither test keeps prev on a zeroed
slice, so prev rotates. A kernel slice survives the expiry and keeps prev.
Only a sub-scheduler's task can hold one, and sub-schedulers are cid-form,
which rejects direct slice writes, so it can be trusted.

Drop the stale pointer to scx_disable_workfn(), update scx_bypass()'s
kerneldoc, and add a TODO noting that the bypass slice handling can be
cleaned up once the cpu form is deprecated and the kfuncs refuse slice
raises from a bypassing scheduler.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c | 32 ++++++++++++++++++++++++--------
 1 file changed, 24 insertions(+), 8 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 9e503510ba1c..dafb40c75cbc 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -3376,18 +3376,33 @@ static enum scx_dsp_verdict dispatch_one(struct rq *rq, struct task_struct *prev
 	if (prev->sched_class == &ext_sched_class) {
 		update_curr_scx(rq);
 
+		/*
+		 * Under bypass, @prev's slice can't be trusted as a cpu-form
+		 * scheduler can write it directly from BPF. Skipping the keep
+		 * below is not enough as the post-dispatch test in
+		 * scx_dispatch_sched() would keep @prev on the same slice.
+		 *
+		 * Expire it as the bypass tick would, so that neither test
+		 * keeps @prev on it and @prev rotates through the bypass
+		 * enqueue. A kernel slice is left alone and the keep below
+		 * keeps @prev on it, see %SCX_TASK_KERNEL_SLICE.
+		 *
+		 * TODO: Once the cpu form is deprecated and the kfuncs refuse
+		 * slice raises from a bypassing scheduler, slices can be
+		 * trusted under bypass. This expiry, the bypass tick's zeroing
+		 * and the related slice handling can then be cleaned up.
+		 */
+		if (scx_bypassing(scx_task_sched(prev), cpu))
+			expire_task_slice(prev);
+
 		/*
 		 * If @prev is runnable & has slice left, it has priority and
 		 * fetching more just increases latency for the fetched tasks.
 		 * Tell pick_task_scx() to keep running @prev. If the BPF
 		 * scheduler wants to handle this explicitly, it should
 		 * implement ->cpu_release().
-		 *
-		 * See scx_disable_workfn() for the explanation on the bypassing
-		 * test.
 		 */
-		if ((prev->scx.flags & SCX_TASK_QUEUED) && prev->scx.slice &&
-		    !scx_bypassing(scx_task_sched(prev), cpu)) {
+		if ((prev->scx.flags & SCX_TASK_QUEUED) && prev->scx.slice) {
 			verdict = SCX_DSP_PREV;
 			goto has_tasks;
 		}
@@ -6917,9 +6932,10 @@ static void unbypass_renotify_idle(struct rq *rq, struct scx_sched *pos,
  * - ops.sub_ecaps_updated() and ops.sub_child_ecaps_updated() are still called.
  *   The bypassed tasks are on the bypass DSQs, out of their reach.
  *
- * - dispatch_one() does not report %SCX_DSP_PREV on non-zero slice as slice
- *   can't be trusted. Whenever a tick triggers, the running task is rotated to
- *   the tail of the queue.
+ * - dispatch_one() expires prev's slice before its keep test as the slice can't
+ *   be trusted. Each tick does the same and rotates the running task to the
+ *   tail of the queue. A kernel slice is left alone, see
+ *   %SCX_TASK_KERNEL_SLICE.
  *
  * - pick_next_task() suppresses zero slice warning.
  *
-- 
2.56.0

[PATCH 13/18] sched_ext: Fix bypassed tasks stalling when the host lacks SCX_CAP_BASE

[Tejun Heo]

A bypassed task stalls on a cpu where its host lacks SCX_CAP_BASE. The host,
the nearest ancestor that is not bypassing, runs a bypassing sub-scheduler's
tasks through its per-cpu bypass DSQs. A cpu's dispatch only descends into
schedulers holding BASE on it, so the host's bypass DSQ on such a cpu is
never consumed until the host regains the cap.

Resolve a bypassed task against its host at enqueue time. If the host lacks
BASE on the cpu, send the task to rescue on that cpu, which runs it at
bounded bandwidth without the host's dispatch having to reach it.

The enqueue does not look above the host for an ancestor that holds BASE.
That ancestor would run the task on cpus it never granted the host.

With rescue disabled the task goes to the host's bypass DSQ as before and
stalls there until the host regains BASE or the bypass ends, or until the
stall watchdog ejects the host, which it blames for a task stuck on its DSQ.
A task that must run on the cpu, migration-disabled or with a migration
pending, is admitted to the local DSQ instead, as any insert of such a task
is.

Bypass also used to end a rescue, sending a preempted rescuee to the bypass
DSQ with every other preempted task. Now that rescue is how a bypassed task
runs while its host lacks BASE, keep a preempted rescuee on the local DSQ
under bypass too. dispatch_one() also keeps such a rescuee running, as an
earlier patch's bypass expiry leaves a kernel slice alone. The grant stays
bounded by served time.

A task scx_fail_parent() punted to its bypassing parent at NONE can be a
rescuee too. switched_from_scx() ends its slice as it does for an enabled
task, so that a NONE rescuee leaving sched_ext while running doesn't leave
rq->scx.rescue.curr pointing at it.

While at it, the rescue test in __scx_resolve_local_dsq() goes through the
new scx_rescue_enabled() helper.

v2:
- The description notes that dispatch_one() keeps a bypassed rescuee too,
  through the new patch before this one (sashiko).
- scx_rescue_keep()'s kerneldoc drops dispatch_one() declining a rescuee
  with slice left, which the new patch removes.

v3:
- With rescue disabled, a task that must run on the cpu is admitted to the
  local DSQ instead of waiting on the host's bypass DSQ.
- The description names the stall watchdog as the way out for a task stuck
  without rescue.

v4:
- switched_from_scx() ends the slice of a NONE task, which can now be a
  rescuee (sashiko).
- scx_punt_task()'s kerneldoc no longer says that switched_from_scx() skips
  a NONE task.
- scx_bypass()'s kerneldoc and the ops.sub_ecaps_updated() doc say where a
  bypassed task can be besides the bypass DSQs.

Fixes: 147d1885f390 ("sched_ext: Add the SCX_CAP_ENQ_IMMED cap")
Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c      | 52 ++++++++++++++++++++++++++++++-------
 kernel/sched/ext/internal.h |  5 ++--
 kernel/sched/ext/sub.c      | 24 +++++++++--------
 kernel/sched/ext/sub.h      |  9 +++++++
 4 files changed, 68 insertions(+), 22 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index dafb40c75cbc..a75d3c3bfad9 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -388,8 +388,12 @@ static const struct sched_class *scx_setscheduler_class(struct task_struct *p)
 	return __setscheduler_class(p->policy, p->prio);
 }
 
-static struct scx_dispatch_q *bypass_enq_target_dsq(struct scx_sched *sch, s32 cpu)
+static struct scx_dispatch_q *bypass_enq_target_dsq(struct scx_sched *sch, struct rq *rq,
+						    struct task_struct *p, u64 *enq_flags)
 {
+	struct scx_sched *host = scx_exec_sched(sch, cpu_of(rq));
+	struct scx_dispatch_q *dsq;
+
 	/*
 	 * If @sch is a sub-sched which is bypassing, its tasks should go into
 	 * the bypass DSQs of the nearest ancestor which is not bypassing. The
@@ -400,8 +404,29 @@ static struct scx_dispatch_q *bypass_enq_target_dsq(struct scx_sched *sch, s32 c
 	 * Whenever a sched starts bypassing, all runnable tasks in its subtree
 	 * are re-enqueued after scx_bypassing() is turned on, guaranteeing that
 	 * all tasks are transferred to the right DSQs.
-	 */
-	return scx_bypass_dsq(scx_exec_sched(sch, cpu), cpu);
+	 *
+	 * A cpu's dispatch only descends into a host holding %SCX_CAP_BASE on
+	 * it, so the bypass DSQ of a host without BASE here is never consumed.
+	 * Send @p to rescue on this cpu instead, resolving against @host.
+	 * scx_dispatch_enqueue() resolves a local DSQ again, against the
+	 * bypassing @sch. %SCX_ENQ_IGNORE_CAPS makes that second resolution
+	 * admit @p.
+	 *
+	 * A migration-disabled task can't leave this cpu, and a task with a
+	 * migration pending must run here until the migration moves it. Either
+	 * would wait on the unconsumed bypass DSQ until the host regains BASE.
+	 * Resolve them too when rescue is disabled. The resolution admits both
+	 * to the local DSQ, as it does any insert of such a task.
+	 */
+	if (unlikely(scx_missing_caps(host, cpu_of(rq), SCX_CAP_BASE)) &&
+	    (scx_rescue_enabled() || is_migration_disabled(p) || p->migration_pending)) {
+		*enq_flags |= SCX_ENQ_RESCUE;
+		dsq = scx_resolve_local_dsq(host, rq, p, enq_flags);
+		*enq_flags |= SCX_ENQ_IGNORE_CAPS;
+	} else {
+		dsq = scx_bypass_dsq(host, cpu_of(rq));
+	}
+	return dsq;
 }
 
 /**
@@ -2353,7 +2378,7 @@ void scx_do_enqueue_task(struct rq *rq, struct task_struct *p, u64 enq_flags,
 	dsq = find_global_dsq(sch, task_cpu(p));
 	goto enqueue;
 bypass:
-	dsq = bypass_enq_target_dsq(sch, task_cpu(p));
+	dsq = bypass_enq_target_dsq(sch, rq, p, &enq_flags);
 	goto enqueue;
 
 enqueue:
@@ -3735,8 +3760,13 @@ static void put_prev_task_scx(struct rq *rq, struct task_struct *p,
 		 * The exception is an IMMED donor put by proxy_resched_idle():
 		 * the CPU isn't busy, proxy resolution is only retrying through
 		 * idle, so keep the donor where it can be picked again.
+		 *
+		 * A task with a kernel slice stays on the local DSQ even while
+		 * @sch bypasses. Under bypass, rescue runs @p when the host
+		 * lacks %SCX_CAP_BASE, see bypass_enq_target_dsq().
 		 */
-		if (p->scx.slice && !scx_bypassing(sch, cpu_of(rq))) {
+		if (p->scx.slice && (!scx_bypassing(sch, cpu_of(rq)) ||
+				     scx_task_has_kernel_slice(p))) {
 			if ((p->scx.flags & SCX_TASK_IMMED) && !proxy_put) {
 				p->scx.flags |= SCX_TASK_REENQ_PREEMPTED;
 				scx_do_enqueue_task(rq, p, SCX_ENQ_REENQ, -1);
@@ -4987,11 +5017,14 @@ static void switched_from_scx(struct rq *rq, struct task_struct *p)
 	 * parent's own teardown). There is nothing to disable; calling
 	 * scx_disable_task() would WARN on the non-%ENABLED state and trigger a
 	 * NONE -> READY validation failure.
+	 *
+	 * @p can still be a rescuee under the parent. End its slice as
+	 * scx_disable_task() does for an enabled task.
 	 */
 	if (scx_get_task_state(p) == SCX_TASK_NONE)
-		return;
-
-	scx_disable_task(scx_task_sched(p), p);
+		scx_task_slice_ended(rq, p);
+	else
+		scx_disable_task(scx_task_sched(p), p);
 }
 
 static void switched_to_scx(struct rq *rq, struct task_struct *p) {}
@@ -6930,7 +6963,8 @@ static void unbypass_renotify_idle(struct rq *rq, struct scx_sched *pos,
  * - ops.dispatch() is ignored.
  *
  * - ops.sub_ecaps_updated() and ops.sub_child_ecaps_updated() are still called.
- *   The bypassed tasks are on the bypass DSQs, out of their reach.
+ *   The bypassed tasks are on the bypass DSQs, under rescue or admitted to a
+ *   local DSQ, out of their reach.
  *
  * - dispatch_one() expires prev's slice before its keep test as the slice can't
  *   be trusted. Each tick does the same and rotates the running task to the
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index ed63e29c9442..ecd677c4b561 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -924,8 +924,9 @@ struct sched_ext_ops {
 	 * ops.sub_attach() and any change since are reported then. A cpu that
 	 * isn't online at that point reports them once it comes online.
 	 * Deliveries continue while this sched bypasses. Its tasks are then on
-	 * the bypass DSQs, where this op can't insert or move them. Deliveries
-	 * stop once this sched is being disabled.
+	 * the bypass DSQs, under rescue or admitted to a local DSQ, where this
+	 * op can't insert or move them. Deliveries stop once this sched is
+	 * being disabled.
 	 *
 	 * Until this sched is being disabled, whenever an op of this sched runs
 	 * under a cid's rq lock, its effective caps on the cid are the @after
diff --git a/kernel/sched/ext/sub.c b/kernel/sched/ext/sub.c
index 6fa8cea1ef88..01215384eaa5 100644
--- a/kernel/sched/ext/sub.c
+++ b/kernel/sched/ext/sub.c
@@ -36,7 +36,7 @@ static DEFINE_STATIC_KEY_FALSE(__scx_ops_cid_sched_updated_enabled);
 static s32 scx_nr_ops_cid_sched_updated;	/* such scheds, under scx_enable_mutex */
 
 /* latched at root enable before any rescue runs */
-static s32 scx_rescue_bw_1024;
+s32 scx_rescue_bw_1024;			/* 0 if disabled, see scx_rescue_enabled() */
 static s64 scx_rescue_quantum_ns;
 static s64 scx_rescue_sat_delta_ns;
 static unsigned long scx_rescue_decay_halflife;
@@ -519,8 +519,8 @@ void scx_rescue_end(struct rq *rq)
  * rescuee. dispatch_one() keeps it running and put_prev_task_scx() reinserts it
  * at the tail of the local DSQ, behind whatever preempted the rescuee.
  *
- * Return %false to end the rescue instead - the slice is already fully served,
- * @p is leaving the rq or bypass is dismantling rescues.
+ * Return %false to end the rescue instead - the slice is already fully served
+ * or @p is leaving the rq.
  */
 bool scx_rescue_keep(struct rq *rq, struct task_struct *p)
 {
@@ -528,8 +528,7 @@ bool scx_rescue_keep(struct rq *rq, struct task_struct *p)
 
 	lockdep_assert_rq_held(rq);
 
-	if (!remaining || !(p->scx.flags & SCX_TASK_QUEUED) ||
-	    scx_bypassing(scx_task_sched(p), cpu_of(rq)))
+	if (!remaining || !(p->scx.flags & SCX_TASK_QUEUED))
 		return false;
 
 	/* never called on a protected rescuee, so the write can't be refused */
@@ -887,9 +886,11 @@ void scx_rescue_init(struct rq *rq)
  * %SCX_ENQ_PREEMPT_LAZY are also cleared on a fallback migration-disabled
  * admission.
  *
- * Bypass doesn't need special-casing as a bypassing sched's tasks are enqueued
- * to and run by its nearest non-bypassing ancestor. If root is bypassing, it
- * always holds all caps.
+ * bypass_enq_target_dsq() also calls this with the host as @sch when the host
+ * lacks %SCX_CAP_BASE on @rq's cid, if rescue is enabled or @p must run on the
+ * cpu, migration-disabled or with a migration pending. The bypassed task is
+ * then admitted or diverted like a capless %SCX_ENQ_RESCUE insert by the host
+ * itself.
  */
 struct scx_dispatch_q *__scx_resolve_local_dsq(struct scx_sched *sch, struct rq *rq,
 					       struct task_struct *p, u64 *enq_flags)
@@ -932,7 +933,7 @@ struct scx_dispatch_q *__scx_resolve_local_dsq(struct scx_sched *sch, struct rq
 	scx_divert_strip_flags(p, enq_flags);
 
 	/* the enqueuer opted for rescue instead of rejection and reenqueue */
-	if ((*enq_flags & SCX_ENQ_RESCUE) && likely(scx_rescue_bw_1024)) {
+	if (scx_rescue_enabled() && (*enq_flags & SCX_ENQ_RESCUE)) {
 		__scx_add_event(sch, SCX_EV_SUB_RESCUE, 1);
 		if (scx_rescue_try_admit(rq, p))
 			return &rq->scx.local_dsq;
@@ -1656,8 +1657,9 @@ static void scx_rehome_task(struct scx_sched *to, struct task_struct *p)
  *
  * Used when @to must take over @p but failed to initialize it. Bypass keeps
  * scheduling decisions away from @to but @p can still trigger its task ops,
- * which may confuse the BPF side. @to is dying anyway. The exit paths skip
- * %NONE tasks (see __scx_disable_and_exit_task() and switched_from_scx()).
+ * which may confuse the BPF side. @to is dying anyway. The exit paths call no
+ * op on a %NONE task. __scx_disable_and_exit_task() skips it and
+ * switched_from_scx() only ends its slice.
  */
 static void scx_punt_task(struct scx_sched *to, struct task_struct *p)
 {
diff --git a/kernel/sched/ext/sub.h b/kernel/sched/ext/sub.h
index de80a8af524c..b6d4a522a5a3 100644
--- a/kernel/sched/ext/sub.h
+++ b/kernel/sched/ext/sub.h
@@ -14,6 +14,8 @@
 
 #ifdef CONFIG_EXT_SUB_SCHED
 
+extern s32 scx_rescue_bw_1024;
+
 struct scx_sched *scx_skip_subtree_pre(struct scx_sched *pos, struct scx_sched *root);
 struct scx_sched *scx_next_descendant_pre(struct scx_sched *pos, struct scx_sched *root);
 void scx_cid_sched_update(struct rq *rq, struct scx_sched *sch);
@@ -51,6 +53,12 @@ void scx_rescue_dump(struct seq_buf *s, struct rq *rq);
 void scx_rescue_set_knobs(struct scx_sched *sch);
 void scx_rescue_init(struct rq *rq);
 
+/* false only when the root disabled rescue, see scx_rescue_set_knobs() */
+static inline bool scx_rescue_enabled(void)
+{
+	return likely(scx_rescue_bw_1024);
+}
+
 /*
  * cgrp->scx_sched is written by root/sub enable/disable under all of
  * scx_enable_mutex, scx_fork_rwsem and cgroup_mutex. A new cgroup inherits the
@@ -128,6 +136,7 @@ static inline void scx_rescue_end(struct rq *rq) {}
 static inline bool scx_rescue_keep(struct rq *rq, struct task_struct *p) { return false; }
 static inline void scx_rescue_flush(struct rq *rq) {}
 static inline void scx_rescue_dump(struct seq_buf *s, struct rq *rq) {}
+static inline bool scx_rescue_enabled(void) { return false; }
 static inline void scx_rescue_set_knobs(struct scx_sched *sch) {}
 static inline void scx_rescue_init(struct rq *rq) {}
 static inline void scx_dec_has_subs(struct scx_sched *sch) {}
-- 
2.56.0

[PATCH 14/18] sched_ext: Fix the bypass load balancer stranding tasks on cpus where the host lacks SCX_CAP_BASE

[Tejun Heo]

The bypass load balancer spreads a host's bypassed tasks over the cpus of a
node by moving them between its per-cpu bypass DSQs. It counts every cpu of
the node as capacity and takes any cpu below the target as a donee. A cpu's
dispatch only descends into schedulers holding SCX_CAP_BASE on it, so a
bypass DSQ on a cpu where the host lacks BASE is never consumed. A task
moved there sits until the host regains the cap or the bypass window ends.

Count only cpus where the host holds BASE as capacity and pick donees from
them. The target rises to spread the node's tasks over the cpus that can run
them, and a capless cpu receives nothing. A capless cpu with tasks is still
an ordinary donor, offloaded down to the donor target when its queue grows
past the threshold; draining it is left to the following patches.

Fixes: 147d1885f390 ("sched_ext: Add the SCX_CAP_ENQ_IMMED cap")
Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index a75d3c3bfad9..ff79d9af0389 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -6741,7 +6741,9 @@ static void bypass_lb_node(struct scx_sched *sch, int node)
 		u32 nr = READ_ONCE(scx_bypass_dsq(sch, cpu)->nr);
 
 		nr_tasks += nr;
-		nr_cpus++;
+		/* only a cpu whose dispatch consumes its DSQ is capacity */
+		if (!scx_missing_caps(sch, cpu, SCX_CAP_BASE))
+			nr_cpus++;
 
 		before_min = min(nr, before_min);
 		before_max = max(nr, before_max);
@@ -6760,7 +6762,9 @@ static void bypass_lb_node(struct scx_sched *sch, int node)
 
 	cpumask_clear(donee_mask);
 	for_each_cpu_and(cpu, cpu_online_mask, node_mask) {
-		if (READ_ONCE(scx_bypass_dsq(sch, cpu)->nr) < nr_target)
+		/* a donee's DSQ is consumed only if @sch holds %SCX_CAP_BASE */
+		if (READ_ONCE(scx_bypass_dsq(sch, cpu)->nr) < nr_target &&
+		    !scx_missing_caps(sch, cpu, SCX_CAP_BASE))
 			cpumask_set_cpu(cpu, donee_mask);
 	}
 
-- 
2.56.0

[PATCH 15/18] sched_ext: Factor out bypass_lb_move() and insert with the donor DSQ unlocked

[Tejun Heo]

Prepare the bypass load balancer for moving tasks homed on other cpus. Such
a move takes the task's own rq lock through the remote consumer's handshake,
which drops the donor DSQ lock on the way, so its insert runs with the donor
DSQ unlocked and under the task's rq lock. The balancer's own move holds the
donor DSQ lock throughout, with the donee DSQ lock nested inside.

Switch the donor-homed move to the same scheme so that one body can serve
both: factor it out of bypass_lb_cpu() into bypass_lb_move(), which takes
the task off the donor DSQ and inserts with the DSQ unlocked. A task that
can't run on the donee goes back to the head of the donor DSQ instead of
being skipped in place.

The walk parks its cursor behind the task before each move and resumes from
it afterwards. It is bounded to the tasks queued when it began, by their DSQ
sequence numbers, so a task it put back is not visited again in the same
pass.

Before, the walk held both the donor rq lock and the donor DSQ lock across
every move. The donor DSQ can be very long and contended, a walk competing
with the cpu's dispatch and the enqueues for the DSQ lock would crawl, and
the walk is what relieves that contention by moving tasks off the DSQ. The
walk now cycles the DSQ lock around each insert but still holds the rq lock
throughout, and that alone keeps dispatch and enqueues away. The DSQ lock is
uncontended when retaken, so cycling it is cheap.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c | 131 ++++++++++++++++++++++++++++-------------
 1 file changed, 91 insertions(+), 40 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index ff79d9af0389..d8e0d99886c6 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -6619,6 +6619,57 @@ bool scx_hardlockup(int cpu)
 	return true;
 }
 
+/**
+ * bypass_lb_move - Move one task off a bypass DSQ for the balancer
+ * @sch: host sched whose bypass DSQs are balanced
+ * @donor_rq: rq of the cpu whose bypass DSQ gives up @p, locked
+ * @donor_dsq: @sch's bypass DSQ on @donor_rq's cpu, locked
+ * @p: task on @donor_dsq, homed on @donor_rq
+ * @donee: cpu whose bypass DSQ receives @p
+ *
+ * Take @p off @donor_dsq and insert it into the bypass DSQ of @donee. The
+ * caller picks @donee from @p's affinity alone and task_can_run_on_remote_rq()
+ * can still refuse the move under @p's rq lock. @p then goes back to the head
+ * of @donor_dsq.
+ *
+ * @donor_rq and @donor_dsq are locked on entry and on return. @donor_dsq is
+ * unlocked across the insert, so the caller's walk must find its place in it
+ * again after the call.
+ *
+ * Return the cpu that received @p, or -1 if @p went back onto @donor_dsq.
+ */
+static s32 bypass_lb_move(struct scx_sched *sch, struct rq *donor_rq,
+			  struct scx_dispatch_q *donor_dsq, struct task_struct *p, s32 donee)
+{
+	struct scx_dispatch_q *dst_dsq;
+	u64 enq_flags;
+	s32 ret;
+
+	lockdep_assert_rq_held(donor_rq);
+	lockdep_assert_held(&donor_dsq->lock);
+
+	scx_wait_dispatching(p);
+	dispatch_dequeue_locked(p, donor_dsq);
+	raw_spin_unlock(&donor_dsq->lock);
+
+	/* @p's affinity is stable under its rq lock */
+	if (task_can_run_on_remote_rq(sch, p, cpu_rq(donee), false)) {
+		dst_dsq = scx_bypass_dsq(sch, donee);
+		enq_flags = 0;
+		ret = donee;
+	} else {
+		/* back where it was, at the head so that the walk skips it */
+		dst_dsq = donor_dsq;
+		enq_flags = SCX_ENQ_HEAD;
+		ret = -1;
+	}
+
+	scx_dispatch_enqueue(sch, donor_rq, dst_dsq, p, 0, 0, enq_flags);
+
+	raw_spin_lock(&donor_dsq->lock);
+	return ret;
+}
+
 static u32 bypass_lb_cpu(struct scx_sched *sch, s32 donor,
 			 struct cpumask *donee_mask, struct cpumask *resched_mask,
 			 u32 nr_donor_target, u32 nr_donee_target)
@@ -6642,14 +6693,15 @@ static u32 bypass_lb_cpu(struct scx_sched *sch, s32 donor,
 
 	raw_spin_rq_lock_irq(donor_rq);
 	raw_spin_lock(&donor_dsq->lock);
+	/* bound the walk to the tasks queued now, see the seq test below */
+	cursor.priv = donor_dsq->seq;
 	list_add(&cursor.node, &donor_dsq->list);
 resume:
 	n = container_of(&cursor, struct task_struct, scx.dsq_list);
 	n = nldsq_next_task(donor_dsq, n, false);
 
 	while ((p = n)) {
-		struct scx_dispatch_q *donee_dsq;
-		int donee;
+		s32 donee, cpu;
 
 		n = nldsq_next_task(donor_dsq, n, false);
 
@@ -6659,6 +6711,13 @@ static u32 bypass_lb_cpu(struct scx_sched *sch, s32 donor,
 		if (cpumask_empty(donee_mask))
 			break;
 
+		/*
+		 * A task queued after the walk began, including one the walk
+		 * itself put back, is left for the next pass.
+		 */
+		if (u32_before(cursor.priv, p->scx.dsq_seq))
+			continue;
+
 		/*
 		 * If an earlier pass placed @p on @donor_dsq from a different
 		 * CPU and the donee hasn't consumed it yet, @p is still on the
@@ -6672,49 +6731,41 @@ static u32 bypass_lb_cpu(struct scx_sched *sch, s32 donor,
 		if (donee >= nr_cpu_ids)
 			continue;
 
-		donee_dsq = scx_bypass_dsq(sch, donee);
-
-		/*
-		 * $p's rq is not locked but $p's DSQ lock protects its
-		 * scheduling properties making this test safe.
-		 */
-		if (!task_can_run_on_remote_rq(sch, p, cpu_rq(donee), false))
-			continue;
-
-		/*
-		 * Moving $p from one non-local DSQ to another. The source rq
-		 * and DSQ are already locked. Do an abbreviated dequeue and
-		 * then perform enqueue without unlocking $donor_dsq.
-		 *
-		 * We don't want to drop and reacquire the lock on each
-		 * iteration as @donor_dsq can be very long and potentially
-		 * highly contended. Donee DSQs are less likely to be contended.
-		 * The nested locking is safe as only this LB moves tasks
-		 * between bypass DSQs.
-		 */
-		dispatch_dequeue_locked(p, donor_dsq);
-		scx_dispatch_enqueue(sch, cpu_rq(donee), donee_dsq, p, 0, 0, SCX_ENQ_NESTED);
-
 		/*
-		 * $donee might have been idle and need to be woken up. No need
-		 * to be clever. Kick every CPU that receives tasks.
+		 * The move drops @donor_dsq's lock. Park the cursor behind @p
+		 * and resume from it afterwards.
 		 */
-		cpumask_set_cpu(donee, resched_mask);
+		list_move(&cursor.node, &p->scx.dsq_list.node);
+		cpu = bypass_lb_move(sch, donor_rq, donor_dsq, p, donee);
+		if (cpu >= 0) {
+			/*
+			 * @cpu might have been idle and need to be woken up.
+			 * No need to be clever. Reschedule every cpu that
+			 * receives tasks.
+			 */
+			cpumask_set_cpu(cpu, resched_mask);
 
-		if (READ_ONCE(donee_dsq->nr) >= nr_donee_target)
-			cpumask_clear_cpu(donee, donee_mask);
+			if (READ_ONCE(scx_bypass_dsq(sch, cpu)->nr) >= nr_donee_target)
+				cpumask_clear_cpu(cpu, donee_mask);
 
-		nr_balanced++;
-		if (!(nr_balanced % SCX_BYPASS_LB_BATCH) && n) {
-			list_move_tail(&cursor.node, &n->scx.dsq_list.node);
-			raw_spin_unlock(&donor_dsq->lock);
-			scx_rq_lock_drop(donor_rq);
-			raw_spin_rq_unlock_irq(donor_rq);
-			cpu_relax();
-			raw_spin_rq_lock_irq(donor_rq);
-			raw_spin_lock(&donor_dsq->lock);
-			goto resume;
+			/*
+			 * The rq lock is held across the walk, not cycled per
+			 * move. @donor_dsq can be very long and contended, and
+			 * the rq lock is what keeps its consumer and enqueuers
+			 * away while the walk cycles the DSQ lock. The walk can
+			 * still take a while, so let the cpu breathe every so
+			 * often.
+			 */
+			if (!(++nr_balanced % SCX_BYPASS_LB_BATCH)) {
+				raw_spin_unlock(&donor_dsq->lock);
+				scx_rq_lock_drop(donor_rq);
+				raw_spin_rq_unlock_irq(donor_rq);
+				cpu_relax();
+				raw_spin_rq_lock_irq(donor_rq);
+				raw_spin_lock(&donor_dsq->lock);
+			}
 		}
+		goto resume;
 	}
 
 	list_del_init(&cursor.node);
-- 
2.56.0

[PATCH 16/18] sched_ext: Let the bypass load balancer move tasks homed on other cpus

[Tejun Heo]

The following patch uses the bypass load balancer to drain a cpu's bypass
DSQ when the host loses SCX_CAP_BASE there. The host no longer consumes that
DSQ, so the drain has to move every task on it. The balancer skips the ones
homed on another cpu.

An earlier pass moved such a task in from its home cpu's DSQ, and the donor
cpu has not consumed it yet. A task's DSQ membership may only change under
the task's rq lock, which the balancer does not hold.

Take such a task off the donor DSQ with unlink_dsq_and_switch_rq_lock(), the
handshake a remote consumer uses. It drops the DSQ lock and switches to the
task's rq lock. A dequeue that gets in between takes the task, and the
balancer moves on. The insert then runs under the task's home rq lock, as it
already does for a donor-homed task.

A task whose donee is its own cpu can always run there.
task_can_run_on_remote_rq() warns when asked about a task's own cpu, so the
balancer skips the test for such a task.

The walk holds the donor rq lock, which keeps the donor cpu's dispatch and
its enqueuers away from the donor DSQ. A move of a task homed elsewhere
switches to the task's rq lock, and both can reach the DSQ meanwhile. The
parked cursor keeps the walk's place through a consume, and the sequence
bound keeps a task enqueued meanwhile out of the pass.

v2: The description says "the following patch", not "patches", uses the
balancer to drain a DSQ.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c | 64 +++++++++++++++++++++++++-----------------
 1 file changed, 39 insertions(+), 25 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index d8e0d99886c6..f02be209c44a 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -6624,7 +6624,7 @@ bool scx_hardlockup(int cpu)
  * @sch: host sched whose bypass DSQs are balanced
  * @donor_rq: rq of the cpu whose bypass DSQ gives up @p, locked
  * @donor_dsq: @sch's bypass DSQ on @donor_rq's cpu, locked
- * @p: task on @donor_dsq, homed on @donor_rq
+ * @p: task on @donor_dsq, homed anywhere
  * @donee: cpu whose bypass DSQ receives @p
  *
  * Take @p off @donor_dsq and insert it into the bypass DSQ of @donee. The
@@ -6632,28 +6632,45 @@ bool scx_hardlockup(int cpu)
  * can still refuse the move under @p's rq lock. @p then goes back to the head
  * of @donor_dsq.
  *
+ * The insert runs under @p's rq lock. A task homed on @donor_rq is dequeued in
+ * place. unlink_dsq_and_switch_rq_lock() takes one homed elsewhere off and
+ * switches the rq lock from @donor_rq to its own. The switch can lose to a
+ * dequeue, and @p then belongs to the dequeuer.
+ *
  * @donor_rq and @donor_dsq are locked on entry and on return. @donor_dsq is
- * unlocked across the insert, so the caller's walk must find its place in it
- * again after the call.
+ * unlocked in between, and @donor_rq too while the lock is switched, so the
+ * caller's walk must find its place in @donor_dsq again after the call.
  *
- * Return the cpu that received @p, or -1 if @p went back onto @donor_dsq.
+ * Return the cpu that received @p, or -1 if @p went back onto @donor_dsq or was
+ * lost to a dequeue.
  */
 static s32 bypass_lb_move(struct scx_sched *sch, struct rq *donor_rq,
 			  struct scx_dispatch_q *donor_dsq, struct task_struct *p, s32 donee)
 {
 	struct scx_dispatch_q *dst_dsq;
 	u64 enq_flags;
+	struct rq *rq = task_rq(p);
 	s32 ret;
 
 	lockdep_assert_rq_held(donor_rq);
 	lockdep_assert_held(&donor_dsq->lock);
 
-	scx_wait_dispatching(p);
-	dispatch_dequeue_locked(p, donor_dsq);
-	raw_spin_unlock(&donor_dsq->lock);
+	if (rq == donor_rq) {
+		scx_wait_dispatching(p);
+		dispatch_dequeue_locked(p, donor_dsq);
+		raw_spin_unlock(&donor_dsq->lock);
+	} else if (!unlink_dsq_and_switch_rq_lock(p, donor_dsq, donor_rq, rq)) {
+		/* lost to a dequeue, @p is the dequeuer's now */
+		ret = -1;
+		goto out;
+	}
 
-	/* @p's affinity is stable under its rq lock */
-	if (task_can_run_on_remote_rq(sch, p, cpu_rq(donee), false)) {
+	/*
+	 * @p's affinity is stable under its rq lock. A task whose donee is its
+	 * own cpu can always run there, and task_can_run_on_remote_rq() warns
+	 * when asked about a task's own cpu.
+	 */
+	if (donee == task_cpu(p) || task_can_run_on_remote_rq(sch, p, cpu_rq(donee), false)) {
 		dst_dsq = scx_bypass_dsq(sch, donee);
 		enq_flags = 0;
 		ret = donee;
@@ -6664,8 +6681,10 @@ static s32 bypass_lb_move(struct scx_sched *sch, struct rq *donor_rq,
 		ret = -1;
 	}
 
-	scx_dispatch_enqueue(sch, donor_rq, dst_dsq, p, 0, 0, enq_flags);
-
+	scx_dispatch_enqueue(sch, rq, dst_dsq, p, 0, 0, enq_flags);
+out:
+	if (rq != donor_rq)
+		switch_rq_lock(rq, donor_rq);
 	raw_spin_lock(&donor_dsq->lock);
 	return ret;
 }
@@ -6718,22 +6737,16 @@ static u32 bypass_lb_cpu(struct scx_sched *sch, s32 donor,
 		if (u32_before(cursor.priv, p->scx.dsq_seq))
 			continue;
 
-		/*
-		 * If an earlier pass placed @p on @donor_dsq from a different
-		 * CPU and the donee hasn't consumed it yet, @p is still on the
-		 * previous CPU and task_rq(@p) != @donor_rq. @p can't be moved
-		 * without its rq locked. Skip.
-		 */
-		if (task_rq(p) != donor_rq)
-			continue;
-
 		donee = cpumask_any_and_distribute(donee_mask, p->cpus_ptr);
 		if (donee >= nr_cpu_ids)
 			continue;
 
 		/*
-		 * The move drops @donor_dsq's lock. Park the cursor behind @p
-		 * and resume from it afterwards.
+		 * The move drops @donor_dsq's lock, and the rq lock too when
+		 * @p is homed elsewhere. The donor cpu's dispatch and its
+		 * enqueuers can reach @donor_dsq then. Park the cursor behind
+		 * @p and resume from it afterwards. The seq test above keeps a
+		 * task enqueued meanwhile out of the pass.
 		 */
 		list_move(&cursor.node, &p->scx.dsq_list.node);
 		cpu = bypass_lb_move(sch, donor_rq, donor_dsq, p, donee);
@@ -6752,9 +6765,10 @@ static u32 bypass_lb_cpu(struct scx_sched *sch, s32 donor,
 			 * The rq lock is held across the walk, not cycled per
 			 * move. @donor_dsq can be very long and contended, and
 			 * the rq lock is what keeps its consumer and enqueuers
-			 * away while the walk cycles the DSQ lock. The walk can
-			 * still take a while, so let the cpu breathe every so
-			 * often.
+			 * away while the walk cycles the DSQ lock. The rare
+			 * move of a task homed elsewhere does switch it away,
+			 * see the cursor comment above. The walk can still take
+			 * a while, so let the cpu breathe every so often.
 			 */
 			if (!(++nr_balanced % SCX_BYPASS_LB_BATCH)) {
 				raw_spin_unlock(&donor_dsq->lock);
-- 
2.56.0

[PATCH 17/18] sched_ext: Fix queued bypassed tasks stalling when the host loses SCX_CAP_BASE

[Tejun Heo]

A task already queued on a bypass DSQ stalls when the host loses
SCX_CAP_BASE on that cpu, as the cpu no longer consumes the DSQ. Sending a
bypassed task to rescue at enqueue time covers only a task enqueued after
the loss, and the balancer no longer donates to a capless cpu.

Have the bypass load balancer drain such a DSQ: a cpu where the host lacks
BASE is a donor with a target of zero, and the donees take its tasks under
their ordinary fill.

A drained task no donee can take is reenqueued at its home: to rescue when
the host lacks BASE there and rescue is enabled, and to the home's bypass
DSQ otherwise.

A balancing donor leaves a task that can't move to its donee where it is.
Reenqueueing it at home would put a task homed on the donor cpu at the tail
of its own DSQ, and a blocked task that proxy execution migrated there,
which every remote cpu refuses, would go to the tail on every pass.

Without rescue, a task homed on the drained cpu goes back onto the drained
DSQ, once per pass, and waits for the host to regain BASE, the bypass to
end, or the stall watchdog to eject the host. Its reenqueues are not counted
against the reenqueue limit, as its scheduler is bypassing.

A BASE loss kicks the balancer so that the drain doesn't wait for the next
interval. With a zero interval the kick is the only pass, and a task that
lands on the drained DSQ after it waits the same way.

v2: The description drops an in-flight affinity change from the causes of a
refused balancer move.

v3:
- The description says that a stuck task's reenqueues are not counted
  against the reenqueue limit (sashiko).
- The description names the stall watchdog as the way out for a task stuck
  without rescue, and the tasks a zero interval leaves on a drained DSQ.
- bypass_lb_reenq_home()'s kerneldoc and the description give the placement
  with rescue disabled and for a task that must run at home, and the other
  two comments on that placement point to the kerneldoc.

v4:
- Skip a task the donee would refuse on a balancing donor instead of
  reenqueueing it at home (Andrea).
- bypass_lb_reenq_home()'s kerneldoc is corrected on a migration-disabled
  task homed on a donee and on the local DSQ admission.
- The description is condensed.

Fixes: 147d1885f390 ("sched_ext: Add the SCX_CAP_ENQ_IMMED cap")
Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c      | 163 ++++++++++++++++++++++++++----------
 kernel/sched/ext/internal.h |   1 +
 kernel/sched/ext/sub.c      |   8 ++
 3 files changed, 129 insertions(+), 43 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index f02be209c44a..c7a5eba4b468 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -6619,36 +6619,76 @@ bool scx_hardlockup(int cpu)
 	return true;
 }
 
+/**
+ * bypass_lb_reenq_home - Reenqueue a task at its home from the balancer
+ * @rq: @p's home rq, locked
+ * @p: task off any DSQ, taken from the bypass DSQ of @dsq_cpu
+ * @dsq_cpu: cpu whose bypass DSQ gave up @p
+ *
+ * @p came off the bypass DSQ of @dsq_cpu with no donee able to take it. A drain
+ * sends a migration-disabled @p here unless the donee picked for it is its
+ * home, as every other cpu refuses it. The enqueue sends @p to rescue when the
+ * host lacks %SCX_CAP_BASE at home and rescue is enabled, and to the home's
+ * bypass DSQ otherwise. When @p must run at home, migration-disabled or with a
+ * migration pending, and the host lacks BASE there, it is admitted to the
+ * home's local DSQ in either case, see bypass_enq_target_dsq(). No DSQ lock may
+ * be held, as the enqueue takes one.
+ */
+static void bypass_lb_reenq_home(struct rq *rq, struct task_struct *p, s32 dsq_cpu)
+{
+	lockdep_assert_rq_held(rq);
+
+	WARN_ON_ONCE(p->scx.flags & SCX_TASK_REENQ_REASON_MASK);
+	p->scx.flags |= SCX_TASK_REENQ_CAP;
+	p->scx.reenq_reason_caps = SCX_CAP_BASE;
+	p->scx.reenq_reason_cid = __scx_cpu_to_cid(dsq_cpu);
+	scx_do_enqueue_task(rq, p, SCX_ENQ_REENQ, -1);
+}
+
+/**
+ * bypass_lb_can_move - Test whether @p can move to the bypass DSQ of @donee
+ * @sch: host sched whose bypass DSQs are balanced
+ * @p: task on a bypass DSQ of @sch, homed anywhere
+ * @donee: cpu whose bypass DSQ would receive @p, -1 for none
+ *
+ * @p's rq lock pins @p's affinity. Without it the answer is a hint. A task
+ * whose donee is its own cpu can always run there, and
+ * task_can_run_on_remote_rq() warns when asked about a task's own cpu.
+ */
+static bool bypass_lb_can_move(struct scx_sched *sch, struct task_struct *p, s32 donee)
+{
+	return donee >= 0 && (donee == task_cpu(p) ||
+			      task_can_run_on_remote_rq(sch, p, cpu_rq(donee), false));
+}
+
 /**
  * bypass_lb_move - Move one task off a bypass DSQ for the balancer
  * @sch: host sched whose bypass DSQs are balanced
  * @donor_rq: rq of the cpu whose bypass DSQ gives up @p, locked
  * @donor_dsq: @sch's bypass DSQ on @donor_rq's cpu, locked
  * @p: task on @donor_dsq, homed anywhere
- * @donee: cpu whose bypass DSQ receives @p
+ * @donee: cpu whose bypass DSQ receives @p, -1 for none
  *
  * Take @p off @donor_dsq and insert it into the bypass DSQ of @donee. The
- * caller picks @donee from @p's affinity alone and task_can_run_on_remote_rq()
- * can still refuse the move under @p's rq lock. @p then goes back to the head
- * of @donor_dsq.
+ * caller tests bypass_lb_can_move() without @p's rq lock for a task homed
+ * elsewhere, and not at all for a drain, so the test is repeated here under
+ * @p's rq lock. A task that fails it is not put back, as the host may no longer
+ * consume @donor_dsq. bypass_lb_reenq_home() reenqueues it at its home instead.
  *
- * The insert runs under @p's rq lock. A task homed on @donor_rq is dequeued in
- * place. unlink_dsq_and_switch_rq_lock() takes one homed elsewhere off and
- * switches the rq lock from @donor_rq to its own. The switch can lose to a
- * dequeue, and @p then belongs to the dequeuer.
+ * The insert and the reenqueue run under @p's rq lock. A task homed on
+ * @donor_rq is dequeued in place. unlink_dsq_and_switch_rq_lock() takes one
+ * homed elsewhere off and switches the rq lock from @donor_rq to its own. The
+ * switch can lose to a dequeue, and @p then belongs to the dequeuer.
  *
  * @donor_rq and @donor_dsq are locked on entry and on return. @donor_dsq is
  * unlocked in between, and @donor_rq too while the lock is switched, so the
  * caller's walk must find its place in @donor_dsq again after the call.
  *
- * Return the cpu that received @p, or -1 if @p went back onto @donor_dsq or was
- * lost to a dequeue.
+ * Return the cpu that received @p, or -1 if @p was lost to a dequeue.
  */
 static s32 bypass_lb_move(struct scx_sched *sch, struct rq *donor_rq,
 			  struct scx_dispatch_q *donor_dsq, struct task_struct *p, s32 donee)
 {
-	struct scx_dispatch_q *dst_dsq;
-	u64 enq_flags;
 	struct rq *rq = task_rq(p);
 	s32 ret;
 
@@ -6665,23 +6705,17 @@ static s32 bypass_lb_move(struct scx_sched *sch, struct rq *donor_rq,
 		goto out;
 	}
 
-	/*
-	 * @p's affinity is stable under its rq lock. A task whose donee is its
-	 * own cpu can always run there, and task_can_run_on_remote_rq() warns
-	 * when asked about a task's own cpu.
-	 */
-	if (donee == task_cpu(p) || task_can_run_on_remote_rq(sch, p, cpu_rq(donee), false)) {
-		dst_dsq = scx_bypass_dsq(sch, donee);
-		enq_flags = 0;
+	if (bypass_lb_can_move(sch, p, donee)) {
+		scx_dispatch_enqueue(sch, rq, scx_bypass_dsq(sch, donee), p, 0, 0, 0);
 		ret = donee;
 	} else {
-		/* back where it was, at the head so that the walk skips it */
-		dst_dsq = donor_dsq;
-		enq_flags = SCX_ENQ_HEAD;
-		ret = -1;
+		/*
+		 * No donee, or @p can't run on it. Reenqueue @p at home
+		 * instead, see bypass_lb_reenq_home().
+		 */
+		bypass_lb_reenq_home(rq, p, cpu_of(donor_rq));
+		ret = cpu_of(rq);
 	}
-
-	scx_dispatch_enqueue(sch, rq, dst_dsq, p, 0, 0, enq_flags);
 out:
 	if (rq != donor_rq)
 		switch_rq_lock(rq, donor_rq);
@@ -6704,10 +6738,11 @@ static u32 bypass_lb_cpu(struct scx_sched *sch, s32 donor,
 	 * All we want to guarantee is reasonable forward progress. No reason to
 	 * fine tune. Assuming every task on @donor_dsq runs their full slice,
 	 * consider offloading iff the total queued duration is over the
-	 * threshold.
+	 * threshold. A zero target is a drain and empties the DSQ regardless.
 	 */
 	min_delta_us = READ_ONCE(scx_bypass_lb_intv_us) / SCX_BYPASS_LB_MIN_DELTA_DIV;
-	if (delta < DIV_ROUND_UP(min_delta_us, READ_ONCE(scx_slice_bypass_us)))
+	if (nr_donor_target &&
+	    delta < DIV_ROUND_UP(min_delta_us, READ_ONCE(scx_slice_bypass_us)))
 		return 0;
 
 	raw_spin_rq_lock_irq(donor_rq);
@@ -6727,18 +6762,29 @@ static u32 bypass_lb_cpu(struct scx_sched *sch, s32 donor,
 		if (donor_dsq->nr <= nr_donor_target)
 			break;
 
-		if (cpumask_empty(donee_mask))
+		/*
+		 * Without a donee, a balancing donor has nothing to do. A drain
+		 * still sends its tasks home.
+		 */
+		if (cpumask_empty(donee_mask) && nr_donor_target)
 			break;
 
 		/*
 		 * A task queued after the walk began, including one the walk
-		 * itself put back, is left for the next pass.
+		 * itself reenqueued, is left for the next pass.
 		 */
 		if (u32_before(cursor.priv, p->scx.dsq_seq))
 			continue;
 
+		/*
+		 * A balancing donor leaves a task that can't move where it is.
+		 * A home reenqueue would push one homed on the donor to the
+		 * back of @donor_dsq on every pass. A drain sends it home.
+		 */
 		donee = cpumask_any_and_distribute(donee_mask, p->cpus_ptr);
 		if (donee >= nr_cpu_ids)
+			donee = -1;
+		if (nr_donor_target && !bypass_lb_can_move(sch, p, donee))
 			continue;
 
 		/*
@@ -6795,6 +6841,7 @@ static void bypass_lb_node(struct scx_sched *sch, int node)
 	const struct cpumask *node_mask = cpumask_of_node(node);
 	struct cpumask *donee_mask = sch->bypass_lb_donee_cpumask;
 	struct cpumask *resched_mask = sch->bypass_lb_resched_cpumask;
+	u32 intv_us = READ_ONCE(scx_bypass_lb_intv_us);
 	u32 nr_tasks = 0, nr_cpus = 0, nr_balanced = 0;
 	u32 nr_target, nr_donor_target;
 	u32 before_min = U32_MAX, before_max = 0;
@@ -6814,7 +6861,7 @@ static void bypass_lb_node(struct scx_sched *sch, int node)
 		before_max = max(nr, before_max);
 	}
 
-	if (!nr_cpus)
+	if (!nr_tasks)
 		return;
 
 	/*
@@ -6822,7 +6869,7 @@ static void bypass_lb_node(struct scx_sched *sch, int node)
 	 * balancing to fill donee CPUs up to $nr_target. Once targets are
 	 * calculated, find the donee CPUs.
 	 */
-	nr_target = DIV_ROUND_UP(nr_tasks, nr_cpus);
+	nr_target = nr_cpus ? DIV_ROUND_UP(nr_tasks, nr_cpus) : 0;
 	nr_donor_target = DIV_ROUND_UP(nr_target * SCX_BYPASS_LB_DONOR_PCT, 100);
 
 	cpumask_clear(donee_mask);
@@ -6833,18 +6880,28 @@ static void bypass_lb_node(struct scx_sched *sch, int node)
 			cpumask_set_cpu(cpu, donee_mask);
 	}
 
-	/* iterate !donee CPUs and see if they should be offloaded */
+	/*
+	 * Offload the !donee cpus. A cpu where the host lost %SCX_CAP_BASE no
+	 * longer consumes its DSQ. Drain it as a donor with a target of zero. A
+	 * zero interval disables the balancing, not the drain.
+	 */
 	cpumask_clear(resched_mask);
 	for_each_cpu_and(cpu, cpu_online_mask, node_mask) {
-		if (cpumask_empty(donee_mask))
-			break;
-		if (cpumask_test_cpu(cpu, donee_mask))
+		u32 target = nr_donor_target;
+
+		if (unlikely(scx_missing_caps(sch, cpu, SCX_CAP_BASE))) {
+			/* the mask may predate the loss, keep this cpu out */
+			cpumask_clear_cpu(cpu, donee_mask);
+			target = 0;
+		} else if (!intv_us || cpumask_empty(donee_mask) ||
+			   cpumask_test_cpu(cpu, donee_mask)) {
 			continue;
-		if (READ_ONCE(scx_bypass_dsq(sch, cpu)->nr) <= nr_donor_target)
+		}
+		if (READ_ONCE(scx_bypass_dsq(sch, cpu)->nr) <= target)
 			continue;
 
 		nr_balanced += bypass_lb_cpu(sch, cpu, donee_mask, resched_mask,
-					     nr_donor_target, nr_target);
+					     target, nr_target);
 	}
 
 	for_each_cpu(cpu, resched_mask)
@@ -6883,9 +6940,28 @@ static void scx_bypass_lb_timerfn(struct timer_list *timer)
 	for_each_node_with_cpus(node)
 		bypass_lb_node(sch, node);
 
+	/*
+	 * A kick during this run can leave the timer pending at now, see
+	 * scx_bypass_lb_kick(). timer_reduce() only brings the expiry forward,
+	 * so that kick is kept.
+	 */
 	intv_us = READ_ONCE(scx_bypass_lb_intv_us);
 	if (intv_us)
-		mod_timer(timer, jiffies + usecs_to_jiffies(intv_us));
+		timer_reduce(timer, jiffies + usecs_to_jiffies(intv_us));
+}
+
+/**
+ * scx_bypass_lb_kick - Run the bypass load balancer now
+ * @sch: sched that lost %SCX_CAP_BASE on a cpu
+ *
+ * @sch no longer consumes its bypass DSQ on that cpu. Have the balancer drain
+ * it without waiting for the next interval, see bypass_lb_node(). A loss during
+ * a run leaves the timer pending at now, and the run that follows sees it.
+ */
+void scx_bypass_lb_kick(struct scx_sched *sch)
+{
+	if (scx_bypass_dsp_enabled(sch))
+		mod_timer(&sch->bypass_lb_timer, jiffies);
 }
 
 static bool inc_bypass_depth(struct scx_sched *sch)
@@ -6954,11 +7030,12 @@ static void enable_bypass_dsp(struct scx_sched *sch)
 
 	/*
 	 * The LB timer will stop running if bypass dispatch is disabled. Start
-	 * after enabling bypass dispatch.
+	 * after enabling bypass dispatch. A kick from a BASE loss can have set
+	 * the timer to expire now in the meantime, see scx_bypass_lb_kick().
+	 * timer_reduce() keeps that expiry.
 	 */
-	if (intv_us && !timer_pending(&host->bypass_lb_timer))
-		mod_timer(&host->bypass_lb_timer,
-			  jiffies + usecs_to_jiffies(intv_us));
+	if (intv_us)
+		timer_reduce(&host->bypass_lb_timer, jiffies + usecs_to_jiffies(intv_us));
 }
 
 /* may be called without holding scx_bypass_lock */
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index ecd677c4b561..ff0a3d508533 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -2332,6 +2332,7 @@ void scx_cgroup_unlock(void);
 s32 scx_alloc_kern_arena_objs(struct scx_sched *sch);
 void scx_disable_bypass_dsp(struct scx_sched *sch);
 void scx_bypass(struct scx_sched *sch, bool bypass);
+void scx_bypass_lb_kick(struct scx_sched *sch);
 s32 scx_link_sched(struct scx_sched *sch);
 void scx_unlink_sched(struct scx_sched *sch);
 void scx_disable_dump(struct scx_sched *sch);
diff --git a/kernel/sched/ext/sub.c b/kernel/sched/ext/sub.c
index 01215384eaa5..ecab82ca5e82 100644
--- a/kernel/sched/ext/sub.c
+++ b/kernel/sched/ext/sub.c
@@ -1243,6 +1243,14 @@ void __scx_process_sync_ecaps(struct rq *rq, struct task_struct *prev)
 		gained = ecaps & ~old;
 		lost_all |= lost;
 
+		/*
+		 * The bypass DSQ on this cpu is no longer consumed, so drain
+		 * it. A bypassing sched's own bypass DSQs are empty, as its
+		 * tasks sit on the host's.
+		 */
+		if ((lost & SCX_CAP_BASE) && !scx_bypassing(pcpu->sch, cpu))
+			scx_bypass_lb_kick(pcpu->sch);
+
 		/*
 		 * Tell the sched and its parent that the sched's effective caps
 		 * on this cid changed. The invocations are equivalent to the
-- 
2.56.0

[PATCH 18/18] sched_ext: Fix a lone IMMED task waiting behind a kept prev

[Tejun Heo]

While dispatch is in progress, rq_is_open() returns true and dsq_inc_nr()
schedules no reenqueue for an IMMED task inserted on an empty local DSQ. The
check at the end of dispatch_one() schedules the reenqueue only when more
than one task is queued, so a dispatch that keeps prev leaves a lone IMMED
task behind:

 1. a parent's ops.dispatch() sub-dispatches into a child, which consumes a
    task with SCX_ENQ_IMMED onto the empty local DSQ
 2. the parent extends prev's slice and dispatch keeps prev
 3. the local DSQ holds one task, so no reenqueue is scheduled

The IMMED task waits on the local DSQ behind prev.

Schedule the reenqueue when dispatch keeps prev and an IMMED task is queued,
as even a lone IMMED task can't run right away behind prev.

v2: The description and dispatch_one()'s comment are reworded.

Fixes: 98d709cba319 ("sched_ext: Implement SCX_ENQ_IMMED")
Reported-by: Sashiko <sashiko-bot@kernel.org>
Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index c7a5eba4b468..b7a78be01b42 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -3500,8 +3500,12 @@ static enum scx_dsp_verdict dispatch_one(struct rq *rq, struct task_struct *prev
 	 *
 	 * - A non-IMMED HEAD task can get queued in front of an IMMED task
 	 *   between the IMMED queueing and the subsequent scheduling event.
+	 *
+	 * - When dispatch keeps @prev, even a lone IMMED task queued during the
+	 *   dispatch can't run right away.
 	 */
-	if (unlikely(rq->scx.local_dsq.nr > 1 && rq->scx.nr_immed))
+	if (unlikely((rq->scx.local_dsq.nr > 1 || verdict == SCX_DSP_PREV) &&
+		     rq->scx.nr_immed))
 		scx_schedule_reenq_local(rq, 0);
 
 	rq->scx.flags &= ~SCX_RQ_IN_DISPATCH;
-- 
2.56.0