[PATCH RFC -next 2/5] tcp: sync memcg budget on protocol buffer updates

[Cai Xinchen <caixinchen1@huawei.com>]

TCP rewrites the socket buffers after sock_init_data() already charged
the generic defaults, and keeps rewriting them during autotuning:

  - tcp_init_sock() overrides sk_sndbuf/sk_rcvbuf with the per-net
    tcp_wmem/tcp_rmem sysctl defaults,
  - tcp_set_rcvlowat() raises sk_rcvbuf from the requested rcv window,
  - tcp_sndbuf_expand(), tcp_clamp_window() and tcp_rcvbuf_grow()
    grow the buffers during autotuning, and
  - tcp_should_expand_sndbuf() caps sk_sndbuf under memory pressure.

Sync the memcg budget tracker with sk_memcg_budget_sync() at each of
these writes so the charged amount follows the actual budget: growth
is charged to the memcg, shrink is refunded.

Assisted-by: opencode:glm-5.3
Signed-off-by: Cai Xinchen <caixinchen1@huawei.com>
---
 net/ipv4/tcp.c       |  5 +++++
 net/ipv4/tcp_input.c | 10 ++++++++--
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/net/ipv4/tcp.c b/net/ipv4/tcp.c
index 3c9e1a88a6bd..0d2a39672aa0 100644
--- a/net/ipv4/tcp.c
+++ b/net/ipv4/tcp.c
@@ -471,6 +471,10 @@ void tcp_init_sock(struct sock *sk)
 
 	WRITE_ONCE(sk->sk_sndbuf, READ_ONCE(sock_net(sk)->ipv4.sysctl_tcp_wmem[1]));
 	WRITE_ONCE(sk->sk_rcvbuf, READ_ONCE(sock_net(sk)->ipv4.sysctl_tcp_rmem[1]));
+	/* The default buffers grew from the generic sock_init_data()
+	 * values: charge the difference to the memcg.
+	 */
+	sk_memcg_budget_sync(sk, gfp_memcg_charge());
 	tcp_scaling_ratio_init(sk);
 
 	set_bit(SOCK_SUPPORT_ZC, &sk->sk_socket->flags);
@@ -1847,6 +1851,7 @@ int tcp_set_rcvlowat(struct sock *sk, int val)
 	space = tcp_space_from_win(sk, val);
 	if (space > sk->sk_rcvbuf) {
 		WRITE_ONCE(sk->sk_rcvbuf, space);
+		sk_memcg_budget_sync(sk, gfp_memcg_charge());
 
 		if (tp->window_clamp && tp->window_clamp < val)
 			WRITE_ONCE(tp->window_clamp, val);
diff --git a/net/ipv4/tcp_input.c b/net/ipv4/tcp_input.c
index 0f60a1dbf927..8abec8ebf916 100644
--- a/net/ipv4/tcp_input.c
+++ b/net/ipv4/tcp_input.c
@@ -629,9 +629,11 @@ static void tcp_sndbuf_expand(struct sock *sk)
 	sndmem = ca_ops->sndbuf_expand ? ca_ops->sndbuf_expand(sk) : 2;
 	sndmem *= nr_segs * per_mss;
 
-	if (sk->sk_sndbuf < sndmem)
+	if (sk->sk_sndbuf < sndmem) {
 		WRITE_ONCE(sk->sk_sndbuf,
 			   min(sndmem, READ_ONCE(sock_net(sk)->ipv4.sysctl_tcp_wmem[2])));
+		sk_memcg_budget_sync(sk, gfp_memcg_charge());
+	}
 }
 
 /* 2. Tuning advertised window (window_clamp, rcv_ssthresh)
@@ -791,6 +793,7 @@ static void tcp_clamp_window(struct sock *sk)
 	    sk_memory_allocated(sk) < sk_prot_mem_limits(sk, 0)) {
 		WRITE_ONCE(sk->sk_rcvbuf,
 			   min(atomic_read(&sk->sk_rmem_alloc), rmem2));
+		sk_memcg_budget_sync(sk, gfp_memcg_charge());
 	}
 	if (atomic_read(&sk->sk_rmem_alloc) > sk->sk_rcvbuf)
 		tp->rcv_ssthresh = min(tp->window_clamp, 2U * tp->advmss);
@@ -948,6 +951,7 @@ void tcp_rcvbuf_grow(struct sock *sk, u32 newval)
 	rcvbuf = min_t(u32, tcp_space_from_win(sk, rcvwin), cap);
 	if (rcvbuf > sk->sk_rcvbuf) {
 		WRITE_ONCE(sk->sk_rcvbuf, rcvbuf);
+		sk_memcg_budget_sync(sk, gfp_memcg_charge());
 		/* Make the window clamp follow along.  */
 		WRITE_ONCE(tp->window_clamp,
 			   tcp_win_from_space(sk, rcvbuf));
@@ -6055,8 +6059,10 @@ static bool tcp_should_expand_sndbuf(struct sock *sk)
 		 * it never goes below SOCK_MIN_SNDBUF.
 		 * See sk_stream_moderate_sndbuf() for more details.
 		 */
-		if (unused_mem > SOCK_MIN_SNDBUF)
+		if (unused_mem > SOCK_MIN_SNDBUF) {
 			WRITE_ONCE(sk->sk_sndbuf, unused_mem);
+			sk_memcg_budget_sync(sk, gfp_memcg_charge());
+		}
 
 		return false;
 	}
-- 
2.18.0.huawei.25