When a kernel thread works for a cgroup under set_active_cgroup(), its stalls, like memory reclaim, should count as that cgroup's pressure.
Make task_psi_group() use the active cgroup. When it changes, move the task's pressure state from the old groups to the new ones, the same way cgroup_move_task() does for a real cgroup move. Signed-off-by: Shakeel Butt <[email protected]> --- kernel/sched/core.c | 2 +- kernel/sched/psi.c | 22 +++++++++++++++++++++- kernel/sched/stats.h | 10 ++++++++++ 3 files changed, 32 insertions(+), 2 deletions(-) diff --git a/kernel/sched/core.c b/kernel/sched/core.c index ccb7ee9f2f26..b9e288b76da9 100644 --- a/kernel/sched/core.c +++ b/kernel/sched/core.c @@ -5768,7 +5768,7 @@ struct cgroup *set_active_cgroup(struct cgroup *cgrp) rq->donor->sched_class->update_curr(rq); old = p->active_cgroup; - p->active_cgroup = cgrp; + psi_set_active_cgroup(p, cgrp); task_rq_unlock(rq, p, &rf); return old; diff --git a/kernel/sched/psi.c b/kernel/sched/psi.c index 4e152410653d..7e03be26b8d3 100644 --- a/kernel/sched/psi.c +++ b/kernel/sched/psi.c @@ -886,7 +886,7 @@ static inline struct psi_group *task_psi_group(struct task_struct *task) { #ifdef CONFIG_CGROUPS if (static_branch_likely(&psi_cgroups_enabled)) - return cgroup_psi(task_dfl_cgroup(task)); + return cgroup_psi(task->active_cgroup ?: task_dfl_cgroup(task)); #endif return &psi_system; } @@ -1213,6 +1213,26 @@ void cgroup_move_task(struct task_struct *task, struct css_set *to) task_rq_unlock(rq, task, &rf); } +/* + * Set @task's active cgroup and move its pressure state along with it. + * The caller holds the task's rq lock. + */ +void psi_set_active_cgroup(struct task_struct *task, struct cgroup *cgrp) +{ + unsigned int task_flags = task->psi_flags; + + lockdep_assert_rq_held(task_rq(task)); + + if (!static_branch_likely(&psi_cgroups_enabled) || !task_flags) { + task->active_cgroup = cgrp; + return; + } + + psi_task_change(task, task_flags, 0); + task->active_cgroup = cgrp; + psi_task_change(task, 0, task_flags); +} + void psi_cgroup_restart(struct psi_group *group) { int cpu; diff --git a/kernel/sched/stats.h b/kernel/sched/stats.h index ebe0a7765f98..60f46b7b598c 100644 --- a/kernel/sched/stats.h +++ b/kernel/sched/stats.h @@ -103,6 +103,9 @@ __schedstats_from_se(struct sched_entity *se) void psi_task_change(struct task_struct *task, int clear, int set); void psi_task_switch(struct task_struct *prev, struct task_struct *next, bool sleep); +#ifdef CONFIG_CGROUPS +void psi_set_active_cgroup(struct task_struct *task, struct cgroup *cgrp); +#endif #ifdef CONFIG_IRQ_TIME_ACCOUNTING void psi_account_irqtime(struct rq *rq, struct task_struct *curr, struct task_struct *prev); #else /* !CONFIG_IRQ_TIME_ACCOUNTING: */ @@ -227,6 +230,13 @@ static inline void psi_sched_switch(struct task_struct *prev, bool sleep) {} static inline void psi_account_irqtime(struct rq *rq, struct task_struct *curr, struct task_struct *prev) {} +#ifdef CONFIG_CGROUPS +static inline void psi_set_active_cgroup(struct task_struct *task, + struct cgroup *cgrp) +{ + task->active_cgroup = cgrp; +} +#endif #endif /* !CONFIG_PSI */ #ifdef CONFIG_SCHED_INFO -- 2.53.0-Meta
