[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()).