[PATCH v2 0/6] memcg: Batch mod_memcg_state() for MEMCG_SOCK.

[PATCH v2 0/6] memcg: Batch mod_memcg_state() for MEMCG_SOCK.

[Kuniyuki Iwashima]

On our production fleet, we saw mem_cgroup_sk_charge() spent over
3x the cycles of try_charge_memcg() although the former should be
a thin wrapper of the latter.

Looking at perf during tcp_rr, most of the cycles were spent in
mod_memcg_state(MEMCG_SOCK), which incurred two heavy memory stalls
due to cache line misses. (see patch 4)

Also, tcp_rr w/ 20000 flows throughput often degrades by ~50% on
many core machine. (see patch 5)

This series adds nr_sock[] to struct memcg_stock_pcp to batch
mod_memcg_state() for MEMCG_SOCK and move heavy operations to
out of local_lock(&memcg_stock.lock).

With this series, tcp_rr w/ 10000 flows achieves 31% more
throughput with lower latency on machine w/ AMD EPYC 7B12 64-Core
Processor (Zen 2 Rome, 2 sockets, 256 vCPU)

  tcp_rr (10000 flows, 128 threads, 60s, average of 3 runs):

  Metric                     |       Baseline |        Patched |    Diff
  ---------------------------+----------------+----------------+--------
  Throughput (RPS)           |     20,884,499 |     27,466,490 | +31.52%
  Latency Mean (us)          |         478.46 |         363.71 | -23.98%
  Latency p50 (us)           |         453.96 |         356.68 | -21.43%
  Latency p90 (us)           |         669.00 |         479.56 | -28.32%
  Latency p99 (us)           |         969.38 |         617.80 | -36.27%


Changes:
  v2:
    * Rebased on mm.git/mm-new
    * Add patch 5 for flaky throughput issue
    * Patch 4:
      * Benchmark again after rebase and update numbers
      * Check in_nmi() in __refill_stock()

  v1: https://lore.kernel.org/linux-mm/20261016052811.3085635-1-kuniyu@google.com/


Kuniyuki Iwashima (6):
  memcg: Remove unnecessary READ_ONCE() under memcg_stock.lock.
  memcg: Reduce nesting in consume_stock() and refill_stock().
  memcg: Annotate slow paths in consume_stock() and refill_stock().
  memcg: Batch mod_memcg_state() for MEMCG_SOCK.
  memcg: Uncharge pages outside memcg_stock.lock in __refill_stock().
  selftests: cgroup: Deflake test_memcg_sock() in test_memcontrol.c.

 mm/memcontrol.c                               | 178 +++++++++++++-----
 .../selftests/cgroup/test_memcontrol.c        |   8 +
 2 files changed, 139 insertions(+), 47 deletions(-)

-- 
2.56.0.313.g2479322e68-goog

[PATCH v2 1/6] memcg: Remove unnecessary READ_ONCE() under memcg_stock.lock.

[Kuniyuki Iwashima]

this_cpu_ptr(&memcg_stock) is always modified under
local_lock(&memcg_stock.lock).

Unlike is_memcg_drain_needed() under RCU, consume_stock(),
drain_stock(), and refill_stock() do not need READ_ONCE().

Especially, in refill_stock(), READ_ONCE() is used twice for
stock->cached[i], which forces the compiler to emit two volatile
loads from stock->cached[i].

Let's remove the unnecessary READ_ONCE() and eliminate the
redundant loads in refill_stock():

