[RFC PATCH 2/7] cgroup: add set_active_cgroup() to charge CPU time to a cgroup

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

A kernel thread that does work for a cgroup, like memcg reclaim from a
workqueue, has that CPU time charged to its own cgroup, usually root.

Add task_struct::active_cgroup. When set, the task's CPU time goes to
that cgroup instead. set_active_cgroup() sets it and returns the old
value, like set_active_memcg(). It first charges the time used so far to
the old cgroup, so even short work is charged to the right place.

Code under set_active_cgroup() is kernel code, so charge its time as
system time right away. Skip the tick and vtime path for it, which could
charge it late, after the old value is back. Forced idle time from core
scheduling comes through that path too. It is not run time, so charge it
to the active cgroup there.

Only allow callers in the root cgroup, like kworkers. The scheduler still
runs the caller in its own cgroup and takes the time from that cgroup's
cpu.max, and on nohz_full CPUs vtime can post system time from inside the
scope to that cgroup after the old value is back. Warn and do nothing for
other callers.

Warn and clear it if a task exits with it still set.

Signed-off-by: Shakeel Butt <shakeel.butt@linux.dev>
---
 include/linux/cgroup.h | 26 +++++++++++++++++++++++++
 include/linux/sched.h  |  2 ++
 kernel/cgroup/cgroup.c |  4 ++++
 kernel/fork.c          |  4 ++++
 kernel/sched/core.c    | 43 ++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 79 insertions(+)

diff --git a/include/linux/cgroup.h b/include/linux/cgroup.h
index 7b19ded2578d..761bd0a81c67 100644
--- a/include/linux/cgroup.h
+++ b/include/linux/cgroup.h
@@ -879,6 +879,8 @@ static inline void cgroup_account_system_time(struct cgroup *cgrp,
 	}
 }
 
