[RFC PATCH 0/7] cgroup: charge kernel work to the cgroup it is done for

[RFC PATCH 0/7] cgroup: charge kernel work to the cgroup it is done for

[Shakeel Butt]

Sometimes a kernel thread does work for a cgroup. The CPU time of that
work is charged to the thread's own cgroup, which is usually the root.
The cgroup that caused the work never sees it, and never pays for it.

This series lets a kernel thread say which cgroup it is working for.
That cgroup then sees the CPU time in its cpu.stat and the stalls in
its memory.pressure, and the CPU time comes out of its cpu.max quota.
The first user is the memcg reclaim that runs from high_work.


The problem
===========

When a memcg goes over memory.high, the task that charged the memory
is made to reclaim on its way back to user space. That does not work
when the charge happens in IRQ context, because the interrupted task
may have nothing to do with the memcg. For that case the kernel
queues memcg->high_work, and a kworker does the reclaim later.

The kworker runs in the root cgroup. So for that reclaim:
 - its CPU time does not show up in the memcg's cpu.stat;
 - its reclaim work does not show up in the memcg's memory.pressure;
 - its CPU time does not count against the memcg's cpu.max limit;
 - the rest of the system sees it as system overhead.

It is one case of a wider gap. For memory, a kernel thread can charge
an allocation to another memcg with set_active_memcg(). For IO there
is kthread_associate_blkcg(). For CPU time and pressure there is
nothing like it. And a kworker cannot just be moved into the cgroup:
kworkers are bound kthreads (PF_NO_SETAFFINITY), and cgroup refuses to
move those.

For our lock isolation work [5, 6], we plan to use async reclaim more
extensively, not just for charge requests from IRQ context, and thus
we want to make the CPU accounting of this work accurate.


Earlier attempts
================

This has come up before, with memory reclaim as one of the main
examples.

In 2019, Daniel Jordan posted "cgroup-aware unbound workqueues" [1].
It moved the worker into the work item's cgroup before running it, and
back to the root after. A move took about 1 us, but every move takes
cgroup_mutex and cgroup_threadgroup_rwsem, so moves cannot run in
parallel. The workers also showed up in cgroup.procs and blocked
rmdir, and cpuset did not work. It was not merged. In his reply [2],
Tejun said the direction for memory and IO is "remote charging, where
a kthread explicitly says who the specific io or allocation is for",
along with back-charging, and that "CPU doesn't have a backcharging
mechanism yet".

In 2021, Hao Lee asked for remote charging of CPU time [3]. His case
is ours: memcg reclaim from a workqueue "will steal cpu time from the
system level, which breaks the resource isolation". Tejun agreed [4]
that something like this is needed for CPU time spent in common code
for a specific cgroup, "e.g. memory reclaim, net packet rx", and said
he knew of no patches for it. Daniel said he was working on one.
Nothing has been merged since.


Options we looked at
====================

A. Keep the shared kworker, and only charge the cgroup.

  A1. Measure how long the reclaim ran, then charge that time to the
      memcg's cgroup.

  A2. Let the kworker say "charge this cgroup" while it works, the
      same way set_active_memcg() works for memory.

B. Also make the cgroup's CPU limits apply.

  B1. Run the kworker in the cgroup's scheduling group while it works,
      so cpu.weight and cpu.max apply to it.

  B2. Run the reclaim at full speed, then take its CPU time out of the
      cgroup's cpu.max quota afterwards (back-charging).

C. Do the reclaim inside the cgroup.

  C1. Record the overage as a debt on the memcg, and let the memcg's
      own tasks pay it the next time they charge memory or return to
      user space.

  C2. Give each memcg its own reclaim thread that lives in the cgroup.

  C3. Use a cgroup-aware workqueue, as in [1], or per-cgroup worker
      pools.


What we chose and why
=====================

We chose A2, together with B2 for cpu.max.

We do not throttle the reclaim itself. Adding limits to it is more
complicated and most probably unneeded, as we envision that we will
need concurrent background reclaimers instead of throttling in a
real-world environment. We are working on developing a system to
balance the rate of allocations/charges with the rate of reclaim, to
keep the system always running effectively.

In addition, reclaim takes sleeping locks, like i_mmap_rwsem and the
anon_vma lock in rmap walks, and fs locks in shrinkers. With a low
cgroup weight, the kworker can be preempted while it holds them, and
then tasks in other cgroups wait on it. It also hurts the workqueue. A
kworker that is runnable but not running still counts as running, and
the workqueue only spots CPU hogs by the CPU time they use. So other
work queued on that CPU's system_wq waits too.

Still, the reclaim should not be free CPU time on top of the cgroup's
limit. With A2 alone, the reclaim shows up in cpu.stat, but the
cgroup's own tasks still get their full cpu.max quota. B2 fixes that
without slowing the reclaim down. The kworker runs at full speed, and
afterwards its time is taken out of the cgroup's cpu.max quota, so the
cgroup's own tasks get less CPU time instead. This is the
back-charging Tejun described [2]. It only covers cpu.max.
Back-charging cpu.weight is future work.

We did not take C2 or C3. A high_work run asks for only 64 pages,
which is far too little to pay for moving a thread into a cgroup
through the global cgroup locks. A thread per memcg means thousands of
threads. cgroup v2 does not allow tasks in a non-leaf domain cgroup.
And a kernel thread in a cgroup shows up in cgroup.procs and blocks
rmdir.

C1 changes behaviour. memory.high would start slowing tasks down for
socket memory that arrived in softirq, and it would still need a
fallback for when no task in the memcg runs. That deserves its own
discussion.

Between A1 and A2: A1 is a few lines and easy to backport, but it only
fixes cpu.stat, it charges the time in one lump at the end, and it
only helps this one user. A2 charges the time as it is used, PSI can
follow it, and any kernel thread that works for a cgroup can use it.


How it works
============

task_struct gets an active_cgroup pointer. set_active_cgroup() sets it
and returns the old value, which the caller puts back when done:

	old = set_active_cgroup(memcg->css.cgroup);
	reclaim_high(memcg, MEMCG_CHARGE_BATCH, GFP_KERNEL);
	set_active_cgroup(old);

Three important details:

 - Before switching, set_active_cgroup() charges the time used so far
   to the old target, under the task's rq lock. A high_work run is
   often shorter than a tick. Without this, its time goes to whatever
   is set the next time the scheduler adds up run time, which is
   usually root again.

 - Code run under set_active_cgroup() is always kernel code. So its
   time is charged as system time right away, and the tick and vtime
   paths skip it. Without this, on nohz_full CPUs vtime posts kernel
   time late, after the old value is back, and system_usec comes out
   as 0. Forced idle time from core scheduling comes through the same
   path. It is not run time, so it goes to the active cgroup there.

 - PSI follows active_cgroup too. When it changes, the task's pressure
   state moves from the old groups to the new ones under the same rq
   lock, the same way cgroup_move_task() does for a real move.

For cpu.max, set_active_cgroup() also notes the task's run time at
each switch. When it switches away from a cgroup, it passes the time
used under that cgroup to cfs_bandwidth_charge():

 - The time is taken from the cpu.max pool of the cgroup and of each
   ancestor that has a limit. These are the same pools the cgroup's
   own tasks draw from.

 - The kworker is never throttled, because the work has already run.
   If a pool runs dry, the rest becomes debt, which is paid first out
   of the next periods' quota. So the cgroup's own tasks get less CPU
   time for a while.

 - The debt is capped at one period's quota, so one long piece of
   work cannot starve the cgroup for long. Time over the cap still
   shows up in cpu.stat. Writing cpu.max or cpu.max.burst clears the
   debt.

 - If neither the cgroup nor any ancestor has a limit, it returns
   right away.

