[PATCH v5 4/4] sched/uclamp: Move uclamp_rq_{inc, dec}() from core.c

[Qais Yousef]

fair and RT throttling can dequeue tasks without going through core
leading to partial dequeue and corrupt rq uclamp state. Fix the problem
by pairing uclamp_rq_{inc, dec}() with {add, sub}_nr_running() so that
each class that cares about uclamp updates the rq state along with the
nr_running.

Remove the sched_delayed guards from uclamp_rq_{inc, dec}() since we now
react at the appropriate time to do the operations and no longer need to
care about this detail. With this gone, the flags parameter can be
dropped too.

Signed-off-by: Qais Yousef <qyousef@xxxxxxxxxxx>
---
 kernel/sched/core.c    | 33 +++++++++------------------------
 kernel/sched/ext/ext.c |  6 ++----
 kernel/sched/fair.c    |  6 ++----
 kernel/sched/rt.c      | 20 ++++++++++++--------
 kernel/sched/sched.h   |  8 ++++----
 5 files changed, 29 insertions(+), 44 deletions(-)

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index a2035d66270b9..bf690d27faf6e 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -1866,7 +1866,7 @@ static inline void uclamp_rq_dec_id(struct rq *rq, struct task_struct *p,
 	}
 }
 
-static inline void uclamp_rq_inc(struct rq *rq, struct task_struct *p, int flags)
+void uclamp_rq_inc(struct rq *rq, struct task_struct *p)
 {
 	enum uclamp_id clamp_id;
 
@@ -1879,13 +1879,6 @@ static inline void uclamp_rq_inc(struct rq *rq, struct task_struct *p, int flags
 	if (!uclamp_is_used())
 		return;
 
-	if (unlikely(!p->sched_class->uclamp_enabled))
-		return;
-
-	/* Only inc the delayed task which being woken up. */
-	if (p->se.sched_delayed && !(flags & ENQUEUE_DELAYED))
-		return;
-
 	for_each_clamp_id(clamp_id)
 		uclamp_rq_inc_id(rq, p, clamp_id);
 
@@ -1894,7 +1887,7 @@ static inline void uclamp_rq_inc(struct rq *rq, struct task_struct *p, int flags
 		uclamp_rq_clear_idle(rq);
 }
 
-static inline void uclamp_rq_dec(struct rq *rq, struct task_struct *p)
+void uclamp_rq_dec(struct rq *rq, struct task_struct *p)
 {
 	enum uclamp_id clamp_id;
 
@@ -1907,12 +1900,6 @@ static inline void uclamp_rq_dec(struct rq *rq, struct task_struct *p)
 	if (!uclamp_is_used())
 		return;
 
-	if (unlikely(!p->sched_class->uclamp_enabled))
-		return;
-
-	if (p->se.sched_delayed)
-		return;
-
 	for_each_clamp_id(clamp_id)
 		uclamp_rq_dec_id(rq, p, clamp_id);
 }
@@ -2151,8 +2138,6 @@ static void __init init_uclamp(void)
 }
 
 #else /* !CONFIG_UCLAMP_TASK: */
-static inline void uclamp_rq_inc(struct rq *rq, struct task_struct *p, int flags) { }
-static inline void uclamp_rq_dec(struct rq *rq, struct task_struct *p) { }
 static inline void uclamp_fork(struct task_struct *p) { }
 static inline void uclamp_post_fork(struct task_struct *p) { }
 static inline void init_uclamp(void) { }
@@ -2188,11 +2173,10 @@ void enqueue_task(struct rq *rq, struct task_struct *p, int flags)
 		update_rq_clock(rq);
 
 	/*
-	 * Can be before ->enqueue_task() because uclamp considers the
-	 * ENQUEUE_DELAYED task before its ->sched_delayed gets cleared
-	 * in ->enqueue_task().
+	 * Need to be careful what operations are done here, some classes like
+	 * fair and RT can enqueue tasks without going through core due to some
+	 * features like throttling.
 	 */
-	uclamp_rq_inc(rq, p, flags);
 
 	p->sched_class->enqueue_task(rq, p, flags);
 
@@ -2222,10 +2206,11 @@ inline bool dequeue_task(struct rq *rq, struct task_struct *p, int flags)
 	psi_dequeue(p, flags);
 
 	/*
-	 * Must be before ->dequeue_task() because ->dequeue_task() can 'fail'
-	 * and mark the task ->sched_delayed.
+	 * Need to be careful what operations are done here, some classes like
+	 * fair and RT can dequeue tasks without going through core due to some
+	 * features like throttling.
 	 */
-	uclamp_rq_dec(rq, p);
+
 	return p->sched_class->dequeue_task(rq, p, flags);
 }
 
diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 0403c46342813..fe0971cdba691 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -2174,6 +2174,7 @@ static void enqueue_task_scx(struct rq *rq, struct task_struct *p, int core_enq_
 	p->scx.flags |= SCX_TASK_QUEUED;
 	rq->scx.nr_running++;
 	add_nr_running(rq, 1);
+	uclamp_rq_inc(rq, p);
 
 	if (SCX_HAS_OP(sch, runnable) && !task_on_rq_migrating(p))
 		SCX_CALL_OP_TASK(sch, runnable, rq, p, enq_flags);
@@ -2319,6 +2320,7 @@ static bool dequeue_task_scx(struct rq *rq, struct task_struct *p, int core_deq_
 	p->scx.flags &= ~SCX_TASK_QUEUED;
 	rq->scx.nr_running--;
 	sub_nr_running(rq, 1);
+	uclamp_rq_dec(rq, p);
 
 	scx_dispatch_dequeue(rq, p);
 
@@ -5068,10 +5070,6 @@ DEFINE_SCHED_CLASS(ext) = {
 	.prio_changed		= prio_changed_scx,
 
 	.update_curr		= update_curr_scx,
-
-#ifdef CONFIG_UCLAMP_TASK
-	.uclamp_enabled		= 1,
-#endif
 };
 
 s32 scx_init_dsq(struct scx_dispatch_q *dsq, u64 dsq_id, struct scx_sched *sch)
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 0feec172fd37e..a518e28538a7f 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -8210,6 +8210,7 @@ enqueue_task_fair(struct rq *rq, struct task_struct *p, int flags)
 
 	/* At this point se is NULL and we are at root level*/
 	add_nr_running(rq, 1);
+	uclamp_rq_inc(rq, p);
 
 	/*
 	 * Since new tasks are assigned an initial util_avg equal to
@@ -8328,6 +8329,7 @@ static bool __dequeue_task(struct rq *rq, struct task_struct *p, int flags)
 		__dequeue_entity(cfs_rq, se);
 
 	sub_nr_running(rq, 1);
+	uclamp_rq_dec(rq, p);
 
 	/* balance early to pull high priority tasks */
 	if (unlikely(!was_sched_idle && sched_idle_rq(rq)))
@@ -15886,10 +15888,6 @@ DEFINE_SCHED_CLASS(fair) = {
 #ifdef CONFIG_SCHED_CORE
 	.task_is_throttled	= task_is_throttled_fair,
 #endif
-
-#ifdef CONFIG_UCLAMP_TASK
-	.uclamp_enabled		= 1,
-#endif
 };
 
 void print_cfs_stats(struct seq_file *m, int cpu)
diff --git a/kernel/sched/rt.c b/kernel/sched/rt.c
index cefb0a8d37e81..96d37a1ccc2a8 100644
--- a/kernel/sched/rt.c
+++ b/kernel/sched/rt.c
@@ -1132,8 +1132,12 @@ void inc_rt_tasks(struct sched_rt_entity *rt_se, struct rt_rq *rt_rq)
 	inc_rt_prio(rt_rq, prio);
 	inc_rt_group(rt_se, rt_rq);
 
-	if (!group_rt_rq(rt_se))
-		add_nr_running(rq_of_rt_rq(rt_rq), 1);
+	if (!group_rt_rq(rt_se)) {
+		struct rq *rq = rq_of_rt_rq(rt_rq);
+
+		add_nr_running(rq, 1);
+		uclamp_rq_inc(rq, rt_task_of(rt_se));
+	}
 }
 
 static inline
@@ -1147,8 +1151,12 @@ void dec_rt_tasks(struct sched_rt_entity *rt_se, struct rt_rq *rt_rq)
 	dec_rt_prio(rt_rq, rt_se_prio(rt_se));
 	dec_rt_group(rt_se, rt_rq);
 
-	if (!group_rt_rq(rt_se))
-		sub_nr_running(rq_of_rt_rq(rt_rq), 1);
+	if (!group_rt_rq(rt_se)) {
+		struct rq *rq = rq_of_rt_rq(rt_rq);
+
+		sub_nr_running(rq, 1);
+		uclamp_rq_dec(rq, rt_task_of(rt_se));
+	}
 }
 
 /*
@@ -2578,10 +2586,6 @@ DEFINE_SCHED_CLASS(rt) = {
 #ifdef CONFIG_SCHED_CORE
 	.task_is_throttled	= task_is_throttled_rt,
 #endif
-
-#ifdef CONFIG_UCLAMP_TASK
-	.uclamp_enabled		= 1,
-#endif
 };
 
 #ifdef CONFIG_RT_GROUP_SCHED
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 0cab55963657c..40876533defdd 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -2650,10 +2650,6 @@ enum snt_e {
 
 struct sched_class {
 
-#ifdef CONFIG_UCLAMP_TASK
-	int uclamp_enabled;
-#endif
-
 	/*
 	 * move_queued_task/activate_task/enqueue_task: rq->lock
 	 * ttwu_do_activate/activate_task/enqueue_task: rq->lock
@@ -3703,6 +3699,8 @@ static inline unsigned long cpu_util_rt(struct rq *rq)
 
 unsigned long uclamp_eff_value(struct task_struct *p, enum uclamp_id clamp_id);
 void init_uclamp_rq(struct rq *rq);
+void uclamp_rq_inc(struct rq *rq, struct task_struct *p);
+void uclamp_rq_dec(struct rq *rq, struct task_struct *p);
 
 /*
  * When uclamp is compiled in, the aggregation at rq level is 'turned off'
@@ -3812,6 +3810,8 @@ uclamp_eff_value(struct task_struct *p, enum uclamp_id clamp_id)
 }
 
 static inline void init_uclamp_rq(struct rq *rq) {}
+static inline void uclamp_rq_inc(struct rq *rq, struct task_struct *p) {}
+static inline void uclamp_rq_dec(struct rq *rq, struct task_struct *p) {}
 
 static inline bool uclamp_rq_is_capped(struct rq *rq) { return false; }
 
-- 
2.53.0