+struct cgroup *set_active_cgroup(struct cgroup *cgrp);
+
 static inline void cgroup_account_cputime(struct task_struct *task,
 					  u64 delta_exec)
 {
@@ -886,6 +888,12 @@ static inline void cgroup_account_cputime(struct task_struct *task,
 
 	cpuacct_charge(task, delta_exec);
 
+	/* Time spent under set_active_cgroup() is all kernel time. */
+	if (task->active_cgroup) {
+		cgroup_account_system_time(task->active_cgroup, delta_exec);
+		return;
+	}
+
 	cgrp = task_dfl_cgroup(task);
 	if (cgroup_parent(cgrp))
 		__cgroup_account_cputime(cgrp, delta_exec);
@@ -899,6 +907,20 @@ static inline void cgroup_account_cputime_field(struct task_struct *task,
 
 	cpuacct_account_field(task, index, delta_exec);
 
+	/*
+	 * cgroup_account_cputime() has charged the run time already. Forced
+	 * idle time from core scheduling is not run time: charge it here.
+	 */
+	if (task->active_cgroup) {
+#ifdef CONFIG_SCHED_CORE
+		if (index == CPUTIME_FORCEIDLE &&
+		    cgroup_parent(task->active_cgroup))
+			__cgroup_account_cputime_field(task->active_cgroup,
+						       index, delta_exec);
+#endif
+		return;
+	}
+
 	cgrp = task_dfl_cgroup(task);
 	if (cgroup_parent(cgrp))
 		__cgroup_account_cputime_field(cgrp, index, delta_exec);
@@ -913,6 +935,10 @@ static inline void cgroup_account_cputime_field(struct task_struct *task,
 						u64 delta_exec) {}
 static inline void cgroup_account_system_time(struct cgroup *cgrp,
 					      u64 delta_exec) {}
+static inline struct cgroup *set_active_cgroup(struct cgroup *cgrp)
+{
+	return NULL;
+}
 
 #endif	/* CONFIG_CGROUPS */
 
diff --git a/include/linux/sched.h b/include/linux/sched.h
index 5ce350a616e4..002941f60e88 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -1356,6 +1356,8 @@ struct task_struct {
 	struct css_set __rcu		*cgroups;
 	/* cg_list protected by css_set_lock and tsk->alloc_lock: */
 	struct list_head		cg_list;
+	/* If set, CPU time is charged here; see set_active_cgroup(): */
+	struct cgroup			*active_cgroup;
 #ifdef CONFIG_PREEMPT_RT
 	struct llist_node		cg_dead_lnode;
 #endif	/* CONFIG_PREEMPT_RT */
diff --git a/kernel/cgroup/cgroup.c b/kernel/cgroup/cgroup.c
index 804318ae160e..f51ed71c122f 100644
--- a/kernel/cgroup/cgroup.c
+++ b/kernel/cgroup/cgroup.c
@@ -7173,6 +7173,10 @@ void cgroup_task_exit(struct task_struct *tsk)
 	struct cgroup_subsys *ss;
 	int i;
 
+	/* The task forgot to undo set_active_cgroup(). */
+	if (WARN_ON_ONCE(tsk->active_cgroup))
+		set_active_cgroup(NULL);
+
 	/* see cgroup_post_fork() for details */
 	do_each_subsys_mask(ss, i, have_exit_callback) {
 		ss->exit(tsk);
diff --git a/kernel/fork.c b/kernel/fork.c
index 39be51048d77..044a4a5de8ef 100644
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -1006,6 +1006,10 @@ static struct task_struct *dup_task_struct(struct task_struct *orig, int node)
 	tsk->active_memcg = NULL;
 #endif
 
+#ifdef CONFIG_CGROUPS
+	tsk->active_cgroup = NULL;
+#endif
+
 #ifdef CONFIG_X86_BUS_LOCK_DETECT
 	tsk->reported_split_lock = 0;
 #endif
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 57dad4fc163e..ccb7ee9f2f26 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -5732,6 +5732,49 @@ unsigned long long task_sched_runtime(struct task_struct *p)
 	return ns;
 }
 
+#ifdef CONFIG_CGROUPS
+/**
+ * set_active_cgroup - charge current's CPU time to another cgroup
+ * @cgrp: the cgroup to charge, or NULL for current's own cgroup
+ *
+ * For kernel code that does work for a cgroup. The time it uses is charged
+ * to @cgrp as kernel time. Returns the old value, which the caller restores
+ * when done. @cgrp must stay alive until then.
+ *
+ * Only for callers in the root cgroup, like kworkers. The scheduler still
+ * runs the caller in its own cgroup, so from any other cgroup the time would
+ * also count against that cgroup's cpu.max, and on nohz_full CPUs some of
+ * it could show up in that cgroup's cpu.stat. It warns and does nothing for
+ * such a caller.
+ */
+struct cgroup *set_active_cgroup(struct cgroup *cgrp)
+{
+	struct task_struct *p = current;
+	struct cgroup *old;
+	struct rq_flags rf;
+	struct rq *rq;
+
+	WARN_ON_ONCE(!in_task());
+	WARN_ON_ONCE(cgrp && cgrp->root != &cgrp_dfl_root);
+
+	scoped_guard(rcu) {
+		if (cgrp && WARN_ON_ONCE(cgroup_parent(task_dfl_cgroup(p))))
+			return p->active_cgroup;
+	}
+
+	rq = task_rq_lock(p, &rf);
+	/* Charge the time used so far to the old cgroup. */
+	update_rq_clock(rq);
+	rq->donor->sched_class->update_curr(rq);
+
+	old = p->active_cgroup;
+	p->active_cgroup = cgrp;
+	task_rq_unlock(rq, p, &rf);
+
+	return old;
+}
+#endif
+
 static u64 cpu_resched_latency(struct rq *rq)
 {
 	int latency_warn_ms = READ_ONCE(sysctl_resched_latency_warn_ms);
-- 
2.53.0-Meta