[RFC PATCH 3/7] psi: charge pressure to the task's active cgroup

[Shakeel Butt <shakeel.butt@linux.dev>]

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 <shakeel.butt@linux.dev>
---
 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