sched_ext: Add bandwidth-limited rescue execution for stranded tasks

diff --git a/include/linux/sched/ext.h b/include/linux/sched/ext.h
index b519fbc88e17..a6aabbefd185 100644
--- a/include/linux/sched/ext.h
+++ b/include/linux/sched/ext.h
@@ -59,6 +59,7 @@ enum scx_dsq_id_flags {
 	SCX_DSQ_LOCAL		= SCX_DSQ_FLAG_BUILTIN | 2,
 	SCX_DSQ_BYPASS		= SCX_DSQ_FLAG_BUILTIN | 3,
 	SCX_DSQ_REJECT		= SCX_DSQ_FLAG_BUILTIN | 4,	/* internal - see find_dsq_for_dispatch() */
+	SCX_DSQ_RESCUE		= SCX_DSQ_FLAG_BUILTIN | 5,	/* internal - see find_dsq_for_dispatch() */
 	SCX_DSQ_LOCAL_ON	= SCX_DSQ_FLAG_BUILTIN | SCX_DSQ_FLAG_LOCAL_ON,

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
@@ -2056,6 +2067,7 @@ void scx_do_enqueue_task(struct rq *rq, struct task_struct *p, u64 enq_flags,
 	if (!(sch->ops.flags & SCX_OPS_ENQ_EXITING) &&
 	    unlikely(p->flags & PF_EXITING)) {
 		__scx_add_event(sch, SCX_EV_ENQ_SKIP_EXITING, 1);
+		enq_flags |= SCX_ENQ_RESCUE;	/* avoid looping on cap rejection */
 		goto local;
 	}

diff --git a/kernel/sched/ext/sub.c b/kernel/sched/ext/sub.c
@@ -282,18 +619,27 @@ struct scx_dispatch_q *scx_resolve_local_dsq(struct scx_sched *sch, struct rq *r
 	/*
-	 * Only local DSQ can honor IMMED and dsq_inc_nr() WARNs on IMMED into
-	 * others. Strip both the enq flag and the sticky task flag - the
-	 * latter can carry in from an earlier admitted IMMED insert. Strip
-	 * PREEMPT too.
+	 * Diverting to rescue or reject, neither of which honors IMMED, PREEMPT
+	 * or HEAD - a diversion has no priority and IMMED is not allowed on
+	 * non-local DSQs. Strip the enq and task flags along with the slice.
 	 */
-	*enq_flags &= ~(SCX_ENQ_IMMED | SCX_ENQ_PREEMPT);
+	*enq_flags &= ~(SCX_ENQ_IMMED | SCX_ENQ_PREEMPT | SCX_ENQ_HEAD |
+			SCX_ENQ_APPLY_SLICE | SCX_ENQ_SLICE_DFL);
 	p->scx.flags &= ~SCX_TASK_IMMED;
 
+	/* the enqueuer opted for rescue instead of rejection and reenqueue */
+	if ((*enq_flags & SCX_ENQ_RESCUE) && likely(scx_rescue_bw_1024)) {
+		__scx_add_event(sch, SCX_EV_SUB_RESCUE, 1);
+		if (scx_rescue_try_admit(rq, p))
+			return &rq->scx.local_dsq;
+		else
+			return &rq->scx.rescue.dsq;
+	}
+
+	p->scx.reenq_reason_caps = missing;
+	p->scx.reenq_reason_cid = cid;
+
 	return &rq->scx.reject_dsq;
 }

+static bool scx_rescue_try_admit(struct rq *rq, struct task_struct *p)
+{
+	scx_rescue_accrue(rq);
+
+	if (!rq->scx.rescue.curr && list_empty(&rq->scx.rescue.dsq.list) &&
+	    rq->scx.rescue.budget >= scx_rescue_quantum_ns) {
+		scx_rescue_admit(rq, p, scx_rescue_quantum_ns);
+		return true;
+	}
+
+	scx_rescue_timer_arm(rq);
+	return false;
+}

+	} else if (p->scx.dsq && rq->scx.rescue.budget > 2 * scx_rescue_quantum_ns) {
+		/*
+		 * The rescuee waited for the CPU for too long. Escalate - grant
+		 * the unserved remainder, protect it from the schedulers and
+		 * preempt the current task. The slice is set before the
+		 * protection. Repeat beats only repeat the head move - the
+		 * slice write is refused on a protected task.
+		 */
+		scx_set_task_slice(p, scx_rescue_slice_remaining(rq));
+		p->scx.flags |= SCX_TASK_PROTECTED;
+		scx_task_unlink_from_dsq(p, &rq->scx.local_dsq);
+		scx_move_local_task_to_local_dsq(scx_task_sched(p), p,
+					SCX_ENQ_HEAD | SCX_ENQ_PREEMPT | SCX_ENQ_IGNORE_CAPS,
+					&rq->scx.local_dsq, rq);
+	}
+out_arm:
+	scx_rescue_timer_arm(rq);
+}

+void scx_rescue_charge(struct rq *rq, s64 delta_exec)
+{
+	lockdep_assert_rq_held(rq);
+
+	/*
+	 * 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.
+	 */
+	delta_exec = min_t(s64, delta_exec, scx_rescue_quantum_ns + TICK_NSEC);
+
+	rq->scx.rescue.budget -= delta_exec;
+
+	if (!scx_rescue_slice_remaining(rq))
+		scx_task_slice_ended(rq, rq->scx.rescue.curr);
+}

+bool scx_rescue_keep(struct rq *rq, struct task_struct *p)
+{
+	s64 remaining = scx_rescue_slice_remaining(rq);
+
+	lockdep_assert_rq_held(rq);
+
+	if (!remaining || !(p->scx.flags & SCX_TASK_QUEUED) ||
+	    scx_bypassing(scx_task_sched(p), cpu_of(rq)))
+		return false;
+
+	scx_set_task_slice(p, remaining);
+	return true;
+}