[LSF/MM/BPF TOPIC] Reimagining Memory Cgroup (memcg_ext)

[Shakeel Butt]

Over the last couple of weeks, I have been brainstorming on how I would go
about redesigning memcg, taking inspiration from sched_ext and bpfoom, with a
focus on existing challenges and issues. This proposal outlines the high-level
direction. Followup emails and patch series will cover and brainstorm the
mechanisms (of course BPF) to achieve these goals.

Memory cgroups provide memory accounting and the ability to control memory usage
of workloads through two categories of limits. Throttling limits (memory.max and
memory.high) cap memory consumption. Protection limits (memory.min and
memory.low) shield a workload's memory from reclaim under external memory
pressure.

Challenges
----------

- Workload owners rarely know their actual memory requirements, leading to
  overprovisioned limits, lower utilization, and higher infrastructure costs.

- Throttling limit enforcement is synchronous in the allocating task's context,
  which can stall latency-sensitive threads.

- The stalled thread may hold shared locks, causing priority inversion -- all
  waiters are blocked regardless of their priority.

- Enforcement is indiscriminate -- there is no way to distinguish a
  performance-critical or latency-critical allocator from a latency-tolerant
  one.

- Protection limits assume static working sets size, forcing owners to either
  overprovision or build complex userspace infrastructure to dynamically adjust
  them.

Feature Wishlist
----------------

Here is the list of features and capabilities I want to enable in the
redesigned memcg limit enforcement world.

Per-Memcg Background Reclaim

In the new memcg world, with the goal of (mostly) eliminating direct synchronous
reclaim for limit enforcement, provide per-memcg background reclaimers which can
scale across CPUs with the allocation rate.

Lock-Aware Throttling

The ability to avoid throttling an allocating task that is holding locks, to
prevent priority inversion. In Meta's fleet, we have observed lock holders stuck
in memcg reclaim, blocking all waiters regardless of their priority or
criticality.

Thread-Level Throttling Control

Workloads should be able to indicate at the thread level which threads can be
synchronously throttled and which cannot. For example, while experimenting with
sched_ext, we drastically improved the performance of AI training workloads by
prioritizing threads interacting with the GPU. Similarly, applications can
identify the threads or thread pools on their performance-critical paths and
the memcg enforcement mechanism should not throttle them.

Combined Memory and Swap Limits

Some users (Google actually) need the ability to enforce limits based on
combined memory and swap usage, similar to cgroup v1's memsw limit, providing a
ceiling on total memory commitment rather than treating memory and swap
independently.

Dynamic Protection Limits

Rather than static protection limits, the kernel should support defining
protection based on the actual working set of the workload, leveraging signals
such as working set estimation, PSI, refault rates, or a combination thereof to
automatically adapt to the workload's current memory needs.

Shared Memory Semantics

With more flexibility in limit enforcement, the kernel should be able to
account for memory shared between workloads (cgroups) during enforcement.
Today, enforcement only looks at each workload's memory usage independently.
Sensible shared memory semantics would allow the enforcer to consider
cross-cgroup sharing when making reclaim and throttling decisions.

Memory Tiering

With a flexible limit enforcement mechanism, the kernel can balance memory
usage of different workloads across memory tiers based on their performance
requirements. Tier accounting and hotness tracking are orthogonal, but the
decisions of when and how to balance memory between tiers should be handled by
the enforcer.

Collaborative Load Shedding

Many workloads communicate with an external entity for load balancing and rely
on their own usage metrics like RSS or memory pressure to signal whether they
can accept more or less work. This is guesswork. Instead of the
workload guessing, the limit enforcer -- which is actually managing the
workload's memory usage -- should be able to communicate available headroom or
request the workload to shed load or reduce memory usage. This collaborative
load shedding mechanism would allow workloads to make informed decisions rather
than reacting to coarse signals.

Cross-Subsystem Collaboration

Finally, the limit enforcement mechanism should collaborate with the CPU
scheduler and other subsystems that can release memory. For example, dirty
memory is not reclaimable and the memory subsystem wakes up flushers to trigger
writeback. However, flushers need CPU to run -- asking the CPU scheduler to
prioritize them ensures the kernel does not lack reclaimable memory under
stressful conditions. Similarly, some subsystems free memory through workqueues
or RCU callbacks. While this may seem orthogonal to limit enforcement, we can
definitely take advantage by having visibility into these situations.

Putting It All Together
-----------------------

To illustrate the end goal, here is an example of the scenario I want to
enable. Suppose there is an AI agent controlling the resources of a host. I
should be able to provide the following policy and everything should work out
of the box:

Policy: "keep system-level memory utilization below 95 percent;
avoid priority inversions by not throttling allocators holding locks; trim each
workload's usage to its working set without regressing its relevant performance
metrics; collaborate with workloads on load shedding and memory trimming
decisions; and under extreme memory pressure, collaborate with the OOM killer
and the central job scheduler to kill and clean up a workload."

Initially I added this example for fun, but from [1] it seems like there is a
real need to enable such capabilities.

[1] https://arxiv.org/abs/2602.09345