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

[Kuniyuki Iwashima <kuniyu@google.com>]

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