The caller must keep the cgroup alive until it puts back the old
value. For high_work this holds: the memcg's css pins its cgroup, and
high_work is cancelled in mem_cgroup_css_free(). If a task exits with
active_cgroup still set, it gets a warning and the field is cleared.

The caller must be in the root cgroup, as kworkers are. The scheduler
still runs the caller in its own cgroup, so from any other cgroup the
time would also be taken out of that cgroup's cpu.max, and on nohz_full
CPUs vtime could post some system time to the wrong cgroup at a switch.
For such a caller, set_active_cgroup() warns and does nothing.

Which cgroup pays for high_work? The work belongs to the memcg that is
over memory.high. A charge in IRQ context walks up from the charging
memcg and queues the high_work of the first memcg it finds over
memory.high. So if a child makes the charge but only its parent is
over memory.high, the reclaim is charged to the parent: its cpu.stat,
its memory.pressure, and the cpu.max of the parent and of each
ancestor above it that has a limit. The child is not charged. Reclaim
in task context is different: the charging task reclaims on its way
back to user space, so the child and every ancestor above it are
charged.


The patches
===========

  1: A helper that charges kernel CPU time to a cgroup.
  2: set_active_cgroup().
  3: PSI following it.
  4: A helper that takes kernel CPU time out of a cgroup's cpu.max.
  5: set_active_cgroup() using it.
  6: high_work using set_active_cgroup().
  7: Selftests for high_work.


Testing
=======

 - End to end, with the setup above. Without the series, the memcg's
   cpu.stat and memory.pressure did not move. With it, the memcg's
   usage_usec matched the CPU time that high_work used, measured with
   bpftrace, to within 1%, all of it as system time, also on nohz_full.
   memory.pressure grew by the reclaim's stall time.

 - A test module that calls set_active_cgroup() directly. The charge
   matches the CPU time used to within 20 us in runs of 200 to 500 ms.
   It also covers nesting, use from a user task in the root cgroup,
   over 30,000 switches while userspace moves the thread between
   cgroups, a task that exits without putting the old value back, and
   rmdir of a cgroup that is still in use. A caller outside the root
   cgroup gets one warning, and its time stays with its own cgroup. It
   passes on normal, psi=0, cgroup_disable=pressure, nohz_full,
   lockdep+KASAN+UBSAN and KCSAN kernels. No pressure state is left
   behind, and KCSAN reports nothing in the new code.

 - Forced idle, on a core scheduling kernel with SMT siblings. The
   module burns 500 ms under set_active_cgroup() next to a busy loop
   with a different core scheduling cookie, so each forces the other's
   CPU idle. The active cgroup's core_sched.force_idle_usec grows by
   exactly the forced idle time charged to the module's thread, about
   250 ms per run.

 - The same module against cpu.max, with a CPU hog in a cgroup whose
   cpu.max is "20000 100000":
    - 100 pieces of 2 ms of work for the cgroup cost the hog 214 to
      233 ms of CPU time for 199 ms of work. The same holds with the
      limit on the parent and the hog in a sibling cgroup.
    - Five 300 ms pieces cost the hog 89 to 104 ms, about five
      periods' quota. When cpu.max is written after each piece, the
      debt is cleared and the hog loses nothing.
    - Without a limit, the hog loses nothing.
    - rmdir of a cgroup that has debt works.
   It passes on normal, lockdep+KASAN+UBSAN and KCSAN kernels.

 - A caller inside a limited cgroup, doing 200 ms of work for that same
   cgroup. set_active_cgroup() refuses it, so the work is charged once:
   a hog there loses 227 to 233 ms, the same as when the caller does
   not use set_active_cgroup() at all (211 to 232 ms). Without the
   root cgroup rule, it loses 412 to 423 ms, as the time is charged
   twice.

 - A stress test that creates, fills, kills and removes memcgs while
   high_work runs, with and without cpu.max limits on the memcgs and
   their parent, on normal and debug kernels. With limits, each run
   takes 4 to 5 seconds of high_work time out of cpu.max, in about
   2,100 pieces. All the dying memcgs get freed within 2 seconds.

 - The two new selftests. Without the series, the first fails and the
   second skips, because no reclaim time shows up at all. With only
   patches 1-3 and 6, the second fails. With the whole series, both
   pass, on normal and debug kernels.

 - Every patch builds, with and without cgroups, and with core
   scheduling. The whole series builds cleanly without PSI, without
   memcg, without cgroups, without CFS bandwidth control, without
   CGROUP_SCHED, on UP, with PREEMPT_RT, with SCHED_CORE, with
   nohz_full, on i386, with clang, on arm64, with allmodconfig, and
   with W=1 on the touched files.

 - No measurable change in perf bench sched pipe or messaging. The hot
   path gains one load and one branch.


Open questions
==============

 - When the memcg's own tasks are idle, the kworker is its only busy
   member while it reclaims. So memory.pressure "full" equals "some"
   for that time. Is that what we want? Tools like oomd will see it.

 - Is set_active_cgroup() the right name and place? It is limited to
   callers in the root cgroup. Other callers would need the vtime state
   split at each switch, and their own cpu.max left alone while the
   override is set. Is that worth doing?

 - Is one period's quota the right cap for the cpu.max debt? A larger
   cap makes the cgroup pay for more of a long run, but can stall its
   tasks for longer.

 - Should high_work charge the memcg that made the charge instead, as
   reclaim in task context does? That would mean queuing the charging
   memcg's own high_work, so several children of an over-high parent
   could reclaim at the same time, instead of one shared work item.

 - Other kernel threads that work for one cgroup could use
   set_active_cgroup() too. Network receive itself is harder, because
   the cgroup is only known after the work starts, as Daniel pointed
   out in the thread at [3].


[1] https://lore.kernel.org/20190605133650.28545-1-daniel.m.jordan@oracle.com/
[2] https://lore.kernel.org/20190605135319.GK374014@devbig004.ftw2.facebook.com/
[3] https://lore.kernel.org/60decdb6.1c69fb81.6130e.7642@mx.google.com/
[4] https://lore.kernel.org/YN+Sne76dhKBzV%2FR@mtj.duckdns.org/
[5] https://lore.kernel.org/20260921192559.2619635-1-shakeel.butt@linux.dev/
[6] https://lore.kernel.org/20260924155025.949998-1-shakeel.butt@linux.dev/

Shakeel Butt (7):
  cgroup: add cgroup_account_system_time()
  cgroup: add set_active_cgroup() to charge CPU time to a cgroup
  psi: charge pressure to the task's active cgroup
  sched/fair: add cfs_bandwidth_charge() for kernel work done for a
    cgroup
  cgroup: take set_active_cgroup() time out of the cgroup's cpu.max
  memcg: charge high_work reclaim to the memcg
  selftests: cgroup: check that high_work reclaim is charged to the
    memcg

 include/linux/cgroup.h                        |  39 ++
 include/linux/sched.h                         |   4 +
 kernel/cgroup/cgroup.c                        |   4 +
 kernel/fork.c                                 |   4 +
 kernel/sched/core.c                           |  52 +++
 kernel/sched/fair.c                           |  43 +++
 kernel/sched/psi.c                            |  22 +-
 kernel/sched/sched.h                          |   8 +
 kernel/sched/stats.h                          |  10 +
 mm/memcontrol.c                               |   5 +
 .../selftests/cgroup/test_memcontrol.c        | 332 ++++++++++++++++++
 11 files changed, 522 insertions(+), 1 deletion(-)


