[PATCH v2] sched_ext: keep running prev when prev->scx.slice != 0

[PATCH v2] sched_ext: keep running prev when prev->scx.slice != 0

[Henry Huang]

In our environment, we have various types of tasks. Some tasks can keep running after their slice is exhausted,
while others need to be dispatched into the global DSQ for rescheduling. Therefore, we set %SCX_OPS_ENQ_LAST.

However, we encountered a problem:
Because put_prev_task_scx() is executed after pick_task_scx(), @prev only has the opportunity to be dispatched
into local DSQ in put_prev_task_scx(). Since pick_task_scx returns NULL, the CPU enters the idle state instead of running @prev.

Our current workaround is to set a value to @prev->scx.slice in ops.dispatch() and call scx_bpf_kick_cpu(cpu, 0) to trigger a reschedule.
Of course, this approach introduces some overhead.

Our solution:
When %SCX_OPS_ENQ_LAST is set and prev->scx.slice != 0, we only need to set %SCX_RQ_BAL_KEEP in blance_one to
ensure pick_task_scx() can pick the correct task.

For v2:
- set %SCX_RQ_BAL_KEEP in dispatch loop.
- cache prev->scx.flags & SCX_TASK_QUEUED result.

Henry Huang (1):
  sched_ext: keep running prev when prev->scx.slice != 0

 kernel/sched/ext.c | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

-- 
Henry

[PATCH v2] sched_ext: keep running prev when prev->scx.slice != 0

[Henry Huang]

When %SCX_OPS_ENQ_LAST is set and prev->scx.slice != 0,
@prev will be dispacthed into the local DSQ in put_prev_task_scx().
However, pick_task_scx() is executed before put_prev_task_scx(),
so it will not pick @prev.
Set %SCX_RQ_BAL_KEEP in balance_one() to ensure that pick_task_scx()
can pick @prev.

Signed-off-by: Henry Huang <henry.hj@antgroup.com>
---
 kernel/sched/ext.c | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index 81da76a..ba74e07 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -2747,6 +2747,7 @@ static int balance_one(struct rq *rq, struct task_struct *prev)
 {
 	struct scx_dsp_ctx *dspc = this_cpu_ptr(scx_dsp_ctx);
 	bool prev_on_scx = prev->sched_class == &ext_sched_class;
+	bool prev_on_rq = (prev->scx.flags & SCX_TASK_QUEUED) != 0;
 	int nr_loops = SCX_DSP_MAX_LOOPS;
 
 	lockdep_assert_rq_held(rq);
@@ -2779,8 +2780,7 @@ static int balance_one(struct rq *rq, struct task_struct *prev)
 		 * See scx_ops_disable_workfn() for the explanation on the
 		 * bypassing test.
 		 */
-		if ((prev->scx.flags & SCX_TASK_QUEUED) &&
-		    prev->scx.slice && !scx_rq_bypassing(rq)) {
+		if (prev_on_rq && prev->scx.slice && !scx_rq_bypassing(rq)) {
 			rq->scx.flags |= SCX_RQ_BAL_KEEP;
 			goto has_tasks;
 		}
@@ -2813,6 +2813,10 @@ static int balance_one(struct rq *rq, struct task_struct *prev)
 
 		flush_dispatch_buf(rq);
 
+		if (prev_on_rq && prev->scx.slice) {
+			rq->scx.flags |= SCX_RQ_BAL_KEEP;
+			goto has_tasks;
+		}
 		if (rq->scx.local_dsq.nr)
 			goto has_tasks;
 		if (consume_global_dsq(rq))
@@ -2838,8 +2842,7 @@ static int balance_one(struct rq *rq, struct task_struct *prev)
 	 * Didn't find another task to run. Keep running @prev unless
 	 * %SCX_OPS_ENQ_LAST is in effect.
 	 */
-	if ((prev->scx.flags & SCX_TASK_QUEUED) &&
-	    (!static_branch_unlikely(&scx_ops_enq_last) ||
+	if (prev_on_rq && (!static_branch_unlikely(&scx_ops_enq_last) ||
 	     scx_rq_bypassing(rq))) {
 		rq->scx.flags |= SCX_RQ_BAL_KEEP;
 		goto has_tasks;
-- 
Henry

Re: [PATCH v2] sched_ext: keep running prev when prev->scx.slice != 0

[Tejun Heo]

On Wed, Jan 08, 2025 at 04:47:10PM +0800, Henry Huang wrote:
> When %SCX_OPS_ENQ_LAST is set and prev->scx.slice != 0,
> @prev will be dispacthed into the local DSQ in put_prev_task_scx().
> However, pick_task_scx() is executed before put_prev_task_scx(),
> so it will not pick @prev.
> Set %SCX_RQ_BAL_KEEP in balance_one() to ensure that pick_task_scx()
> can pick @prev.
> 
> Signed-off-by: Henry Huang <henry.hj@antgroup.com>

Applied to sched_ext/for-6.13-fixes with a minor edit (see below).

> @@ -2747,6 +2747,7 @@ static int balance_one(struct rq *rq, struct task_struct *prev)
>  {
>  	struct scx_dsp_ctx *dspc = this_cpu_ptr(scx_dsp_ctx);
>  	bool prev_on_scx = prev->sched_class == &ext_sched_class;
> +	bool prev_on_rq = (prev->scx.flags & SCX_TASK_QUEUED) != 0;

I dropped `!= 0` as the bool assignment cast already implies boolean
conversion.

Thanks.

-- 
tejun