sched_ext: Handle cases where pick_task_scx() is called without preceding balance_scx()

sched_ext dispatches tasks from the BPF scheduler from balance_scx() and thus every pick_task_scx() call must be preceded by balance_scx(). While this usually holds, due to a bug, there are cases where the fair class's balance() returns true indicating that it has tasks to run on the CPU and thus terminating balance() calls but fails to actually find the next task to run when pick_task() is called. In such cases, pick_task_scx() can be called without preceding balance_scx(). Detect this condition using SCX_RQ_BAL_PENDING flags. If detected, keep running the previous task if possible and avoid stalling from entering idle without balancing.

Tejun Heo committed 2 years ago

Commit a6250aa1 parent a759bf03 files · +42 −20Browse files

Unified · Split

3 files

kernel/sched/core.c +8 −5
kernel/sched/ext.c +31 −13
kernel/sched/sched.h +3 −2

kernel/sched/core.c+8−5

@@ -5914,12 +5914,15 @@ static void prev_balance(struct rq rq, struct task_struct prev, · 5913 unmodified lines
5914	5914
5915	5915	#ifdef CONFIG_SCHED_CLASS_EXT
5916	5916	/*
5917		− * SCX requires a balance() call before every pick_next_task() including
5918		− * when waking up from SCHED_IDLE. If @start_class is below SCX, start
5919		− * from SCX instead.
	5917	+ * SCX requires a balance() call before every pick_task() including when
	5918	+ * waking up from SCHED_IDLE. If @start_class is below SCX, start from
	5919	+ * SCX instead. Also, set a flag to detect missing balance() call.
5920	5920	*/
5921		− if (scx_enabled() && sched_class_above(&ext_sched_class, start_class))
5922		− start_class = &ext_sched_class;
	5921	+ if (scx_enabled()) {
	5922	+ rq->scx.flags \|= SCX_RQ_BAL_PENDING;
	5923	+ if (sched_class_above(&ext_sched_class, start_class))
	5924	+ start_class = &ext_sched_class;
	5925	+ }
5923	5926	#endif
5924	5927
5925	5928	/*

kernel/sched/ext.c+31−13

@@ -2634,7 +2634,7 @@ static int balance_one(struct rq rq, struct task_struct prev) · 2633 unmodified lines
2634	2634
2635	2635	lockdep_assert_rq_held(rq);
2636	2636	rq->scx.flags \|= SCX_RQ_IN_BALANCE;
2637		− rq->scx.flags &= ~SCX_RQ_BAL_KEEP;
	2637	+ rq->scx.flags &= ~(SCX_RQ_BAL_PENDING \| SCX_RQ_BAL_KEEP);
2638	2638
2639	2639	if (static_branch_unlikely(&scx_ops_cpu_preempt) &&
2640	2640	unlikely(rq->scx.cpu_released)) {
@@ -2948,12 +2948,11 @@ static struct task_struct pick_task_scx(struct rq rq) · 307 unmodified lines
2948	2948	{
2949	2949	struct task_struct *prev = rq->curr;
2950	2950	struct task_struct *p;
	2951	+ bool prev_on_scx = prev->sched_class == &ext_sched_class;
	2952	+ bool keep_prev = rq->scx.flags & SCX_RQ_BAL_KEEP;
	2953	+ bool kick_idle = false;
2951	2954
2952	2955	/*
2953		− * If balance_scx() is telling us to keep running @prev, replenish slice
2954		− * if necessary and keep running @prev. Otherwise, pop the first one
2955		− * from the local DSQ.
2956		− *
2957	2956	* WORKAROUND:
2958	2957	*
2959	2958	* %SCX_RQ_BAL_KEEP should be set iff $prev is on SCX as it must just
@@ -2962,22 +2961,41 @@ static struct task_struct pick_task_scx(struct rq rq) · 2 unmodified lines
2962	2961	* which then ends up calling pick_task_scx() without preceding
2963	2962	* balance_scx().
2964	2963	*
2965		− * For now, ignore cases where $prev is not on SCX. This isn't great and
2966		− * can theoretically lead to stalls. However, for switch_all cases, this
2967		− * happens only while a BPF scheduler is being loaded or unloaded, and,
2968		− * for partial cases, fair will likely keep triggering this CPU.
	2964	+ * Keep running @prev if possible and avoid stalling from entering idle
	2965	+ * without balancing.
2969	2966	*
2970		− * Once fair is fixed, restore WARN_ON_ONCE().
	2967	+ * Once fair is fixed, remove the workaround and trigger WARN_ON_ONCE()
	2968	+ * if pick_task_scx() is called without preceding balance_scx().
2971	2969	*/
2972		− if ((rq->scx.flags & SCX_RQ_BAL_KEEP) &&
2973		− prev->sched_class == &ext_sched_class) {
	2970	+ if (unlikely(rq->scx.flags & SCX_RQ_BAL_PENDING)) {
	2971	+ if (prev_on_scx) {
	2972	+ keep_prev = true;
	2973	+ } else {
	2974	+ keep_prev = false;
	2975	+ kick_idle = true;
	2976	+ }
	2977	+ } else if (unlikely(keep_prev && !prev_on_scx)) {
	2978	+ /* only allowed during transitions */
	2979	+ WARN_ON_ONCE(scx_ops_enable_state() == SCX_OPS_ENABLED);
	2980	+ keep_prev = false;
	2981	+ }
	2982	+
	2983	+ /*
	2984	+ * If balance_scx() is telling us to keep running @prev, replenish slice
	2985	+ * if necessary and keep running @prev. Otherwise, pop the first one
	2986	+ * from the local DSQ.
	2987	+ */
	2988	+ if (keep_prev) {
2974	2989	p = prev;
2975	2990	if (!p->scx.slice)
2976	2991	p->scx.slice = SCX_SLICE_DFL;
2977	2992	} else {
2978	2993	p = first_local_task(rq);
2979		− if (!p)
	2994	+ if (!p) {
	2995	+ if (kick_idle)
	2996	+ scx_bpf_kick_cpu(cpu_of(rq), SCX_KICK_IDLE);
2980	2997	return NULL;
	2998	+ }
2981	2999
2982	3000	if (unlikely(!p->scx.slice)) {
2983	3001	if (!scx_rq_bypassing(rq) && !scx_warned_zero_slice) {

kernel/sched/sched.h+3−2

@@ -751,8 +751,9 @@ enum scx_rq_flags { · 750 unmodified lines
751	751	*/
752	752	SCX_RQ_ONLINE = 1 << 0,
753	753	SCX_RQ_CAN_STOP_TICK = 1 << 1,
754		− SCX_RQ_BAL_KEEP = 1 << 2, /* balance decided to keep current */
755		− SCX_RQ_BYPASSING = 1 << 3,
	754	+ SCX_RQ_BAL_PENDING = 1 << 2, /* balance hasn't run yet */
	755	+ SCX_RQ_BAL_KEEP = 1 << 3, /* balance decided to keep current */
	756	+ SCX_RQ_BYPASSING = 1 << 4,
756	757
757	758	SCX_RQ_IN_WAKEUP = 1 << 16,
758	759	SCX_RQ_IN_BALANCE = 1 << 17,