[PATCH 14/16] sched/deadline: De-couple balance and pick_task

[Andrea Righi <arighi@nvidia.com>]

From: Tejun Heo <tj@kernel.org>

Allow a dl_server to trigger ->balance() from balance_dl() for sched
classes that are always expecting a ->balance() call before
->pick_task(), e.g. sched_ext.

[ arighi:
    - adjust patch after dropping @rf from pick_task()
    - update dl_server_init() to take an additional @balance parameter
    - activate DL server balance only if there's any pending work ]

Co-developed-by: Andrea Righi <arighi@nvidia.com>
Signed-off-by: Andrea Righi <arighi@nvidia.com>
Signed-off-by: Tejun Heo <tj@kernel.org>
---
 include/linux/sched.h   |  2 ++
 kernel/sched/core.c     | 14 +++++++++++---
 kernel/sched/deadline.c | 16 ++++++++++------
 kernel/sched/ext.c      | 17 ++++++++++-------
 kernel/sched/fair.c     |  2 +-
 kernel/sched/sched.h    |  8 +++++++-
 6 files changed, 41 insertions(+), 18 deletions(-)

diff --git a/include/linux/sched.h b/include/linux/sched.h
index 2b272382673d6..aa3ae42da51a9 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -635,6 +635,7 @@ struct sched_rt_entity {
 } __randomize_layout;
 
 typedef bool (*dl_server_has_tasks_f)(struct sched_dl_entity *);
+typedef void (*dl_server_balance_f)(struct sched_dl_entity *, void *);
 typedef struct task_struct *(*dl_server_pick_f)(struct sched_dl_entity *);
 
 struct sched_dl_entity {
@@ -734,6 +735,7 @@ struct sched_dl_entity {
 	 */
 	struct rq			*rq;
 	dl_server_has_tasks_f		server_has_tasks;
+	dl_server_balance_f		server_balance;
 	dl_server_pick_f		server_pick_task;
 
 #ifdef CONFIG_RT_MUTEXES
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index f1a7ad7e560fb..3c2863d961f38 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -5950,14 +5950,22 @@ static void prev_balance(struct rq *rq, struct task_struct *prev,
 
 #ifdef CONFIG_SCHED_CLASS_EXT
 	/*
-	 * SCX requires a balance() call before every pick_task() including when
-	 * waking up from SCHED_IDLE. If @start_class is below SCX, start from
-	 * SCX instead. Also, set a flag to detect missing balance() call.
+	 * SCX requires a balance() call before every pick_task() including
+	 * when waking up from SCHED_IDLE.
+	 *
+	 * If @start_class is below SCX, start balancing from SCX. If the
+	 * DL server has any pending work, start from the DL class instead.
+	 * This ensures the DL server is given a chance to trigger its own
+	 * balance() pass on every prev_balance() invocation.
+	 *
+	 * Also, set a flag to detect missing balance() call.
 	 */
 	if (scx_enabled()) {
 		rq->scx.flags |= SCX_RQ_BAL_PENDING;
 		if (sched_class_above(&ext_sched_class, start_class))
 			start_class = &ext_sched_class;
+		if (on_dl_rq(&rq->ext_server))
+			start_class = &dl_sched_class;
 	}
 #endif
 
diff --git a/kernel/sched/deadline.c b/kernel/sched/deadline.c
index 84c7172ee805c..1f79b1e49b49c 100644
--- a/kernel/sched/deadline.c
+++ b/kernel/sched/deadline.c
@@ -88,11 +88,6 @@ static inline struct dl_rq *dl_rq_of_se(struct sched_dl_entity *dl_se)
 	return &rq_of_dl_se(dl_se)->dl;
 }
 
-static inline int on_dl_rq(struct sched_dl_entity *dl_se)
-{
-	return !RB_EMPTY_NODE(&dl_se->rb_node);
-}
-
 #ifdef CONFIG_RT_MUTEXES
 static inline struct sched_dl_entity *pi_of(struct sched_dl_entity *dl_se)
 {
@@ -1650,11 +1645,13 @@ static bool dl_server_stopped(struct sched_dl_entity *dl_se)
 
 void dl_server_init(struct sched_dl_entity *dl_se, struct rq *rq,
 		    dl_server_has_tasks_f has_tasks,
-		    dl_server_pick_f pick_task)
+		    dl_server_pick_f pick_task,
+		    dl_server_balance_f balance)
 {
 	dl_se->rq = rq;
 	dl_se->server_has_tasks = has_tasks;
 	dl_se->server_pick_task = pick_task;
+	dl_se->server_balance = balance;
 }
 
 void sched_init_dl_servers(void)
@@ -2349,8 +2346,12 @@ static void check_preempt_equal_dl(struct rq *rq, struct task_struct *p)
 	resched_curr(rq);
 }
 
+static struct sched_dl_entity *pick_next_dl_entity(struct dl_rq *dl_rq);
+
 static int balance_dl(struct rq *rq, struct task_struct *p, struct rq_flags *rf)
 {
+	struct sched_dl_entity *dl_se;
+
 	if (!on_dl_rq(&p->dl) && need_pull_dl_task(rq, p)) {
 		/*
 		 * This is OK, because current is on_cpu, which avoids it being
@@ -2363,6 +2364,9 @@ static int balance_dl(struct rq *rq, struct task_struct *p, struct rq_flags *rf)
 		rq_repin_lock(rq, rf);
 	}
 
+	dl_se = pick_next_dl_entity(&rq->dl);
+	if (dl_se && dl_server(dl_se) && dl_se->server_balance)
+		dl_se->server_balance(dl_se, rf);
 	return sched_stop_runnable(rq) || sched_dl_runnable(rq);
 }
 
diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index 69163927a29cd..e6d84b9aa70dc 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -7715,16 +7715,19 @@ static bool ext_server_has_tasks(struct sched_dl_entity *dl_se)
 	return !!dl_se->rq->scx.nr_running;
 }
 
-/*
- * Select the next task to run from the ext scheduling class.
- */
-static struct task_struct *ext_server_pick_task(struct sched_dl_entity *dl_se,
-						void *flags)
+static void ext_server_balance(struct sched_dl_entity *dl_se, void *flags)
 {
 	struct rq_flags *rf = flags;
 
 	balance_scx(dl_se->rq, dl_se->rq->curr, rf);
-	return pick_task_scx(dl_se->rq, rf);
+}
+
+/*
+ * Select the next task to run from the ext scheduling class.
+ */
+static struct task_struct *ext_server_pick_task(struct sched_dl_entity *dl_se)
+{
+	return pick_task_scx(dl_se->rq);
 }
 
 /*
@@ -7736,7 +7739,7 @@ void ext_server_init(struct rq *rq)
 
 	init_dl_entity(dl_se);
 
-	dl_server_init(dl_se, rq, ext_server_has_tasks, ext_server_pick_task);
+	dl_server_init(dl_se, rq, ext_server_has_tasks, ext_server_pick_task, ext_server_balance);
 }
 
 static const struct btf_kfunc_id_set scx_kfunc_set_any = {
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 7573baca9a85a..0c16944d43db8 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -8875,7 +8875,7 @@ void fair_server_init(struct rq *rq)
 
 	init_dl_entity(dl_se);
 
-	dl_server_init(dl_se, rq, fair_server_has_tasks, fair_server_pick_task);
+	dl_server_init(dl_se, rq, fair_server_has_tasks, fair_server_pick_task, NULL);
 }
 
 /*
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 1fbf4ffbcb208..a8615bdd6bdfa 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -384,7 +384,8 @@ extern void dl_server_start(struct sched_dl_entity *dl_se);
 extern void dl_server_stop(struct sched_dl_entity *dl_se);
 extern void dl_server_init(struct sched_dl_entity *dl_se, struct rq *rq,
 		    dl_server_has_tasks_f has_tasks,
-		    dl_server_pick_f pick_task);
+		    dl_server_pick_f pick_task,
+		    dl_server_balance_f balance);
 extern void sched_init_dl_servers(void);
 
 extern void dl_server_update_idle_time(struct rq *rq,
@@ -403,6 +404,11 @@ static inline bool dl_server_active(struct sched_dl_entity *dl_se)
 	return dl_se->dl_server_active;
 }
 
+static inline int on_dl_rq(struct sched_dl_entity *dl_se)
+{
+	return !RB_EMPTY_NODE(&dl_se->rb_node);
+}
+
 #ifdef CONFIG_CGROUP_SCHED
 
 extern struct list_head task_groups;
-- 
2.51.0