base-commit: a8c591ed6b672915e0be57843f943a2a723aff40
-- 
2.53.0-Meta

[RFC PATCH 1/7] cgroup: add cgroup_account_system_time()

[Shakeel Butt]

Add a helper that charges some kernel CPU time to a cgroup's cpu.stat.
The next patch uses it for the time a kernel thread spends working for
another cgroup.

Signed-off-by: Shakeel Butt <shakeel.butt@linux.dev>
---
 include/linux/cgroup.h | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/include/linux/cgroup.h b/include/linux/cgroup.h
index 2afb4cb2bb4f..7b19ded2578d 100644
--- a/include/linux/cgroup.h
+++ b/include/linux/cgroup.h
@@ -868,6 +868,17 @@ void __cgroup_account_cputime(struct cgroup *cgrp, u64 delta_exec);
 void __cgroup_account_cputime_field(struct cgroup *cgrp,
 				    enum cpu_usage_stat index, u64 delta_exec);
 
+/* Charge @delta_exec of kernel CPU time to @cgrp. */
+static inline void cgroup_account_system_time(struct cgroup *cgrp,
+					      u64 delta_exec)
+{
+	if (cgroup_parent(cgrp)) {
+		__cgroup_account_cputime(cgrp, delta_exec);
+		__cgroup_account_cputime_field(cgrp, CPUTIME_SYSTEM,
+					       delta_exec);
+	}
+}
+
 static inline void cgroup_account_cputime(struct task_struct *task,
 					  u64 delta_exec)
 {
@@ -900,6 +911,8 @@ static inline void cgroup_account_cputime(struct task_struct *task,
 static inline void cgroup_account_cputime_field(struct task_struct *task,
 						enum cpu_usage_stat index,
 						u64 delta_exec) {}
+static inline void cgroup_account_system_time(struct cgroup *cgrp,
+					      u64 delta_exec) {}
 
 #endif	/* CONFIG_CGROUPS */
 
-- 
2.53.0-Meta

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

[Shakeel Butt]

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

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

[Shakeel Butt]

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

[RFC PATCH 4/7] sched/fair: add cfs_bandwidth_charge() for kernel work done for a cgroup

[Shakeel Butt]

set_active_cgroup() charges a kernel thread's CPU time to the cgroup it
works for, but only in cpu.stat. The work still does not count against
that cgroup's cpu.max.

Add cfs_bandwidth_charge(). It takes the time out of the cpu.max pool of
the cgroup's task group and of each limited ancestor, the same way the
group's own run time is taken. The work itself is never throttled: it
has already run. Instead the group's own tasks get less time afterwards.

What the pool cannot cover now becomes debt, paid out of the next
refills. The debt is capped at one period's quota, so a burst of work
cannot starve the group for long. Changing cpu.max clears it.

Signed-off-by: Shakeel Butt <shakeel.butt@linux.dev>
---
 kernel/sched/core.c  |  1 +
 kernel/sched/fair.c  | 43 +++++++++++++++++++++++++++++++++++++++++++
 kernel/sched/sched.h |  8 ++++++++
 3 files changed, 52 insertions(+)

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index b9e288b76da9..a487da494795 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -9809,6 +9809,7 @@ static int tg_set_cfs_bandwidth(struct task_group *tg,
 		cfs_b->period = ns_to_ktime(period);
 		cfs_b->quota = quota;
 		cfs_b->burst = burst;
+		cfs_b->debt = 0;
 
 		__refill_cfs_bandwidth_runtime(cfs_b);
 
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 69145dda0df5..84250bd5caa2 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -6626,6 +6626,7 @@ static inline u64 sched_cfs_bandwidth_slice(void)
 void __refill_cfs_bandwidth_runtime(struct cfs_bandwidth *cfs_b)
 {
 	s64 runtime;
+	u64 pay;
 
 	if (unlikely(cfs_b->quota == RUNTIME_INF))
 		return;
@@ -6638,9 +6639,51 @@ void __refill_cfs_bandwidth_runtime(struct cfs_bandwidth *cfs_b)
 	}
 
 	cfs_b->runtime = min(cfs_b->runtime, cfs_b->quota + cfs_b->burst);
+
+	/* Pay back the kernel work charged by cfs_bandwidth_charge(). */
+	pay = min(cfs_b->runtime, cfs_b->debt);
+	cfs_b->runtime -= pay;
+	cfs_b->debt -= pay;
+
 	cfs_b->runtime_snap = cfs_b->runtime;
 }
 
+/*
+ * Kernel work used @delta of CPU time for @cgrp, see set_active_cgroup().
+ * Take it out of the quota of @cgrp's task group and of each ancestor with
+ * a limit, the same way the group's own run time is taken. What the pool
+ * cannot cover now becomes debt, paid out of the next refills. The debt is
+ * capped at one period's quota, so the work cannot starve the group for
+ * long.
+ */
+void cfs_bandwidth_charge(struct cgroup *cgrp, u64 delta)
+{
+	struct task_group *tg;
+
+	if (!cfs_bandwidth_used())
+		return;
+
+	guard(rcu)();
+	tg = css_tg(cgroup_e_css(cgrp, &cpu_cgrp_subsys));
+
+	/* No limit here or above. */
+	if (READ_ONCE(tg->cfs_bandwidth.hierarchical_quota) == RUNTIME_INF)
+		return;
+
+	for (; tg; tg = tg->parent) {
+		struct cfs_bandwidth *cfs_b = &tg->cfs_bandwidth;
+		u64 take;
+
+		guard(raw_spinlock_irqsave)(&cfs_b->lock);
+		if (cfs_b->quota == RUNTIME_INF)
+			continue;
+
+		take = min(delta, cfs_b->runtime);
+		cfs_b->runtime -= take;
+		cfs_b->debt = min(cfs_b->debt + delta - take, cfs_b->quota);
+	}
+}
+
 static inline struct cfs_bandwidth *tg_cfs_bandwidth(struct task_group *tg)
 {
 	return &tg->cfs_bandwidth;
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 6c3ad70e58b8..2d1adfd7ad37 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -456,6 +456,8 @@ struct cfs_bandwidth {
 	u64			runtime;
 	u64			burst;
 	u64			runtime_snap;
+	/* Kernel work charged by cfs_bandwidth_charge(), not paid yet: */
+	u64			debt;
 	s64			hierarchical_quota;
 
 	u8			idle;
@@ -618,6 +620,12 @@ static inline bool cfs_task_bw_constrained(struct task_struct *p) { return false
 
 #endif /* !CONFIG_CGROUP_SCHED */
 
+#ifdef CONFIG_CFS_BANDWIDTH
+void cfs_bandwidth_charge(struct cgroup *cgrp, u64 delta);
+#else
+static inline void cfs_bandwidth_charge(struct cgroup *cgrp, u64 delta) { }
+#endif
+
 /*
  * A weight of 0 or 1 can cause arithmetics problems.
  * A weight of a cfs_rq is the sum of weights of which entities
-- 
2.53.0-Meta

[RFC PATCH 5/7] cgroup: take set_active_cgroup() time out of the cgroup's cpu.max

[Shakeel Butt]

Kernel work done under set_active_cgroup() shows up in the cgroup's
cpu.stat, but the cgroup's tasks still get their full cpu.max quota.

Remember the task's run time at each set_active_cgroup(). When the
active cgroup changes, charge the time used under the old one to its
cpu.max quota with cfs_bandwidth_charge(). The caller is in the root
cgroup, which has no cpu.max, so no quota is charged twice.

Signed-off-by: Shakeel Butt <shakeel.butt@linux.dev>
---
 include/linux/sched.h |  2 ++
 kernel/sched/core.c   | 12 ++++++++++--
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/include/linux/sched.h b/include/linux/sched.h
index 002941f60e88..7ecea9cfa702 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -1358,6 +1358,8 @@ struct task_struct {
 	struct list_head		cg_list;
 	/* If set, CPU time is charged here; see set_active_cgroup(): */
 	struct cgroup			*active_cgroup;
+	/* se.sum_exec_runtime at the last set_active_cgroup(): */
+	u64				active_cgroup_start;
 #ifdef CONFIG_PREEMPT_RT
 	struct llist_node		cg_dead_lnode;
 #endif	/* CONFIG_PREEMPT_RT */
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index a487da494795..91fec6461e4f 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -5738,8 +5738,9 @@ unsigned long long task_sched_runtime(struct task_struct *p)
  * @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.
+ * to @cgrp as kernel time, and taken out of @cgrp's cpu.max quota. 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
@@ -5753,6 +5754,7 @@ struct cgroup *set_active_cgroup(struct cgroup *cgrp)
 	struct cgroup *old;
 	struct rq_flags rf;
 	struct rq *rq;
+	u64 used;
 
 	WARN_ON_ONCE(!in_task());
 	WARN_ON_ONCE(cgrp && cgrp->root != &cgrp_dfl_root);
@@ -5768,9 +5770,15 @@ struct cgroup *set_active_cgroup(struct cgroup *cgrp)
 	rq->donor->sched_class->update_curr(rq);
 
 	old = p->active_cgroup;
+	used = p->se.sum_exec_runtime - p->active_cgroup_start;
+	p->active_cgroup_start = p->se.sum_exec_runtime;
 	psi_set_active_cgroup(p, cgrp);
 	task_rq_unlock(rq, p, &rf);
 
+	/* Take that time out of the old cgroup's cpu.max quota as well. */
+	if (old)
+		cfs_bandwidth_charge(old, used);
+
 	return old;
 }
 #endif
-- 
2.53.0-Meta

[RFC PATCH 6/7] memcg: charge high_work reclaim to the memcg

[Shakeel Butt]

When a charge from IRQ context puts a memcg over memory.high, the
reclaim runs later from high_work on a kworker. The kworker is in the
root cgroup, so the reclaim does not show up in the memcg's cpu.stat
or memory.pressure, and does not count against its cpu.max.

Run the reclaim under set_active_cgroup(). Its CPU time and memory
stalls are then charged to the memcg's cgroup, and its CPU time comes
out of that cgroup's cpu.max quota.

Signed-off-by: Shakeel Butt <shakeel.butt@linux.dev>
---
 mm/memcontrol.c | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 4d00748c8a5b..f1d0ba8b48cb 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -2456,9 +2456,14 @@ static unsigned long reclaim_high(struct mem_cgroup *memcg,
 static void high_work_func(struct work_struct *work)
 {
 	struct mem_cgroup *memcg;
+	struct cgroup *old;
 
 	memcg = container_of(work, struct mem_cgroup, high_work);
+
+	/* Charge the reclaim to the memcg's cgroup, not to the root. */
+	old = set_active_cgroup(memcg->css.cgroup);
 	reclaim_high(memcg, MEMCG_CHARGE_BATCH, GFP_KERNEL);
+	set_active_cgroup(old);
 }
 
 static void high_irq_work_func(struct irq_work *work)
-- 
2.53.0-Meta

[RFC PATCH 7/7] selftests: cgroup: check that high_work reclaim is charged to the memcg

[Shakeel Butt]

Fill a memcg's TCP receive buffers from softirq while its only task
sleeps, so the memcg goes over memory.high and high_work reclaims.

test_memcg_high_work checks that the memcg's cpu.stat and
memory.pressure grow, which only happens if that reclaim is charged to
the memcg.

test_memcg_high_work_cpu_max also runs a CPU hog in the memcg under a
cpu.max limit. The reclaim time should come out of the memcg's quota,
so the hog should get less CPU while high_work is busy.

Both need swap, so the anon memory can be reclaimed.

Signed-off-by: Shakeel Butt <shakeel.butt@linux.dev>
---
 .../selftests/cgroup/test_memcontrol.c        | 332 ++++++++++++++++++
 1 file changed, 332 insertions(+)

diff --git a/tools/testing/selftests/cgroup/test_memcontrol.c b/tools/testing/selftests/cgroup/test_memcontrol.c
index 0ed82347044e..00802804c0d9 100644
--- a/tools/testing/selftests/cgroup/test_memcontrol.c
+++ b/tools/testing/selftests/cgroup/test_memcontrol.c
@@ -18,6 +18,7 @@
 #include <netdb.h>
 #include <errno.h>
 #include <sys/mman.h>
+#include <time.h>
 
 #include "kselftest.h"
 #include "cgroup_util.h"
@@ -1474,6 +1475,335 @@ static int test_memcg_sock(const char *root)
 	return ret;
 }
 
+#define HIGH_WORK_NCONN		16
+
+struct high_work_args {
+	int ctl[2];
+};
+
+/*
+ * Fill some anon memory, accept HIGH_WORK_NCONN connections, then sleep
+ * without ever reading, so all received data stays charged to the memcg.
+ */
+static int high_work_receiver(const char *cgroup, void *arg)
+{
+	struct high_work_args *args = arg;
+	struct sockaddr_in sa = { .sin_family = AF_INET };
+	socklen_t len = sizeof(sa);
+	int sk, i, rcvbuf = MB(8);
+	unsigned long long x = 1, *p;
+	size_t size = MB(256);
+
+	close(args->ctl[0]);
+
+	/* Hard to compress, so reclaiming it to swap takes real work. */
+	p = mmap(NULL, size, PROT_READ | PROT_WRITE,
+		 MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+	if (p == MAP_FAILED)
+		return -1;
+	for (size_t j = 0; j < size / sizeof(*p); j++) {
+		x = x * 6364136223846793005ULL + 1442695040888963407ULL;
+		p[j] = x;
+	}
+
+	sk = socket(AF_INET, SOCK_STREAM, 0);
+	if (sk < 0)
+		return -1;
+	if (setsockopt(sk, SOL_SOCKET, SO_RCVBUFFORCE, &rcvbuf, sizeof(rcvbuf)))
+		return -1;
+	sa.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
+	if (bind(sk, (struct sockaddr *)&sa, len) || listen(sk, HIGH_WORK_NCONN))
+		return -1;
+	if (getsockname(sk, (struct sockaddr *)&sa, &len))
+		return -1;
+	if (write(args->ctl[1], &sa.sin_port, sizeof(sa.sin_port)) !=
+	    sizeof(sa.sin_port))
+		return -1;
+
+	for (i = 0; i < HIGH_WORK_NCONN; i++)
+		if (accept(sk, NULL, NULL) < 0)
+			return -1;
+	if (write(args->ctl[1], "A", 1) != 1)
+		return -1;
+
+	for (;;)
+		pause();
+	return 0;
+}
+
+static long psi_some_total(const char *cgroup, const char *control)
+{
+	char buf[BUF_SIZE];
+	long total;
+
+	if (cg_read(cgroup, control, buf, sizeof(buf)))
+		return -1;
+	if (sscanf(buf, "some avg10=%*f avg60=%*f avg300=%*f total=%ld",
+		   &total) != 1)
+		return -1;
+	return total;
+}
+
+/*
+ * Socket memory charged from softirq can push a memcg over memory.high.
+ * The reclaim then runs from high_work on a kworker. The memcg's only task
+ * sleeps the whole time here, so its cpu.stat and memory.pressure only move
+ * if that reclaim is charged to the memcg.
+ */
+static int test_memcg_high_work(const char *root)
+{
+	int ret = KSFT_FAIL, pid = -1, i, sk[HIGH_WORK_NCONN];
+	long usage, high, some, current, sent = 0;
+	struct high_work_args args;
+	char *memcg, *buf = NULL, c;
+	in_port_t port;
+	time_t last;
+
+	for (i = 0; i < HIGH_WORK_NCONN; i++)
+		sk[i] = -1;
+
+	if (!is_swap_enabled())
+		return KSFT_SKIP;
+
+	memcg = cg_name(root, "memcg_test");
+	if (!memcg)
+		goto cleanup;
+	if (cg_create(memcg))
+		goto cleanup;
+	if (pipe(args.ctl))
+		goto cleanup;
+
+	pid = cg_run_nowait(memcg, high_work_receiver, &args);
+	if (pid < 0)
+		goto cleanup;
+	close(args.ctl[1]);
+	if (read(args.ctl[0], &port, sizeof(port)) != sizeof(port))
+		goto cleanup;
+
+	for (i = 0; i < HIGH_WORK_NCONN; i++) {
+		struct sockaddr_in sa = { .sin_family = AF_INET, .sin_port = port };
+
+		sa.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
+		sk[i] = socket(AF_INET, SOCK_STREAM, 0);
+		if (sk[i] < 0 || connect(sk[i], (struct sockaddr *)&sa, sizeof(sa)))
+			goto cleanup;
+	}
+	if (read(args.ctl[0], &c, 1) != 1)
+		goto cleanup;
+
+	/* Just above the current usage, so this write does not reclaim. */
+	current = cg_read_long(memcg, "memory.current");
+	if (current < 0 || cg_write_numeric(memcg, "memory.high", current + MB(2)))
+		goto cleanup;
+
+	usage = cg_read_key_long(memcg, "cpu.stat", "usage_usec ");
+	high = cg_read_key_long(memcg, "memory.events", "high ");
+	some = psi_some_total(memcg, "memory.pressure");
+
+	/* Fill all sockets until they stop taking data. */
+	buf = calloc(1, MB(1));
+	if (!buf)
+		goto cleanup;
+	last = time(NULL);
+	while (time(NULL) - last < 2) {
+		for (i = 0; i < HIGH_WORK_NCONN; i++) {
+			long n = send(sk[i], buf, MB(1), MSG_DONTWAIT);
+
+			if (n > 0) {
+				sent += n;
+				last = time(NULL);
+			}
+		}
+	}
+	sleep(2);
+
+	if (cg_read_key_long(memcg, "memory.events", "high ") <= high) {
+		ksft_print_msg("high_work did not run, sent %ld bytes\n", sent);
+		goto cleanup;
+	}
+	if (cg_read_key_long(memcg, "cpu.stat", "usage_usec ") <= usage) {
+		ksft_print_msg("high_work CPU time not charged to the memcg\n");
+		goto cleanup;
+	}
+	if (some >= 0 && psi_some_total(memcg, "memory.pressure") <= some) {
+		ksft_print_msg("high_work stall not in memcg memory.pressure\n");
+		goto cleanup;
+	}
+
+	ret = KSFT_PASS;
+
+cleanup:
+	for (i = 0; i < HIGH_WORK_NCONN; i++)
+		if (sk[i] >= 0)
+			close(sk[i]);
+	if (pid > 0) {
+		kill(pid, SIGKILL);
+		waitpid(pid, NULL, 0);
+	}
+	free(buf);
+	cg_destroy(memcg);
+	free(memcg);
+
+	return ret;
+}
+
+static int high_work_spinner(const char *cgroup, void *arg)
+{
+	for (;;)
+		;
+	return 0;
+}
+
+/* CPU time @pid has used so far, in microseconds. */
+static long task_cpu_usec(int pid)
+{
+	unsigned long long ns;
+	char buf[128];
+
+	if (proc_read_text(pid, false, "schedstat", buf, sizeof(buf)) <= 0 ||
+	    sscanf(buf, "%llu", &ns) != 1)
+		return -1;
+	return ns / 1000;
+}
+
+static long now_usec(void)
+{
+	struct timespec ts;
+
+	clock_gettime(CLOCK_MONOTONIC, &ts);
+	return ts.tv_sec * 1000000L + ts.tv_nsec / 1000;
+}
+
+/*
+ * The setup of test_memcg_high_work, plus a CPU hog in the memcg and a
+ * cpu.max limit. The reclaim that high_work does for the memcg should come
+ * out of the memcg's CPU quota, so the hog should get less CPU while
+ * high_work is busy than it gets otherwise.
+ */
+static int test_memcg_high_work_cpu_max(const char *root)
+{
+	int ret = KSFT_FAIL, pid = -1, hog = -1, i, sk[HIGH_WORK_NCONN];
+	long current, h0, h1, h2, t0, t1, t2, u1, u2, reclaim, lost;
+	bool cpu_enabled = false;
+	struct high_work_args args;
+	char *memcg = NULL, *buf = NULL, c;
+	in_port_t port;
+	time_t last;
+
+	for (i = 0; i < HIGH_WORK_NCONN; i++)
+		sk[i] = -1;
+
+	if (!is_swap_enabled())
+		return KSFT_SKIP;
+	if (cg_read_strstr(root, "cgroup.controllers", "cpu"))
+		return KSFT_SKIP;
+	if (cg_read_strstr(root, "cgroup.subtree_control", "cpu")) {
+		if (cg_write(root, "cgroup.subtree_control", "+cpu"))
+			return KSFT_SKIP;
+		cpu_enabled = true;
+	}
+
+	memcg = cg_name(root, "memcg_test");
+	if (!memcg || cg_create(memcg))
+		goto cleanup;
+	/* 20% of a CPU, in short periods so the numbers come out smooth. */
+	if (cg_write(memcg, "cpu.max", "2000 10000"))
+		goto cleanup;
+	if (pipe(args.ctl))
+		goto cleanup;
+
+	pid = cg_run_nowait(memcg, high_work_receiver, &args);
+	if (pid < 0)
+		goto cleanup;
+	close(args.ctl[1]);
+	if (read(args.ctl[0], &port, sizeof(port)) != sizeof(port))
+		goto cleanup;
+
+	for (i = 0; i < HIGH_WORK_NCONN; i++) {
+		struct sockaddr_in sa = { .sin_family = AF_INET, .sin_port = port };
+
+		sa.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
+		sk[i] = socket(AF_INET, SOCK_STREAM, 0);
+		if (sk[i] < 0 || connect(sk[i], (struct sockaddr *)&sa, sizeof(sa)))
+			goto cleanup;
+	}
+	if (read(args.ctl[0], &c, 1) != 1)
+		goto cleanup;
+
+	hog = cg_run_nowait(memcg, high_work_spinner, NULL);
+	if (hog < 0)
+		goto cleanup;
+
+	current = cg_read_long(memcg, "memory.current");
+	if (current < 0 || cg_write_numeric(memcg, "memory.high", current + MB(2)))
+		goto cleanup;
+
+	/* How much CPU the hog gets without high_work. */
+	sleep(1);
+	h0 = task_cpu_usec(hog);
+	t0 = now_usec();
+	sleep(3);
+	h1 = task_cpu_usec(hog);
+	t1 = now_usec();
+	u1 = cg_read_key_long(memcg, "cpu.stat", "usage_usec ");
+
+	/* Fill the sockets slowly, so high_work runs many times. */
+	buf = calloc(1, 64 << 10);
+	if (!buf)
+		goto cleanup;
+	last = time(NULL);
+	while (time(NULL) - last < 2) {
+		for (i = 0; i < HIGH_WORK_NCONN; i++)
+			if (send(sk[i], buf, 64 << 10, MSG_DONTWAIT) > 0)
+				last = time(NULL);
+		usleep(15000);
+	}
+	h2 = task_cpu_usec(hog);
+	t2 = now_usec();
+	u2 = cg_read_key_long(memcg, "cpu.stat", "usage_usec ");
+
+	if (h0 < 0 || h1 <= h0 || h2 < 0 || u1 < 0 || u2 < 0)
+		goto cleanup;
+
+	/* The receiver sleeps, so the rest of the memcg's time is high_work. */
+	reclaim = (u2 - u1) - (h2 - h1);
+	lost = (h1 - h0) * (t2 - t1) / (t1 - t0) - (h2 - h1);
+	if (reclaim < 30000) {
+		ksft_print_msg("high_work used only %ld us, too little to test\n",
+			       reclaim);
+		ret = KSFT_SKIP;
+		goto cleanup;
+	}
+	if (lost < reclaim / 3) {
+		ksft_print_msg("hog lost %ld us, high_work used %ld us\n",
+			       lost, reclaim);
+		goto cleanup;
+	}
+
+	ret = KSFT_PASS;
+
+cleanup:
+	for (i = 0; i < HIGH_WORK_NCONN; i++)
+		if (sk[i] >= 0)
+			close(sk[i]);
+	if (hog > 0) {
+		kill(hog, SIGKILL);
+		waitpid(hog, NULL, 0);
+	}
+	if (pid > 0) {
+		kill(pid, SIGKILL);
+		waitpid(pid, NULL, 0);
+	}
+	free(buf);
+	if (memcg)
+		cg_destroy(memcg);
+	free(memcg);
+	if (cpu_enabled)
+		cg_write(root, "cgroup.subtree_control", "-cpu");
+
+	return ret;
+}
+
 /*
  * This test disables swapping and tries to allocate anonymous memory
  * up to OOM with memory.group.oom set. Then it checks that all
@@ -1780,6 +2110,8 @@ struct memcg_test {
 	T(test_memcg_oom_events),
 	T(test_memcg_swap_max_peak),
 	T(test_memcg_sock),
+	T(test_memcg_high_work),
+	T(test_memcg_high_work_cpu_max),
 	T(test_memcg_oom_group_leaf_events),
 	T(test_memcg_oom_group_parent_events),
 	T(test_memcg_oom_group_score_events),
-- 
2.53.0-Meta

Re: [RFC PATCH 0/7] cgroup: charge kernel work to the cgroup it is done for

[Tejun Heo]

Hello, Shakeel.

On Thu, Sep 24, 2026 at 11:47:04AM -0700, Shakeel Butt wrote:
> This series lets a kernel thread say which cgroup it is working for.
> That cgroup then sees the CPU time in its cpu.stat and the stalls in
> its memory.pressure, and the CPU time comes out of its cpu.max quota.
> The first user is the memcg reclaim that runs from high_work.

This doesn't translate to net rx, which is another major source of
displaced CPU usage. Switching membership on each packet isn't going to
work there. Attribution can't happen that way. We'd much rather count
per-cgroup received packets and prorate the CPU consumption. If at all
possible, I think it'd be better to adopt an approach which can cover
both use cases.

>  - The debt is capped at one period's quota, so one long piece of
>    work cannot starve the cgroup for long. Time over the cap still
>    shows up in cpu.stat. Writing cpu.max or cpu.max.burst clears the
>    debt.

I don't like the debt capping. Having debt doesn't have to mean that
the cgroup doesn't get any bandwidth at all. The cgroup just needs to
be slowed down enough that the generation of new work is throttled and
the whole thing doesn't go out of control. IO control already does
this: when IO debt is accumulated, userspace is heavily throttled, but
not completely stalled, until the whole cgroup's consumption comes
under control. I don't see why the debts would need to be forgiven
unconditionally. The cgroup can keep paying them while running at a
minimal rate to avoid triggering stall failures, and if the situation
doesn't resolve quickly, that will most likely trigger pressure based
kills in any reasonable setup anyway.

Thanks.

--
tejun

Re: [RFC PATCH 0/7] cgroup: charge kernel work to the cgroup it is done for

[Shakeel Butt]

On Thu, Sep 24, 2026 at 10:28:10AM -1000, Tejun Heo wrote:
> Hello, Shakeel.
> 
> On Thu, Sep 24, 2026 at 11:47:04AM -0700, Shakeel Butt wrote:
> > This series lets a kernel thread say which cgroup it is working for.
> > That cgroup then sees the CPU time in its cpu.stat and the stalls in
> > its memory.pressure, and the CPU time comes out of its cpu.max quota.
> > The first user is the memcg reclaim that runs from high_work.
> 
> This doesn't translate to net rx, which is another major source of
> displaced CPU usage. Switching membership on each packet isn't going to
> work there. Attribution can't happen that way. We'd much rather count
> per-cgroup received packets and prorate the CPU consumption. If at all
> possible, I think it'd be better to adopt an approach which can cover
> both use cases.

Very good point. I think we need to think for net rx for two scenarios. First,
the modern NICs with rx steering support and second, old NIC with shared queues.

On the modern NICs where the workloads get their own rx queues, I think the
proposed mechanism can help to do the accurate accounting (I have to extend this
to softirqs as it is limited to kthreads atm) very easily.

The challenge you mentioned is for the old NICs where we get to know the cgroup
assosiation of the rx packets very late (or deep) in the stack. Also CPU spent
on each packet is not necessarily uniform (due to out-of-order or drops or
checksum errors or window shrinking) but for simplicity we can assume uniform
CPU. Maybe the right place to charge for such scenario might be at application
receiving those packets into the memory. I feel like we might need a very
special way to account for this case (maybe through BPF or something). At
the moment to me it seems very hard to have a universal solution which helps
this case and the reclaim case I am targetted.

If you don't mind, I think having solution for modern NICs i.e. dedicated rx
queues, should suffice for now (unless you want the solution for old NICs as
well). Let me know what you think.

> 
> >  - The debt is capped at one period's quota, so one long piece of
> >    work cannot starve the cgroup for long. Time over the cap still
> >    shows up in cpu.stat. Writing cpu.max or cpu.max.burst clears the
> >    debt.
> 
> I don't like the debt capping. Having debt doesn't have to mean that
> the cgroup doesn't get any bandwidth at all. The cgroup just needs to
> be slowed down enough that the generation of new work is throttled and
> the whole thing doesn't go out of control. IO control already does
> this: when IO debt is accumulated, userspace is heavily throttled, but
> not completely stalled, until the whole cgroup's consumption comes
> under control. I don't see why the debts would need to be forgiven
> unconditionally. The cgroup can keep paying them while running at a
> minimal rate to avoid triggering stall failures, and if the situation
> doesn't resolve quickly, that will most likely trigger pressure based
> kills in any reasonable setup anyway.

Sounds good, I will remove this capping in the next version.

Thanks for taking a look.

Re: [RFC PATCH 0/7] cgroup: charge kernel work to the cgroup it is done for

[Shakeel Butt]

Somehow the git send-email got killed in between. Please ignore this.

The full RFC is at
https://lore.kernel.org/20260924184714.912181-1-shakeel.butt@linux.dev

[RFC PATCH 0/7] cgroup: charge kernel work to the cgroup it is done for

[Shakeel Butt]

Sometimes a kernel thread does work for a cgroup. The CPU time of that
work is charged to the thread's own cgroup, which is usually the root.
The cgroup that caused the work never sees it, and never pays for it.

This series lets a kernel thread say which cgroup it is working for.
That cgroup then sees the CPU time in its cpu.stat and the stalls in
its memory.pressure, and the CPU time comes out of its cpu.max quota.
The first user is the memcg reclaim that runs from high_work.


The problem
===========

When a memcg goes over memory.high, the task that charged the memory
is made to reclaim on its way back to user space. That does not work
when the charge happens in IRQ context, because the interrupted task
may have nothing to do with the memcg. For that case the kernel
queues memcg->high_work, and a kworker does the reclaim later.

The kworker runs in the root cgroup. So for that reclaim:
 - its CPU time does not show up in the memcg's cpu.stat;
 - its reclaim work does not show up in the memcg's memory.pressure;
 - its CPU time does not count against the memcg's cpu.max limit;
 - the rest of the system sees it as system overhead.

It is one case of a wider gap. For memory, a kernel thread can charge
an allocation to another memcg with set_active_memcg(). For IO there
is kthread_associate_blkcg(). For CPU time and pressure there is
nothing like it. And a kworker cannot just be moved into the cgroup:
kworkers are bound kthreads (PF_NO_SETAFFINITY), and cgroup refuses to
move those.

For our lock isolation work [5, 6], we plan to use async reclaim more
extensively, not just for charge requests from IRQ context, and thus
we want to make the CPU accounting of this work accurate.


Earlier attempts
================

This has come up before, with memory reclaim as one of the main
examples.

In 2019, Daniel Jordan posted "cgroup-aware unbound workqueues" [1].
It moved the worker into the work item's cgroup before running it, and
back to the root after. A move took about 1 us, but every move takes
cgroup_mutex and cgroup_threadgroup_rwsem, so moves cannot run in
parallel. The workers also showed up in cgroup.procs and blocked
rmdir, and cpuset did not work. It was not merged. In his reply [2],
Tejun said the direction for memory and IO is "remote charging, where
a kthread explicitly says who the specific io or allocation is for",
along with back-charging, and that "CPU doesn't have a backcharging
mechanism yet".

In 2021, Hao Lee asked for remote charging of CPU time [3]. His case
is ours: memcg reclaim from a workqueue "will steal cpu time from the
system level, which breaks the resource isolation". Tejun agreed [4]
that something like this is needed for CPU time spent in common code
for a specific cgroup, "e.g. memory reclaim, net packet rx", and said
he knew of no patches for it. Daniel said he was working on one.
Nothing has been merged since.


Options we looked at
====================

A. Keep the shared kworker, and only charge the cgroup.

  A1. Measure how long the reclaim ran, then charge that time to the
      memcg's cgroup.

  A2. Let the kworker say "charge this cgroup" while it works, the
      same way set_active_memcg() works for memory.

B. Also make the cgroup's CPU limits apply.

  B1. Run the kworker in the cgroup's scheduling group while it works,
      so cpu.weight and cpu.max apply to it.

  B2. Run the reclaim at full speed, then take its CPU time out of the
      cgroup's cpu.max quota afterwards (back-charging).

C. Do the reclaim inside the cgroup.

  C1. Record the overage as a debt on the memcg, and let the memcg's
      own tasks pay it the next time they charge memory or return to
      user space.

  C2. Give each memcg its own reclaim thread that lives in the cgroup.

  C3. Use a cgroup-aware workqueue, as in [1], or per-cgroup worker
      pools.


What we chose and why
=====================

We chose A2, together with B2 for cpu.max.

We do not throttle the reclaim itself. Adding limits to it is more
complicated and most probably unneeded, as we envision that we will
need concurrent background reclaimers instead of throttling in a
real-world environment. We are working on developing a system to
balance the rate of allocations/charges with the rate of reclaim, to
keep the system always running effectively.

In addition, reclaim takes sleeping locks, like i_mmap_rwsem and the
anon_vma lock in rmap walks, and fs locks in shrinkers. With a low
cgroup weight, the kworker can be preempted while it holds them, and
then tasks in other cgroups wait on it. It also hurts the workqueue. A
kworker that is runnable but not running still counts as running, and
the workqueue only spots CPU hogs by the CPU time they use. So other
work queued on that CPU's system_wq waits too.

Still, the reclaim should not be free CPU time on top of the cgroup's
limit. With A2 alone, the reclaim shows up in cpu.stat, but the
cgroup's own tasks still get their full cpu.max quota. B2 fixes that
without slowing the reclaim down. The kworker runs at full speed, and
afterwards its time is taken out of the cgroup's cpu.max quota, so the
cgroup's own tasks get less CPU time instead. This is the
back-charging Tejun described [2]. It only covers cpu.max.
Back-charging cpu.weight is future work.

We did not take C2 or C3. A high_work run asks for only 64 pages,
which is far too little to pay for moving a thread into a cgroup
through the global cgroup locks. A thread per memcg means thousands of
threads. cgroup v2 does not allow tasks in a non-leaf domain cgroup.
And a kernel thread in a cgroup shows up in cgroup.procs and blocks
rmdir.

C1 changes behaviour. memory.high would start slowing tasks down for
socket memory that arrived in softirq, and it would still need a
fallback for when no task in the memcg runs. That deserves its own
discussion.

Between A1 and A2: A1 is a few lines and easy to backport, but it only
fixes cpu.stat, it charges the time in one lump at the end, and it
only helps this one user. A2 charges the time as it is used, PSI can
follow it, and any kernel thread that works for a cgroup can use it.


How it works
============

task_struct gets an active_cgroup pointer. set_active_cgroup() sets it
and returns the old value, which the caller puts back when done:

	old = set_active_cgroup(memcg->css.cgroup);
	reclaim_high(memcg, MEMCG_CHARGE_BATCH, GFP_KERNEL);
	set_active_cgroup(old);

Three important details:

 - Before switching, set_active_cgroup() charges the time used so far
   to the old target, under the task's rq lock. A high_work run is
   often shorter than a tick. Without this, its time goes to whatever
   is set the next time the scheduler adds up run time, which is
   usually root again.

 - Code run under set_active_cgroup() is always kernel code. So its
   time is charged as system time right away, and the tick and vtime
   paths skip it. Without this, on nohz_full CPUs vtime posts kernel
   time late, after the old value is back, and system_usec comes out
   as 0. Forced idle time from core scheduling comes through the same
   path. It is not run time, so it goes to the active cgroup there.

 - PSI follows active_cgroup too. When it changes, the task's pressure
   state moves from the old groups to the new ones under the same rq
   lock, the same way cgroup_move_task() does for a real move.

For cpu.max, set_active_cgroup() also notes the task's run time at
each switch. When it switches away from a cgroup, it passes the time
used under that cgroup to cfs_bandwidth_charge():

 - The time is taken from the cpu.max pool of the cgroup and of each
   ancestor that has a limit. These are the same pools the cgroup's
   own tasks draw from.

 - The kworker is never throttled, because the work has already run.
   If a pool runs dry, the rest becomes debt, which is paid first out
   of the next periods' quota. So the cgroup's own tasks get less CPU
   time for a while.

 - The debt is capped at one period's quota, so one long piece of
   work cannot starve the cgroup for long. Time over the cap still
   shows up in cpu.stat. Writing cpu.max or cpu.max.burst clears the
   debt.

 - If neither the cgroup nor any ancestor has a limit, it returns
   right away.

The caller must keep the cgroup alive until it puts back the old
value. For high_work this holds: the memcg's css pins its cgroup, and
high_work is cancelled in mem_cgroup_css_free(). If a task exits with
active_cgroup still set, it gets a warning and the field is cleared.

The caller must be in the root cgroup, as kworkers are. The scheduler
still runs the caller in its own cgroup, so from any other cgroup the
time would also be taken out of that cgroup's cpu.max, and on nohz_full
CPUs vtime could post some system time to the wrong cgroup at a switch.
For such a caller, set_active_cgroup() warns and does nothing.

Which cgroup pays for high_work? The work belongs to the memcg that is
over memory.high. A charge in IRQ context walks up from the charging
memcg and queues the high_work of the first memcg it finds over
memory.high. So if a child makes the charge but only its parent is
over memory.high, the reclaim is charged to the parent: its cpu.stat,
its memory.pressure, and the cpu.max of the parent and of each
ancestor above it that has a limit. The child is not charged. Reclaim
in task context is different: the charging task reclaims on its way
back to user space, so the child and every ancestor above it are
charged.


The patches
===========

  1: A helper that charges kernel CPU time to a cgroup.
  2: set_active_cgroup().
  3: PSI following it.
  4: A helper that takes kernel CPU time out of a cgroup's cpu.max.
  5: set_active_cgroup() using it.
  6: high_work using set_active_cgroup().
  7: Selftests for high_work.


Testing
=======

 - End to end, with the setup above. Without the series, the memcg's
   cpu.stat and memory.pressure did not move. With it, the memcg's
   usage_usec matched the CPU time that high_work used, measured with
   bpftrace, to within 1%, all of it as system time, also on nohz_full.
   memory.pressure grew by the reclaim's stall time.

 - A test module that calls set_active_cgroup() directly. The charge
   matches the CPU time used to within 20 us in runs of 200 to 500 ms.
   It also covers nesting, use from a user task in the root cgroup,
   over 30,000 switches while userspace moves the thread between
   cgroups, a task that exits without putting the old value back, and
   rmdir of a cgroup that is still in use. A caller outside the root
   cgroup gets one warning, and its time stays with its own cgroup. It
   passes on normal, psi=0, cgroup_disable=pressure, nohz_full,
   lockdep+KASAN+UBSAN and KCSAN kernels. No pressure state is left
   behind, and KCSAN reports nothing in the new code.

 - Forced idle, on a core scheduling kernel with SMT siblings. The
   module burns 500 ms under set_active_cgroup() next to a busy loop
   with a different core scheduling cookie, so each forces the other's
   CPU idle. The active cgroup's core_sched.force_idle_usec grows by
   exactly the forced idle time charged to the module's thread, about
   250 ms per run.

 - The same module against cpu.max, with a CPU hog in a cgroup whose
   cpu.max is "20000 100000":
    - 100 pieces of 2 ms of work for the cgroup cost the hog 214 to
      233 ms of CPU time for 199 ms of work. The same holds with the
      limit on the parent and the hog in a sibling cgroup.
    - Five 300 ms pieces cost the hog 89 to 104 ms, about five
      periods' quota. When cpu.max is written after each piece, the
      debt is cleared and the hog loses nothing.
    - Without a limit, the hog loses nothing.
    - rmdir of a cgroup that has debt works.
   It passes on normal, lockdep+KASAN+UBSAN and KCSAN kernels.

 - A caller inside a limited cgroup, doing 200 ms of work for that same
   cgroup. set_active_cgroup() refuses it, so the work is charged once:
   a hog there loses 227 to 233 ms, the same as when the caller does
   not use set_active_cgroup() at all (211 to 232 ms). Without the
   root cgroup rule, it loses 412 to 423 ms, as the time is charged
   twice.

 - A stress test that creates, fills, kills and removes memcgs while
   high_work runs, with and without cpu.max limits on the memcgs and
   their parent, on normal and debug kernels. With limits, each run
   takes 4 to 5 seconds of high_work time out of cpu.max, in about
   2,100 pieces. All the dying memcgs get freed within 2 seconds.

 - The two new selftests. Without the series, the first fails and the
   second skips, because no reclaim time shows up at all. With only
   patches 1-3 and 6, the second fails. With the whole series, both
   pass, on normal and debug kernels.

 - Every patch builds, with and without cgroups, and with core
   scheduling. The whole series builds cleanly without PSI, without
   memcg, without cgroups, without CFS bandwidth control, without
   CGROUP_SCHED, on UP, with PREEMPT_RT, with SCHED_CORE, with
   nohz_full, on i386, with clang, on arm64, with allmodconfig, and
   with W=1 on the touched files.

 - No measurable change in perf bench sched pipe or messaging. The hot
   path gains one load and one branch.


Open questions
==============

 - When the memcg's own tasks are idle, the kworker is its only busy
   member while it reclaims. So memory.pressure "full" equals "some"
   for that time. Is that what we want? Tools like oomd will see it.

 - Is set_active_cgroup() the right name and place? It is limited to
   callers in the root cgroup. Other callers would need the vtime state
   split at each switch, and their own cpu.max left alone while the
   override is set. Is that worth doing?

 - Is one period's quota the right cap for the cpu.max debt? A larger
   cap makes the cgroup pay for more of a long run, but can stall its
   tasks for longer.

 - Should high_work charge the memcg that made the charge instead, as
   reclaim in task context does? That would mean queuing the charging
   memcg's own high_work, so several children of an over-high parent
   could reclaim at the same time, instead of one shared work item.

 - Other kernel threads that work for one cgroup could use
   set_active_cgroup() too. Network receive itself is harder, because
   the cgroup is only known after the work starts, as Daniel pointed
   out in the thread at [3].


[1] https://lore.kernel.org/20190605133650.28545-1-daniel.m.jordan@oracle.com/
[2] https://lore.kernel.org/20190605135319.GK374014@devbig004.ftw2.facebook.com/
[3] https://lore.kernel.org/60decdb6.1c69fb81.6130e.7642@mx.google.com/
[4] https://lore.kernel.org/YN+Sne76dhKBzV%2FR@mtj.duckdns.org/
[5] https://lore.kernel.org/20260921192559.2619635-1-shakeel.butt@linux.dev/
[6] https://lore.kernel.org/20260924155025.949998-1-shakeel.butt@linux.dev/

Shakeel Butt (7):
  cgroup: add cgroup_account_system_time()
  cgroup: add set_active_cgroup() to charge CPU time to a cgroup
  psi: charge pressure to the task's active cgroup
  sched/fair: add cfs_bandwidth_charge() for kernel work done for a
    cgroup
  cgroup: take set_active_cgroup() time out of the cgroup's cpu.max
  memcg: charge high_work reclaim to the memcg
  selftests: cgroup: check that high_work reclaim is charged to the
    memcg

 include/linux/cgroup.h                        |  39 ++
 include/linux/sched.h                         |   4 +
 kernel/cgroup/cgroup.c                        |   4 +
 kernel/fork.c                                 |   4 +
 kernel/sched/core.c                           |  52 +++
 kernel/sched/fair.c                           |  43 +++
 kernel/sched/psi.c                            |  22 +-
 kernel/sched/sched.h                          |   8 +
 kernel/sched/stats.h                          |  10 +
 mm/memcontrol.c                               |   5 +
 .../selftests/cgroup/test_memcontrol.c        | 332 ++++++++++++++++++
 11 files changed, 522 insertions(+), 1 deletion(-)


base-commit: a8c591ed6b672915e0be57843f943a2a723aff40
-- 
2.53.0-Meta