Before:

  cached = READ_ONCE(stock->cached[i]);
   4c 8b 80 88 d7 bd 84 	movq	-0x7b422878(%rax), %r8
  if (memcg == READ_ONCE(stock->cached[i])) {
   48 3b 98 88 d7 bd 84 	cmpq	-0x7b422878(%rax), %rbx
   0f 84 86 01 00 00    	je	0xffffffff817d01e0 <refill_stock+0x210>

After:

  cached = stock->cached[i];
   4c 8b 88 88 d7 bd 84 	movq	-0x7b422878(%rax), %r9
  if (memcg == stock->cached[i]) {
   4c 39 cb             	cmpq	%r9, %rbx
   0f 84 7a 01 00 00    	je	0xffffffff817d01d0 <refill_stock+0x200>

Signed-off-by: Kuniyuki Iwashima <kuniyu@google.com>
---
 mm/memcontrol.c | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 84f4a8fd4d16..8a401550dfaa 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -2164,10 +2164,10 @@ static bool consume_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
 	stock = this_cpu_ptr(&memcg_stock);
 
 	for (i = 0; i < NR_MEMCG_STOCK; ++i) {
-		if (memcg != READ_ONCE(stock->cached[i]))
+		if (memcg != stock->cached[i])
 			continue;
 
-		stock_pages = READ_ONCE(stock->nr_pages[i]);
+		stock_pages = stock->nr_pages[i];
 		if (stock_pages >= nr_pages) {
 			stock_pages -= nr_pages;
 			WRITE_ONCE(stock->nr_pages[i], stock_pages);
@@ -2197,13 +2197,13 @@ static void memcg_uncharge(struct mem_cgroup *memcg, unsigned int nr_pages)
  */
 static void drain_stock(struct memcg_stock_pcp *stock, int i)
 {
-	struct mem_cgroup *old = READ_ONCE(stock->cached[i]);
+	struct mem_cgroup *old = stock->cached[i];
 	uint8_t stock_pages;
 
 	if (!old)
 		return;
 
-	stock_pages = READ_ONCE(stock->nr_pages[i]);
+	stock_pages = stock->nr_pages[i];
 	if (stock_pages) {
 		memcg_uncharge(old, stock_pages);
 		WRITE_ONCE(stock->nr_pages[i], 0);
@@ -2284,11 +2284,11 @@ static void refill_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
 
 	stock = this_cpu_ptr(&memcg_stock);
 	for (i = 0; i < NR_MEMCG_STOCK; ++i) {
-		cached = READ_ONCE(stock->cached[i]);
+		cached = stock->cached[i];
 		if (!cached && empty_slot == -1)
 			empty_slot = i;
-		if (memcg == READ_ONCE(stock->cached[i])) {
-			stock_pages = READ_ONCE(stock->nr_pages[i]) + nr_pages;
+		if (memcg == stock->cached[i]) {
+			stock_pages = stock->nr_pages[i] + nr_pages;
 			if (stock_pages > MEMCG_STOCK_HIGH) {
 				memcg_uncharge(memcg,
 					       stock_pages - MEMCG_STOCK_LOW);
-- 
2.56.0.313.g2479322e68-goog

[PATCH v2 2/6] memcg: Reduce nesting in consume_stock() and refill_stock().

[Kuniyuki Iwashima]

A later patch adds more code to the most nested block in
consume_stock() and refill_stock().

Let's reduce the nesting to make the later change cleaner.

Signed-off-by: Kuniyuki Iwashima <kuniyu@google.com>
---
 mm/memcontrol.c | 64 ++++++++++++++++++++++++++-----------------------
 1 file changed, 34 insertions(+), 30 deletions(-)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 8a401550dfaa..6666acc9cee8 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -2168,15 +2168,17 @@ static bool consume_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
 			continue;
 
 		stock_pages = stock->nr_pages[i];
-		if (stock_pages >= nr_pages) {
-			stock_pages -= nr_pages;
-			WRITE_ONCE(stock->nr_pages[i], stock_pages);
-			if (!stock_pages) {
-				WRITE_ONCE(stock->cached[i], NULL);
-				css_put(&memcg->css);
-			}
-			ret = true;
+		if (stock_pages < nr_pages)
+			break;
+
+		stock_pages -= nr_pages;
+		WRITE_ONCE(stock->nr_pages[i], stock_pages);
+		if (!stock_pages) {
+			WRITE_ONCE(stock->cached[i], NULL);
+			css_put(&memcg->css);
 		}
+
+		ret = true;
 		break;
 	}
 
@@ -2258,7 +2260,6 @@ static void refill_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
 	struct memcg_stock_pcp *stock;
 	struct mem_cgroup *cached;
 	unsigned int stock_pages;
-	bool success = false;
 	int empty_slot = -1;
 	int i;
 
@@ -2287,32 +2288,35 @@ static void refill_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
 		cached = stock->cached[i];
 		if (!cached && empty_slot == -1)
 			empty_slot = i;
-		if (memcg == stock->cached[i]) {
-			stock_pages = stock->nr_pages[i] + nr_pages;
-			if (stock_pages > MEMCG_STOCK_HIGH) {
-				memcg_uncharge(memcg,
-					       stock_pages - MEMCG_STOCK_LOW);
-				stock_pages = MEMCG_STOCK_LOW;
-			}
-			WRITE_ONCE(stock->nr_pages[i], stock_pages);
-			success = true;
-			break;
+
+		if (memcg != cached)
+			continue;
+
+		stock_pages = stock->nr_pages[i] + nr_pages;
+		if (stock_pages > MEMCG_STOCK_HIGH) {
+			memcg_uncharge(memcg,
+				       stock_pages - MEMCG_STOCK_LOW);
+			stock_pages = MEMCG_STOCK_LOW;
 		}
+
+		WRITE_ONCE(stock->nr_pages[i], stock_pages);
+
+		local_unlock(&memcg_stock.lock);
+		return;
 	}
 
-	if (!success) {
-		i = empty_slot;
-		if (i == -1) {
-			i = stock->drain_idx++;
-			if (stock->drain_idx == NR_MEMCG_STOCK)
-				stock->drain_idx = 0;
-			drain_stock(stock, i);
-		}
-		css_get(&memcg->css);
-		WRITE_ONCE(stock->cached[i], memcg);
-		WRITE_ONCE(stock->nr_pages[i], nr_pages);
+	i = empty_slot;
+	if (i == -1) {
+		i = stock->drain_idx++;
+		if (stock->drain_idx == NR_MEMCG_STOCK)
+			stock->drain_idx = 0;
+		drain_stock(stock, i);
 	}
 
+	css_get(&memcg->css);
+	WRITE_ONCE(stock->cached[i], memcg);
+	WRITE_ONCE(stock->nr_pages[i], nr_pages);
+
 	local_unlock(&memcg_stock.lock);
 }
 
-- 
2.56.0.313.g2479322e68-goog

[PATCH v2 3/6] memcg: Annotate slow paths in consume_stock() and refill_stock().

[Kuniyuki Iwashima]

In consume_stock() and refill_stock(), each per-CPU stock slot
holds up to MEMCG_STOCK_HIGH (64) pages, so depleting a slot in
consume_stock() (!stock_pages) or trimming a slot in refill_stock()
(stock_pages > MEMCG_STOCK_HIGH) happens at most once per batch.

Similarly, requests larger than MEMCG_CHARGE_BATCH, local_trylock()
failure, and insufficient cached pages (stock_pages < nr_pages)
are all slow paths.

Without unlikely() annotations (when building without FDO, or
when an FDO profile does not cover a caller), the compiler treats
these branches as likely and inlines css_put() in consume_stock()
and memcg_uncharge() in refill_stock() right into the middle of
the hot path.

Consequently, the fast path has to jump over the inlined slow
path and jump again to the function epilogue after local_unlock().

Let's annotate the slow paths in consume_stock() and refill_stock()
with unlikely() so that the compiler keeps css_put() and
memcg_uncharge() out-of-line at the cold tail and lets local_unlock()
and the function epilogue fall through on the fast path:

Signed-off-by: Kuniyuki Iwashima <kuniyu@google.com>
---
 mm/memcontrol.c | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 6666acc9cee8..a04b44656b38 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -2157,8 +2157,8 @@ static bool consume_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
 	bool ret = false;
 	int i;
 
-	if (nr_pages > MEMCG_CHARGE_BATCH ||
-	    !local_trylock(&memcg_stock.lock))
+	if (unlikely(nr_pages > MEMCG_CHARGE_BATCH ||
+		     !local_trylock(&memcg_stock.lock)))
 		return ret;
 
 	stock = this_cpu_ptr(&memcg_stock);
@@ -2168,12 +2168,12 @@ static bool consume_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
 			continue;
 
 		stock_pages = stock->nr_pages[i];
-		if (stock_pages < nr_pages)
+		if (unlikely(stock_pages < nr_pages))
 			break;
 
 		stock_pages -= nr_pages;
 		WRITE_ONCE(stock->nr_pages[i], stock_pages);
-		if (!stock_pages) {
+		if (unlikely(!stock_pages)) {
 			WRITE_ONCE(stock->cached[i], NULL);
 			css_put(&memcg->css);
 		}
@@ -2273,8 +2273,8 @@ static void refill_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
 
 	VM_WARN_ON_ONCE(mem_cgroup_is_root(memcg));
 
-	if (nr_pages > MEMCG_CHARGE_BATCH ||
-	    !local_trylock(&memcg_stock.lock)) {
+	if (unlikely(nr_pages > MEMCG_CHARGE_BATCH ||
+		     !local_trylock(&memcg_stock.lock))) {
 		/*
 		 * In case of larger than batch refill or unlikely failure to
 		 * lock the percpu memcg_stock.lock, uncharge memcg directly.
@@ -2293,7 +2293,7 @@ static void refill_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
 			continue;
 
 		stock_pages = stock->nr_pages[i] + nr_pages;
-		if (stock_pages > MEMCG_STOCK_HIGH) {
+		if (unlikely(stock_pages > MEMCG_STOCK_HIGH)) {
 			memcg_uncharge(memcg,
 				       stock_pages - MEMCG_STOCK_LOW);
 			stock_pages = MEMCG_STOCK_LOW;
-- 
2.56.0.313.g2479322e68-goog

[PATCH v2 4/6] memcg: Batch mod_memcg_state() for MEMCG_SOCK.

[Kuniyuki Iwashima]

On our production fleet, we saw mem_cgroup_sk_charge() spent over
3x the cycles of try_charge_memcg() although the former should be
a thin wrapper of the latter.

Looking at perf during tcp_rr, most of the cycles were spent in
mod_memcg_state(MEMCG_SOCK), which incurred two heavy memory stalls
due to cache line misses:

  mem_cgroup_sk_charge  vmlinux [Percent: local period]
  ...
     1.74 │       mov     cpu_number,%esi
     0.70 │       movslq  %ebx,%rcx
     0.31 │       mov     0x518(%r14),%rdx
     8.99 │       add     %rcx,%gs:0x18(%rdx,%rax,8)
                  ^-- this_cpu_add(memcg->vmstats_percpu->state[i], val);
                      in mod_memcg_state()
 >> 31.21 │       test    %ebx,%ebx
     0.05 │     ↓ je      97
          │       mov     %r14,%rdi
          │     → call    __css_rstat_updated
     0.31 │       mov     0x518(%r14),%rdx
          │ 72:   test    %rdx,%rdx
     0.06 │     ↓ je      97
          │       mov     this_cpu_off,%rsi
     0.17 │       mov     0x10(%rsi,%rdx,1),%rax
     6.59 │       mov     0xc18(%rax),%eax
                  ^-- atomic_read(&vmstats->stats_updates)
                      in memcg_vmstats_needs_flush()
 >> 26.02 │       mov     __num_online_cpus,%ecx
 ...
     0.65 │ e5:   mov     %r14,%rdi
          │       mov     %edx,%esi
     1.88 │       mov     %ebx,%edx
          │     → call    try_charge_memcg
     1.56 │       xor     %ebp,%ebp
     0.18 │       test    %eax,%eax
     0.22 │     ↑ jne     9f
          │     ↑ jmp     31

Even worse, after the 2nd cache line miss, the subsequent
atomic_add() for statc->vmstats->stats_updates invalidates the
cache for other CPUs.

Moreover, try_charge_memcg() returned just after consume_stock()
in most cases, but it was not inlined even with FDO.

Usually, skb is charged and uncharged immediately, and caching
the counter in struct memcg_stock_pcp helps.

Let's reduce NR_MEMCG_STOCK from 7 to 6 to fit nr_sock[] in the
single cacheline of struct memcg_stock_pcp, inline the fast path
of consume_stock(), and batch mod_memcg_state() for MEMCG_SOCK
(-64, 64).

With this change, tcp_rr w/ 10000 flows achieves ~29% more
throughput with lower latency on machine w/ AMD EPYC 7B12 64-Core
Processor (Zen 2 Rome, 2 sockets, 256 vCPU), and the TCP stack
spends far fewer cycles on memcg.

  tcp_rr (10000 flows, 128 threads, 60s, average of 3 runs):

  Metric                     |       Baseline |        Patched |    Diff
  ---------------------------+----------------+----------------+--------
  Throughput (RPS)           |     20,884,499 |     26,960,730 | +29.09%
  Latency Mean (us)          |         478.46 |         370.50 | -22.56%
  Latency p50 (us)           |         453.96 |         361.80 | -20.30%
  Latency p90 (us)           |         669.00 |         482.98 | -27.81%
  Latency p99 (us)           |         969.38 |         622.92 | -35.74%

  perf cycles (10000 flows, 10s, average of 3 runs):

  Server (M cycles / %)      |           Baseline |            Patched |    Diff
  ---------------------------+--------------------+--------------------+--------
  Host Total Cycles          |        4,566,914 M |        4,575,636 M |   +0.2%
  tcp_rcv_established()      | 158,009 M (3.460%) | 152,520 M (3.333%) |   -3.5%
  mem_cgroup_sk_charge()     |  18,420 M (0.403%) |   7,474 M (0.163%) |  -59.4%
  mem_cgroup_sk_uncharge()   |  16,593 M (0.363%) |   8,541 M (0.187%) |  -48.5%
  __css_rstat_updated()      |   3,502 M (0.077%) |       0 M (0.000%) | -100.0%
  Memcg Sock Total           |  38,515 M (0.843%) |  16,015 M (0.350%) |  -58.4%

As a side effect, "sock" in memory.stat can be off by up to
MEMCG_CHARGE_BATCH pages per CPU until the slot is drained,
like memory.current with nr_pages[].  This is acceptable, as
flushing nr_sock[] every 2 seconds (FLUSH_TIME) across all CPUs
would rather cause cache line invalidation and hurt performance.

Note that MEMCG_CHARGE_BATCH is unsigned int, so (int) cast is
needed, and stock_sock is declared as int rather than int8_t
to avoid potential overflow for MEMCG_CHARGE_BATCH > 64 and to
generate better asm without 8-bit partial-register operations
or extra sign-extension before mod_memcg_state().

Also, __always_inline is needed for both __consume_stock() and
__refill_stock(), otherwise, FDO branch weight is calculated
from two callers with mixed states (nr_sock == 0 vs nr_sock != 0).

Signed-off-by: Kuniyuki Iwashima <kuniyu@google.com>
---
v2
  * Benchmark again after rebase and update numbers
  * Check in_nmi() in __refill_stock()
---
 mm/memcontrol.c | 85 ++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 74 insertions(+), 11 deletions(-)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index a04b44656b38..b92e952c076a 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -2065,9 +2065,9 @@ void mem_cgroup_print_oom_group(struct mem_cgroup *memcg)
 
 /*
  * The value of NR_MEMCG_STOCK is selected to keep the cached memcgs and their
- * nr_pages in a single cacheline. This may change in future.
+ * nr_pages and nr_sock in a single cacheline. This may change in future.
  */
-#define NR_MEMCG_STOCK 7
+#define NR_MEMCG_STOCK 6
 
 /*
  * Watermarks for a charge stock slot, in the spirit of pcp->high and
@@ -2081,6 +2081,7 @@ void mem_cgroup_print_oom_group(struct mem_cgroup *memcg)
 struct memcg_stock_pcp {
 	local_trylock_t lock;
 	uint8_t nr_pages[NR_MEMCG_STOCK];
+	int8_t nr_sock[NR_MEMCG_STOCK];
 	struct mem_cgroup *cached[NR_MEMCG_STOCK];
 
 	struct work_struct work;
@@ -2140,9 +2141,10 @@ static bool obj_stock_flush_required(struct obj_stock_pcp *stock,
 				     struct mem_cgroup *root_memcg);
 
 /**
- * consume_stock: Try to consume stocked charge on this cpu.
+ * __consume_stock: Try to consume stocked charge on this cpu.
  * @memcg: memcg to consume from.
  * @nr_pages: how many pages to charge.
+ * @nr_sock: how many pages to account for MEMCG_SOCK.
  *
  * Consume the cached charge if enough nr_pages are present otherwise return
  * failure. Also return failure for charge request larger than
@@ -2150,10 +2152,13 @@ static bool obj_stock_flush_required(struct obj_stock_pcp *stock,
  *
  * returns true if successful, false otherwise.
  */
-static bool consume_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
+static __always_inline bool __consume_stock(struct mem_cgroup *memcg,
+					    unsigned int nr_pages,
+					    unsigned int nr_sock)
 {
 	struct memcg_stock_pcp *stock;
 	uint8_t stock_pages;
+	int stock_sock = 0;
 	bool ret = false;
 	int i;
 
@@ -2173,7 +2178,19 @@ static bool consume_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
 
 		stock_pages -= nr_pages;
 		WRITE_ONCE(stock->nr_pages[i], stock_pages);
-		if (unlikely(!stock_pages)) {
+
+		if (nr_sock) {
+			stock_sock = stock->nr_sock[i] + nr_sock;
+
+			if (likely(stock_sock < (int)MEMCG_CHARGE_BATCH)) {
+				WRITE_ONCE(stock->nr_sock[i], stock_sock);
+				stock_sock = 0;
+			} else {
+				WRITE_ONCE(stock->nr_sock[i], 0);
+			}
+		}
+
+		if (unlikely(!stock_pages && !stock->nr_sock[i])) {
 			WRITE_ONCE(stock->cached[i], NULL);
 			css_put(&memcg->css);
 		}
@@ -2184,9 +2201,17 @@ static bool consume_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
 
 	local_unlock(&memcg_stock.lock);
 
+	if (unlikely(stock_sock))
+		mod_memcg_state(memcg, MEMCG_SOCK, stock_sock);
+
 	return ret;
 }
 
+static bool consume_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
+{
+	return __consume_stock(memcg, nr_pages, 0);
+}
+
 static void memcg_uncharge(struct mem_cgroup *memcg, unsigned int nr_pages)
 {
 	page_counter_uncharge(&memcg->memory, nr_pages);
@@ -2201,6 +2226,7 @@ static void drain_stock(struct memcg_stock_pcp *stock, int i)
 {
 	struct mem_cgroup *old = stock->cached[i];
 	uint8_t stock_pages;
+	int stock_sock;
 
 	if (!old)
 		return;
@@ -2211,6 +2237,12 @@ static void drain_stock(struct memcg_stock_pcp *stock, int i)
 		WRITE_ONCE(stock->nr_pages[i], 0);
 	}
 
+	stock_sock = stock->nr_sock[i];
+	if (stock_sock) {
+		mod_memcg_state(old, MEMCG_SOCK, stock_sock);
+		WRITE_ONCE(stock->nr_sock[i], 0);
+	}
+
 	WRITE_ONCE(stock->cached[i], NULL);
 	css_put(&old->css);
 }
@@ -2255,12 +2287,15 @@ static void drain_local_obj_stock(struct work_struct *dummy)
 	local_unlock(&obj_stock.lock);
 }
 
-static void refill_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
+static __always_inline void __refill_stock(struct mem_cgroup *memcg,
+					   unsigned int nr_pages,
+					   unsigned int nr_sock)
 {
 	struct memcg_stock_pcp *stock;
 	struct mem_cgroup *cached;
 	unsigned int stock_pages;
 	int empty_slot = -1;
+	int stock_sock = 0;
 	int i;
 
 	/*
@@ -2274,13 +2309,15 @@ static void refill_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
 	VM_WARN_ON_ONCE(mem_cgroup_is_root(memcg));
 
 	if (unlikely(nr_pages > MEMCG_CHARGE_BATCH ||
+		     (IS_ENABLED(CONFIG_MEMCG_NMI_SAFETY_REQUIRES_ATOMIC) &&
+		      in_nmi()) ||
 		     !local_trylock(&memcg_stock.lock))) {
 		/*
 		 * In case of larger than batch refill or unlikely failure to
 		 * lock the percpu memcg_stock.lock, uncharge memcg directly.
 		 */
 		memcg_uncharge(memcg, nr_pages);
-		return;
+		goto update_sock;
 	}
 
 	stock = this_cpu_ptr(&memcg_stock);
@@ -2301,7 +2338,22 @@ static void refill_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
 
 		WRITE_ONCE(stock->nr_pages[i], stock_pages);
 
+		if (nr_sock) {
+			stock_sock = stock->nr_sock[i] - nr_sock;
+
+			if (likely(stock_sock > -(int)MEMCG_CHARGE_BATCH)) {
+				WRITE_ONCE(stock->nr_sock[i], stock_sock);
+				stock_sock = 0;
+			} else {
+				WRITE_ONCE(stock->nr_sock[i], 0);
+			}
+		}
+
 		local_unlock(&memcg_stock.lock);
+
+		if (unlikely(stock_sock))
+			mod_memcg_state(memcg, MEMCG_SOCK, stock_sock);
+
 		return;
 	}
 
@@ -2318,6 +2370,15 @@ static void refill_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
 	WRITE_ONCE(stock->nr_pages[i], nr_pages);
 
 	local_unlock(&memcg_stock.lock);
+
+update_sock:
+	if (nr_sock)
+		mod_memcg_state(memcg, MEMCG_SOCK, -nr_sock);
+}
+
+static void refill_stock(struct mem_cgroup *memcg, unsigned int nr_pages)
+{
+	__refill_stock(memcg, nr_pages, 0);
 }
 
 static bool is_memcg_drain_needed(struct memcg_stock_pcp *stock,
@@ -2333,7 +2394,8 @@ static bool is_memcg_drain_needed(struct memcg_stock_pcp *stock,
 		if (!memcg)
 			continue;
 
-		if (READ_ONCE(stock->nr_pages[i]) &&
+		if ((READ_ONCE(stock->nr_pages[i]) ||
+		     READ_ONCE(stock->nr_sock[i])) &&
 		    mem_cgroup_is_descendant(memcg, root_memcg)) {
 			flush = true;
 			break;
@@ -5798,6 +5860,9 @@ bool mem_cgroup_sk_charge(const struct sock *sk, unsigned int nr_pages,
 	if (!cgroup_subsys_on_dfl(memory_cgrp_subsys))
 		return memcg1_charge_skmem(memcg, nr_pages, gfp_mask);
 
+	if (likely(__consume_stock(memcg, nr_pages, nr_pages)))
+		return true;
+
 	if (try_charge_memcg(memcg, gfp_mask, nr_pages) == 0) {
 		mod_memcg_state(memcg, MEMCG_SOCK, nr_pages);
 		return true;
@@ -5820,9 +5885,7 @@ void mem_cgroup_sk_uncharge(const struct sock *sk, unsigned int nr_pages)
 		return;
 	}
 
-	mod_memcg_state(memcg, MEMCG_SOCK, -nr_pages);
-
-	refill_stock(memcg, nr_pages);
+	__refill_stock(memcg, nr_pages, nr_pages);
 }
 
 void mem_cgroup_flush_workqueue(void)
-- 
2.56.0.313.g2479322e68-goog

[PATCH v2 5/6] memcg: Uncharge pages outside memcg_stock.lock in __refill_stock().

[Kuniyuki Iwashima]

On host w/ AMD EPYC 9B45 128-Core Processor (Zen 5 Turin,
2 sockets, 512 vCPUs), tcp_rr w/ 20000 flows throughput often
degrades by ~50%, 32M -> 17M~27M.

In the bad case, the number of involuntary context switches
was extremely high by 12x (241733 -> 3093367).

In perf record, 96% samples had the same pattern, where the
thread was interrupted during page_counter_uncharge() by softirq,
and later ksoftirqd was woken up:

  # perf record -C 0-15 -e sched:sched_wakeup -g \
      -o /tmp/sched_wakeup.data -- sleep 1
  # perf report -i /tmp/sched_wakeup.data --stdio --no-children
    96.36%  tcp_rr   [kernel.kallsyms]  [k] wake_up_process
            |
            ---wake_up_process
               |
                --96.00%--irq_exit_rcu
                          |
                           --95.82%--sysvec_call_function_single
                                     asm_sysvec_call_function_single
                                     |
                                      --95.58%--page_counter_uncharge
                                                mem_cgroup_sk_uncharge
                                                sock_rfree
                                                tcp_recvmsg
                                                __x64_sys_recvfrom
                                                do_syscall_64
                                                entry_SYSCALL_64_after_hwframe

  # perf script -i /tmp/sched_wakeup.data | grep -A 20 "comm=ksoftirqd"
  tcp_rr  101952 [001]  9679.340985: sched:sched_wakeup: comm=ksoftirqd/1 pid=23 prio=120 target_cpu=001
  	ffffffffb64cd0cc wake_up_process+0xe6c ([kernel.kallsyms])
  	ffffffffb64cb034 irq_exit_rcu+0x264 ([kernel.kallsyms])
  	ffffffffb68ad31c sysvec_call_function_single+0xcc ([kernel.kallsyms])
  	ffffffffb560158a asm_sysvec_call_function_single+0x1a ([kernel.kallsyms])
  	ffffffffb653c4ce page_counter_uncharge+0x2e ([kernel.kallsyms])
  	ffffffffb6541b28 mem_cgroup_sk_uncharge+0x3c8 ([kernel.kallsyms])
  	ffffffffb6561383 sock_rfree+0x83 ([kernel.kallsyms])
  	ffffffffb6577faf tcp_recvmsg+0x2ff ([kernel.kallsyms])
  	ffffffffb65604d9 __x64_sys_recvfrom+0x159 ([kernel.kallsyms])

The problem is that when stock_pages exceeds MEMCG_STOCK_HIGH
or all slots in struct memcg_stock_pcp are full, __refill_stock()
calls memcg_uncharge() or drain_stock() while holding
memcg_stock.lock.

Because local_trylock() only disables preemption and leaves
interrupts and softirqs enabled, an RPS IPI arriving while
__refill_stock() stalls on page_counter_uncharge() invokes
NET_RX_SOFTIRQ (process_backlog()) on the same CPU while
memcg_stock.lock is held.

Every incoming packet in process_backlog() then fails
local_trylock() in __consume_stock() and falls back to
try_charge_memcg() and __refill_stock(), hammering the global
page_counter atomic twice per packet, and slow processing
finally wakes up ksoftirqd.

Let's move memcg_uncharge(), mod_memcg_state(), and css_put()
outside local_unlock(&memcg_stock.lock) in __refill_stock().

With the change, tcp_rr w/ 20000 flows gains ~4% more throughput
and is stable around the highest value.

Signed-off-by: Kuniyuki Iwashima <kuniyu@google.com>
---
 mm/memcontrol.c | 29 +++++++++++++++++++++++------
 1 file changed, 23 insertions(+), 6 deletions(-)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index b92e952c076a..87b189aba07d 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -2331,13 +2331,13 @@ static __always_inline void __refill_stock(struct mem_cgroup *memcg,
 
 		stock_pages = stock->nr_pages[i] + nr_pages;
 		if (unlikely(stock_pages > MEMCG_STOCK_HIGH)) {
-			memcg_uncharge(memcg,
-				       stock_pages - MEMCG_STOCK_LOW);
-			stock_pages = MEMCG_STOCK_LOW;
+			WRITE_ONCE(stock->nr_pages[i], MEMCG_STOCK_LOW);
+			stock_pages -= MEMCG_STOCK_LOW;
+		} else {
+			WRITE_ONCE(stock->nr_pages[i], stock_pages);
+			stock_pages = 0;
 		}
 
-		WRITE_ONCE(stock->nr_pages[i], stock_pages);
-
 		if (nr_sock) {
 			stock_sock = stock->nr_sock[i] - nr_sock;
 
@@ -2353,16 +2353,25 @@ static __always_inline void __refill_stock(struct mem_cgroup *memcg,
 
 		if (unlikely(stock_sock))
 			mod_memcg_state(memcg, MEMCG_SOCK, stock_sock);
+		if (unlikely(stock_pages))
+			memcg_uncharge(memcg, stock_pages);
 
 		return;
 	}
 
+	cached = NULL;
+	stock_pages = 0;
+
 	i = empty_slot;
 	if (i == -1) {
 		i = stock->drain_idx++;
 		if (stock->drain_idx == NR_MEMCG_STOCK)
 			stock->drain_idx = 0;
-		drain_stock(stock, i);
+
+		cached = stock->cached[i];
+		stock_pages = stock->nr_pages[i];
+		stock_sock = stock->nr_sock[i];
+		WRITE_ONCE(stock->nr_sock[i], 0);
 	}
 
 	css_get(&memcg->css);
@@ -2371,6 +2380,14 @@ static __always_inline void __refill_stock(struct mem_cgroup *memcg,
 
 	local_unlock(&memcg_stock.lock);
 
+	if (unlikely(cached)) {
+		if (stock_pages)
+			memcg_uncharge(cached, stock_pages);
+		if (stock_sock)
+			mod_memcg_state(cached, MEMCG_SOCK, stock_sock);
+		css_put(&cached->css);
+	}
+
 update_sock:
 	if (nr_sock)
 		mod_memcg_state(memcg, MEMCG_SOCK, -nr_sock);
-- 
2.56.0.313.g2479322e68-goog

[PATCH v2 6/6] selftests: cgroup: Deflake test_memcg_sock() in test_memcontrol.c.

[Kuniyuki Iwashima]

Now, MEMCG_SOCK stats updates are batched in per-CPU
memcg_stock->nr_sock[] up to MEMCG_CHARGE_BATCH (64) pages.

After closing the TCP connection, test_memcg_sock() runs
alloc_anon() to force an rstat flush and verifies that "sock"
in memory.stat drops to 0.

However, alloc_anon() only charges and uncharges anonymous pages
(nr_sock == 0), which trims stock->nr_pages[] without draining
stock->nr_sock[].

If a residual negative delta (-63 to -1 pages) remains cached
in stock->nr_sock[] when the socket closes, "sock" in memory.stat
stays positive even after rstat is flushed.

Let's write "0" and then "max" to memory.high in the flush_rounds
loop to trigger drain_all_stock() and flush stock->nr_sock[] into
per-CPU vmstats.

Signed-off-by: Kuniyuki Iwashima <kuniyu@google.com>
---
 tools/testing/selftests/cgroup/test_memcontrol.c | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tools/testing/selftests/cgroup/test_memcontrol.c b/tools/testing/selftests/cgroup/test_memcontrol.c
index bb673012dcf0..62cb0278168d 100644
--- a/tools/testing/selftests/cgroup/test_memcontrol.c
+++ b/tools/testing/selftests/cgroup/test_memcontrol.c
@@ -1483,6 +1483,14 @@ static int test_memcg_sock(const char *root)
 
 	/* the uncharge is done, but a flush still walking the CPUs can hide it from a read */
 	while (flush_rounds--) {
+		/*
+		 * Drain per-CPU memcg_stock so cached MEMCG_SOCK deltas in
+		 * stock->nr_sock[] are flushed to vmstats.
+		 */
+		if (cg_write(memcg, "memory.high", "0") ||
+		    cg_write(memcg, "memory.high", "max"))
+			goto cleanup;
+
 		if (cg_run(memcg, alloc_anon, (void *)MEMCG_FLUSH_BYTES))
 			goto cleanup;
 
-- 
2.56.0.313.g2479322e68-goog

Re: [PATCH v2 0/6] memcg: Batch mod_memcg_state() for MEMCG_SOCK.

[Shakeel Butt]

On Wed, Oct 28, 2026 at 06:50:12PM +0000, Kuniyuki Iwashima wrote:
> On our production fleet, we saw mem_cgroup_sk_charge() spent over
> 3x the cycles of try_charge_memcg() although the former should be
> a thin wrapper of the latter.
> 
> Looking at perf during tcp_rr, most of the cycles were spent in
> mod_memcg_state(MEMCG_SOCK), which incurred two heavy memory stalls
> due to cache line misses. (see patch 4)
> 
> Also, tcp_rr w/ 20000 flows throughput often degrades by ~50% on
> many core machine. (see patch 5)
> 
> This series adds nr_sock[] to struct memcg_stock_pcp to batch
> mod_memcg_state() for MEMCG_SOCK and move heavy operations to
> out of local_lock(&memcg_stock.lock).
> 
> With this series, tcp_rr w/ 10000 flows achieves 31% more
> throughput with lower latency on machine w/ AMD EPYC 7B12 64-Core
> Processor (Zen 2 Rome, 2 sockets, 256 vCPU)
> 
>   tcp_rr (10000 flows, 128 threads, 60s, average of 3 runs):
> 
>   Metric                     |       Baseline |        Patched |    Diff
>   ---------------------------+----------------+----------------+--------
>   Throughput (RPS)           |     20,884,499 |     27,466,490 | +31.52%
>   Latency Mean (us)          |         478.46 |         363.71 | -23.98%
>   Latency p50 (us)           |         453.96 |         356.68 | -21.43%
>   Latency p90 (us)           |         669.00 |         479.56 | -28.32%
>   Latency p99 (us)           |         969.38 |         617.80 | -36.27%
> 
> 
> Changes:
>   v2:
>     * Rebased on mm.git/mm-new
>     * Add patch 5 for flaky throughput issue
>     * Patch 4:
>       * Benchmark again after rebase and update numbers
>       * Check in_nmi() in __refill_stock()
> 
>   v1: https://lore.kernel.org/linux-mm/20261016052811.3085635-1-kuniyu@google.com/
> 
> 
> Kuniyuki Iwashima (6):
>   memcg: Remove unnecessary READ_ONCE() under memcg_stock.lock.
>   memcg: Reduce nesting in consume_stock() and refill_stock().
>   memcg: Annotate slow paths in consume_stock() and refill_stock().
>   memcg: Batch mod_memcg_state() for MEMCG_SOCK.
>   memcg: Uncharge pages outside memcg_stock.lock in __refill_stock().
>   selftests: cgroup: Deflake test_memcg_sock() in test_memcontrol.c.

Decouple the 4th patch (batch MEMCG_STOCK) from the remaining and send the
remaining patches separately and upfront.

For the 4th patch, I would prefer we first try to move the network sockets
from memcg to objcg path. Then I would prefer we optimize the stats update for
all stats instead of special casing for MEMCG_STOCK.

Re: [PATCH v2 1/6] memcg: Remove unnecessary READ_ONCE() under memcg_stock.lock.

[Shakeel Butt]

On Wed, Oct 28, 2026 at 06:50:13PM +0000, Kuniyuki Iwashima wrote:
> this_cpu_ptr(&memcg_stock) is always modified under
> local_lock(&memcg_stock.lock).

You need to be more precise here. Mainly the cached[] and nr_pages[] are only
modified under the local lock. There is flags field which is being modified
without it.

> 
> Unlike is_memcg_drain_needed() under RCU, consume_stock(),
> drain_stock(), and refill_stock() do not need READ_ONCE().
> 
> Especially, in refill_stock(), READ_ONCE() is used twice for
> stock->cached[i], which forces the compiler to emit two volatile
> loads from stock->cached[i].
> 
> Let's remove the unnecessary READ_ONCE() and eliminate the
> redundant loads in refill_stock():

What are we gaining with that? Is this measurable? With READ_ONCE() at least we
are very explicit that this field can read and modified concurrently by
different CPUs i.e. better readability.

Re: [PATCH v2 2/6] memcg: Reduce nesting in consume_stock() and refill_stock().

[Shakeel Butt]

On Wed, Oct 28, 2026 at 06:50:14PM +0000, Kuniyuki Iwashima wrote:
> A later patch adds more code to the most nested block in
> consume_stock() and refill_stock().
> 
> Let's reduce the nesting to make the later change cleaner.
> 
> Signed-off-by: Kuniyuki Iwashima <kuniyu@google.com>

Acked-by: Shakeel Butt <shakeel.butt@linux.dev>

Re: [PATCH v2 0/6] memcg: Batch mod_memcg_state() for MEMCG_SOCK.

[Shakeel Butt]

On Wed, Oct 28, 2026 at 05:41:04PM -0700, Shakeel Butt wrote:
> On Wed, Oct 28, 2026 at 06:50:12PM +0000, Kuniyuki Iwashima wrote:
> > On our production fleet, we saw mem_cgroup_sk_charge() spent over
> > 3x the cycles of try_charge_memcg() although the former should be
> > a thin wrapper of the latter.
> > 
> > Looking at perf during tcp_rr, most of the cycles were spent in
> > mod_memcg_state(MEMCG_SOCK), which incurred two heavy memory stalls
> > due to cache line misses. (see patch 4)
> > 
> > Also, tcp_rr w/ 20000 flows throughput often degrades by ~50% on
> > many core machine. (see patch 5)
> > 
> > This series adds nr_sock[] to struct memcg_stock_pcp to batch
> > mod_memcg_state() for MEMCG_SOCK and move heavy operations to
> > out of local_lock(&memcg_stock.lock).
> > 
> > With this series, tcp_rr w/ 10000 flows achieves 31% more
> > throughput with lower latency on machine w/ AMD EPYC 7B12 64-Core
> > Processor (Zen 2 Rome, 2 sockets, 256 vCPU)
> > 
> >   tcp_rr (10000 flows, 128 threads, 60s, average of 3 runs):
> > 
> >   Metric                     |       Baseline |        Patched |    Diff
> >   ---------------------------+----------------+----------------+--------
> >   Throughput (RPS)           |     20,884,499 |     27,466,490 | +31.52%
> >   Latency Mean (us)          |         478.46 |         363.71 | -23.98%
> >   Latency p50 (us)           |         453.96 |         356.68 | -21.43%
> >   Latency p90 (us)           |         669.00 |         479.56 | -28.32%
> >   Latency p99 (us)           |         969.38 |         617.80 | -36.27%
> > 
> > 
> > Changes:
> >   v2:
> >     * Rebased on mm.git/mm-new
> >     * Add patch 5 for flaky throughput issue
> >     * Patch 4:
> >       * Benchmark again after rebase and update numbers
> >       * Check in_nmi() in __refill_stock()
> > 
> >   v1: https://lore.kernel.org/linux-mm/20261016052811.3085635-1-kuniyu@google.com/
> > 
> > 
> > Kuniyuki Iwashima (6):
> >   memcg: Remove unnecessary READ_ONCE() under memcg_stock.lock.
> >   memcg: Reduce nesting in consume_stock() and refill_stock().
> >   memcg: Annotate slow paths in consume_stock() and refill_stock().
> >   memcg: Batch mod_memcg_state() for MEMCG_SOCK.
> >   memcg: Uncharge pages outside memcg_stock.lock in __refill_stock().
> >   selftests: cgroup: Deflake test_memcg_sock() in test_memcontrol.c.
> 
> Decouple the 4th patch (batch MEMCG_STOCK) from the remaining and send the
> remaining patches separately and upfront.
> 
> For the 4th patch, I would prefer we first try to move the network sockets
> from memcg to objcg path. Then I would prefer we optimize the stats update for
> all stats instead of special casing for MEMCG_STOCK.
> 

MEMCG_SOCK instead of MEMCG_STOCK

Re: [PATCH v2 3/6] memcg: Annotate slow paths in consume_stock() and refill_stock().

[Shakeel Butt]

On Wed, Oct 28, 2026 at 06:50:15PM +0000, Kuniyuki Iwashima wrote:
> In consume_stock() and refill_stock(), each per-CPU stock slot
> holds up to MEMCG_STOCK_HIGH (64) pages, so depleting a slot in
> consume_stock() (!stock_pages) or trimming a slot in refill_stock()
> (stock_pages > MEMCG_STOCK_HIGH) happens at most once per batch.
> 
> Similarly, requests larger than MEMCG_CHARGE_BATCH, local_trylock()
> failure, and insufficient cached pages (stock_pages < nr_pages)
> are all slow paths.
> 
> Without unlikely() annotations (when building without FDO, or
> when an FDO profile does not cover a caller), the compiler treats
> these branches as likely and inlines css_put() in consume_stock()
> and memcg_uncharge() in refill_stock() right into the middle of
> the hot path.

Please state the compiler along with its version here. I didn't see this behavior
on gcc 14.

Re: [PATCH v2 0/6] memcg: Batch mod_memcg_state() for MEMCG_SOCK.

[Kuniyuki Iwashima]

On Wed, 28 Oct 2026 17:52:32 -0700, Shakeel Butt wrote:
> On Wed, Oct 28, 2026 at 05:41:04PM -0700, Shakeel Butt wrote:
> > On Wed, Oct 28, 2026 at 06:50:12PM +0000, Kuniyuki Iwashima wrote:
> > > On our production fleet, we saw mem_cgroup_sk_charge() spent over
> > > 3x the cycles of try_charge_memcg() although the former should be
> > > a thin wrapper of the latter.
> > > 
> > > Looking at perf during tcp_rr, most of the cycles were spent in
> > > mod_memcg_state(MEMCG_SOCK), which incurred two heavy memory stalls
> > > due to cache line misses. (see patch 4)
> > > 
> > > Also, tcp_rr w/ 20000 flows throughput often degrades by ~50% on
> > > many core machine. (see patch 5)
> > > 
> > > This series adds nr_sock[] to struct memcg_stock_pcp to batch
> > > mod_memcg_state() for MEMCG_SOCK and move heavy operations to
> > > out of local_lock(&memcg_stock.lock).
> > > 
> > > With this series, tcp_rr w/ 10000 flows achieves 31% more
> > > throughput with lower latency on machine w/ AMD EPYC 7B12 64-Core
> > > Processor (Zen 2 Rome, 2 sockets, 256 vCPU)
> > > 
> > >   tcp_rr (10000 flows, 128 threads, 60s, average of 3 runs):
> > > 
> > >   Metric                     |       Baseline |        Patched |    Diff
> > >   ---------------------------+----------------+----------------+--------
> > >   Throughput (RPS)           |     20,884,499 |     27,466,490 | +31.52%
> > >   Latency Mean (us)          |         478.46 |         363.71 | -23.98%
> > >   Latency p50 (us)           |         453.96 |         356.68 | -21.43%
> > >   Latency p90 (us)           |         669.00 |         479.56 | -28.32%
> > >   Latency p99 (us)           |         969.38 |         617.80 | -36.27%
> > > 
> > > 
> > > Changes:
> > >   v2:
> > >     * Rebased on mm.git/mm-new
> > >     * Add patch 5 for flaky throughput issue
> > >     * Patch 4:
> > >       * Benchmark again after rebase and update numbers
> > >       * Check in_nmi() in __refill_stock()
> > > 
> > >   v1: https://lore.kernel.org/linux-mm/20261016052811.3085635-1-kuniyu@google.com/
> > > 
> > > 
> > > Kuniyuki Iwashima (6):
> > >   memcg: Remove unnecessary READ_ONCE() under memcg_stock.lock.
> > >   memcg: Reduce nesting in consume_stock() and refill_stock().
> > >   memcg: Annotate slow paths in consume_stock() and refill_stock().
> > >   memcg: Batch mod_memcg_state() for MEMCG_SOCK.
> > >   memcg: Uncharge pages outside memcg_stock.lock in __refill_stock().
> > >   selftests: cgroup: Deflake test_memcg_sock() in test_memcontrol.c.
> > 
> > Decouple the 4th patch (batch MEMCG_STOCK) from the remaining and send the
> > remaining patches separately and upfront.

Sure.


> > 
> > For the 4th patch, I would prefer we first try to move the network sockets
> > from memcg to objcg path.

sk_forward_alloc is already batched per page, and I think
migrating to objcg would just add an indirect layer and some
costs for css_tryget() / css_put(), and still inlining
__consume_stock() or batching in memcg_stock would be crucial.


> > Then I would prefer we optimize the stats update for
> > all stats instead of special casing for MEMCG_STOCK.
> > 
> 
> MEMCG_SOCK instead of MEMCG_STOCK

I guess you mean optimising MEMCG_SOCK along with MEMCG_KMEM (or
mod_memcg_state() / memcg_rstat_updated() in general).

I can give it a try, though some stats already have specialised
batching (e.g. NR_SLAB_RECLAIMABLE_B in obj_stock and folio stats
in uncharge_batch()).

Re: [PATCH v2 1/6] memcg: Remove unnecessary READ_ONCE() under memcg_stock.lock.

[Kuniyuki Iwashima]

On Wed, 28 Oct 2026 17:48:26 -0700, Shakeel Butt wrote:
> On Wed, Oct 28, 2026 at 06:50:13PM +0000, Kuniyuki Iwashima wrote:
> > this_cpu_ptr(&memcg_stock) is always modified under
> > local_lock(&memcg_stock.lock).
> 
> You need to be more precise here. Mainly the cached[] and nr_pages[] are only
> modified under the local lock. There is flags field which is being modified
> without it.

Will update.

> 
> > 
> > Unlike is_memcg_drain_needed() under RCU, consume_stock(),
> > drain_stock(), and refill_stock() do not need READ_ONCE().
> > 
> > Especially, in refill_stock(), READ_ONCE() is used twice for
> > stock->cached[i], which forces the compiler to emit two volatile
> > loads from stock->cached[i].
> > 
> > Let's remove the unnecessary READ_ONCE() and eliminate the
> > redundant loads in refill_stock():
> 
> What are we gaining with that? Is this measurable?

I haven't measured, but here clang unrolled 7 loops and
each iteration had two loads, where it should be a single load.


> With READ_ONCE() at least we
> are very explicit that this field can read and modified concurrently by
> different CPUs i.e. better readability.

For me, this is rather confusing.

READ_ONCE() means this field is modified concurrently,
and WRITE_ONCE() means this field is read concurrently
(or modified if the writer is not under lock).

So the genral rule is "use READ_ONCE() in a lockless context
and WRITE_ONCE() under lock", and at least we should not use
READ_ONCE() under lock, which adds unnecessary volatile load
and implies non-existing concurrent writer.

Re: [PATCH v2 3/6] memcg: Annotate slow paths in consume_stock() and refill_stock().

[Kuniyuki Iwashima]

On Wed, 28 Oct 2026 17:55:31 -0700, Shakeel Butt wrote:
> On Wed, Oct 28, 2026 at 06:50:15PM +0000, Kuniyuki Iwashima wrote:
> > In consume_stock() and refill_stock(), each per-CPU stock slot
> > holds up to MEMCG_STOCK_HIGH (64) pages, so depleting a slot in
> > consume_stock() (!stock_pages) or trimming a slot in refill_stock()
> > (stock_pages > MEMCG_STOCK_HIGH) happens at most once per batch.
> > 
> > Similarly, requests larger than MEMCG_CHARGE_BATCH, local_trylock()
> > failure, and insufficient cached pages (stock_pages < nr_pages)
> > are all slow paths.
> > 
> > Without unlikely() annotations (when building without FDO, or
> > when an FDO profile does not cover a caller), the compiler treats
> > these branches as likely and inlines css_put() in consume_stock()
> > and memcg_uncharge() in refill_stock() right into the middle of
> > the hot path.
> 
> Please state the compiler along with its version here. I didn't see this behavior
> on gcc 14.

Sure, actually I included asm in v1 but thought no one is
interested in it and this might make my patch look AI-sloppy,
but I'm happy that it wasn't the case :)

https://lore.kernel.org/linux-mm/20261016052811.3085635-4-kuniyu@google.com/

This was generated w/ Clang 23, but I was able to see the
same result w/ Clang 21 too.  I'll include them in v3 again.

Before (refill_stock):

  stock_pages = stock->nr_pages[i] + nr_pages;
   7ab2: 43 0f b6 44 26 01       movzbl  0x1(%r14,%r12), %eax
   7ab8: 01 c6                   addl    %eax, %esi
  if (stock_pages > MEMCG_STOCK_HIGH) {
   7aba: 83 fe 41                cmpl    $0x41, %esi
   7abd: 72 1d                   jb      0x7adc <refill_stock+0x21c>
  memcg_uncharge(memcg, stock_pages - MEMCG_STOCK_LOW);   # inlined
   7abf: 41 89 f7                movl    %esi, %r15d
   7ac2: 41 83 c7 e0             addl    $-0x20, %r15d
   7ac6: 48 8d bb 00 01 00 00    leaq    0x100(%rbx), %rdi
   7acd: 4c 89 fe                movq    %r15, %rsi
   7ad0: e8 00 00 00 00          callq   <page_counter_uncharge>
   7ad5: be 20 00 00 00          movl    $0x20, %esi
   7ada: eb 41                   jmp     0x7b1d <refill_stock+0x25d>
  WRITE_ONCE(stock->nr_pages[i], stock_pages);
   7adc: 43 88 74 26 01          movb    %sil, 0x1(%r14,%r12)
  local_unlock(&memcg_stock.lock);
   7ae1: 65 48 8b 05 00 00 00 00 movq    %gs:(%rip), %rax
   7ae9: c6 80 00 00 00 00 00    movb    $0x0, (%rax)
   7af0: e9 03 fe ff ff          jmp     0x78f8 <refill_stock+0x38>

After (refill_stock):

  stock_pages = stock->nr_pages[i] + nr_pages;
   7b76: 42 0f b6 44 33 01       movzbl  0x1(%rbx,%r14), %eax
   7b7c: 01 c6                   addl    %eax, %esi
  if (unlikely(stock_pages > MEMCG_STOCK_HIGH)) {
   7b7e: 83 fe 41                cmpl    $0x41, %esi
   7b81: 73 20                   jae     0x7ba3 <refill_stock+0x233>
  WRITE_ONCE(stock->nr_pages[i], stock_pages);
   7b83: 42 88 74 33 01          movb    %sil, 0x1(%rbx,%r14)
  local_unlock(&memcg_stock.lock);
   7b88: 65 48 8b 05 00 00 00 00 movq    %gs:(%rip), %rax
   7b90: c6 80 00 00 00 00 00    movb    $0x0, (%rax)
  }
   7b97: 5b                      popq    %rbx
   7b98: 41 5e                   popq    %r14
   7b9a: 41 5f                   popq    %r15
   7b9c: 5d                      popq    %rbp
   7b9d: 2e e9 00 00 00 00       jmp     <__x86_return_thunk>
  memcg_uncharge(memcg, stock_pages - MEMCG_STOCK_LOW);
   7ba3: 83 c6 e0                addl    $-0x20, %esi
   7ba6: e8 15 f2 ff ff          callq   0x6dc0 <memcg_uncharge>
   7bab: be 20 00 00 00          movl    $0x20, %esi
   7bb0: eb d1                   jmp     0x7b83 <refill_stock+0x213>

Re: [PATCH v2 0/6] memcg: Batch mod_memcg_state() for MEMCG_SOCK.

[Kuniyuki Iwashima]

On Thu, 29 Oct 2026 01:35:18 +0000, Kuniyuki Iwashima wrote:
> On Wed, 28 Oct 2026 17:52:32 -0700, Shakeel Butt wrote:
> > On Wed, Oct 28, 2026 at 05:41:04PM -0700, Shakeel Butt wrote:
> > > On Wed, Oct 28, 2026 at 06:50:12PM +0000, Kuniyuki Iwashima wrote:
> > > > On our production fleet, we saw mem_cgroup_sk_charge() spent over
> > > > 3x the cycles of try_charge_memcg() although the former should be
> > > > a thin wrapper of the latter.
> > > > 
> > > > Looking at perf during tcp_rr, most of the cycles were spent in
> > > > mod_memcg_state(MEMCG_SOCK), which incurred two heavy memory stalls
> > > > due to cache line misses. (see patch 4)
> > > > 
> > > > Also, tcp_rr w/ 20000 flows throughput often degrades by ~50% on
> > > > many core machine. (see patch 5)
> > > > 
> > > > This series adds nr_sock[] to struct memcg_stock_pcp to batch
> > > > mod_memcg_state() for MEMCG_SOCK and move heavy operations to
> > > > out of local_lock(&memcg_stock.lock).
> > > > 
> > > > With this series, tcp_rr w/ 10000 flows achieves 31% more
> > > > throughput with lower latency on machine w/ AMD EPYC 7B12 64-Core
> > > > Processor (Zen 2 Rome, 2 sockets, 256 vCPU)
> > > > 
> > > >   tcp_rr (10000 flows, 128 threads, 60s, average of 3 runs):
> > > > 
> > > >   Metric                     |       Baseline |        Patched |    Diff
> > > >   ---------------------------+----------------+----------------+--------
> > > >   Throughput (RPS)           |     20,884,499 |     27,466,490 | +31.52%
> > > >   Latency Mean (us)          |         478.46 |         363.71 | -23.98%
> > > >   Latency p50 (us)           |         453.96 |         356.68 | -21.43%
> > > >   Latency p90 (us)           |         669.00 |         479.56 | -28.32%
> > > >   Latency p99 (us)           |         969.38 |         617.80 | -36.27%
> > > > 
> > > > 
> > > > Changes:
> > > >   v2:
> > > >     * Rebased on mm.git/mm-new
> > > >     * Add patch 5 for flaky throughput issue
> > > >     * Patch 4:
> > > >       * Benchmark again after rebase and update numbers
> > > >       * Check in_nmi() in __refill_stock()
> > > > 
> > > >   v1: https://lore.kernel.org/linux-mm/20261016052811.3085635-1-kuniyu@google.com/
> > > > 
> > > > 
> > > > Kuniyuki Iwashima (6):
> > > >   memcg: Remove unnecessary READ_ONCE() under memcg_stock.lock.
> > > >   memcg: Reduce nesting in consume_stock() and refill_stock().
> > > >   memcg: Annotate slow paths in consume_stock() and refill_stock().
> > > >   memcg: Batch mod_memcg_state() for MEMCG_SOCK.
> > > >   memcg: Uncharge pages outside memcg_stock.lock in __refill_stock().
> > > >   selftests: cgroup: Deflake test_memcg_sock() in test_memcontrol.c.
> > > 
> > > Decouple the 4th patch (batch MEMCG_STOCK) from the remaining and send the
> > > remaining patches separately and upfront.
> 
> Sure.
> 
> 
> > > 
> > > For the 4th patch, I would prefer we first try to move the network sockets
> > > from memcg to objcg path.
> 
> sk_forward_alloc is already batched per page, and I think
> migrating to objcg would just add an indirect layer and some
> costs for css_tryget() / css_put(), and still inlining
> __consume_stock() or batching in memcg_stock would be crucial.

Interestingly, patch 5 only (without patch 4) gives even more
throughput gain on the platform mentioned in patch 4.

But as expected, __css_rstat_updated() spends more cycles.

  tcp_rr (10000 flows, 128 threads, 60s, average of 3 runs):

  Metric                     |       Baseline |   Patched  |    Diff
  ---------------------------+----------------+---------------------+--------
  Throughput (RPS)           |     20,884,499 | 28,453,051 | +36.24%
  Latency Mean (us)          |         478.46 |     351.07 | -26.63%
  Latency p50 (us)           |         453.96 |     343.04 | -24.43%
  Latency p90 (us)           |         669.00 |     479.57 | -28.32%
  Latency p99 (us)           |         969.38 |     593.92 | -38.73%

  perf cycles (10000 flows, 10s, average of 3 runs):

  Server (M cycles / %)      |           Baseline |            Patched |    Diff
  ---------------------------+--------------------+-----------------------+--------
  Host Total Cycles          |        4,566,914 M |        4,548,747 M |   -0.4%
  tcp_rcv_established()      | 158,009 M (3.460%) | 158,773 M (3.490%) |   +0.5%
  mem_cgroup_sk_charge()     |  18,420 M (0.403%) |  16,211 M (0.356%) |  -12.0%
  mem_cgroup_sk_uncharge()   |  16,593 M (0.363%) |  12,121 M (0.266%) |  -27.0%
  __css_rstat_updated()      |   3,502 M (0.077%) |   5,606 M (0.123%) |  +60.0%
  Memcg Sock Total           |  38,515 M (0.843%) |  33,937 M (0.746%) |  -11.9%


> 
> 
> > > Then I would prefer we optimize the stats update for
> > > all stats instead of special casing for MEMCG_STOCK.
> > > 
> > 
> > MEMCG_SOCK instead of MEMCG_STOCK
> 
> I guess you mean optimising MEMCG_SOCK along with MEMCG_KMEM (or
> mod_memcg_state() / memcg_rstat_updated() in general).
> 
> I can give it a try, though some stats already have specialised
> batching (e.g. NR_SLAB_RECLAIMABLE_B in obj_stock and folio stats
> in uncharge_batch()).

Re: [PATCH v2 0/6] memcg: Batch mod_memcg_state() for MEMCG_SOCK.

[Shakeel Butt]

On Thu, Oct 29, 2026 at 01:35:18AM +0000, Kuniyuki Iwashima wrote:
> On Wed, 28 Oct 2026 17:52:32 -0700, Shakeel Butt wrote:
[...]
> > > 
> > > For the 4th patch, I would prefer we first try to move the network sockets
> > > from memcg to objcg path.
> 
> sk_forward_alloc is already batched per page, and I think
> migrating to objcg would just add an indirect layer and some
> costs for css_tryget() / css_put(), and still inlining
> __consume_stock() or batching in memcg_stock would be crucial.
> 

We are over time moving kernel objects holding direct reference to memcg to the
objcg infra and the reason is such objects can become source of zombie memcgs.
There is no urgency for the network sockets but eventually it will happen. We
will see if indirection is really measurable at that time.