[PATCH net] net: memcontrol: charge allocated memory after mem_cgroup_sk_alloc()

[Roman Gushchin <guro@fb.com>]

We've catched several cgroup css refcounting issues on 4.15-rc7,
triggered from different release paths. We've used cgroups v2.
I've added a temporarily per-memcg sockmem atomic counter,
and found, that we're sometimes falling below 0. It was easy
to reproduce, so I was able to bisect the problem.

It was introduced by the commit 9f1c2674b328 ("net: memcontrol:
defer call to mem_cgroup_sk_alloc()"), which moved
the mem_cgroup_sk_alloc() call from the BH context
into inet_csk_accept().

The problem is that all the memory allocated before
mem_cgroup_sk_alloc() is charged to the socket,
but not charged to the memcg. So, when we're releasing
the socket, we're uncharging more, than we've charged.

Fix this by charging the cgroup by the amount of already
allocated memory right after mem_cgroup_sk_alloc() in
inet_csk_accept().

Fixes: 9f1c2674b328 ("net: memcontrol: defer call to mem_cgroup_sk_alloc()")
Signed-off-by: Roman Gushchin <guro@fb.com>
Cc: Eric Dumazet <edumazet@google.com>
Cc: Johannes Weiner <hannes@cmpxchg.org>
Cc: Tejun Heo <tj@kernel.org>
Cc: David S. Miller <davem@davemloft.net>
---
 net/ipv4/inet_connection_sock.c | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/net/ipv4/inet_connection_sock.c b/net/ipv4/inet_connection_sock.c
index 4ca46dc08e63..f439162c2ea2 100644
--- a/net/ipv4/inet_connection_sock.c
+++ b/net/ipv4/inet_connection_sock.c
@@ -434,6 +434,7 @@ struct sock *inet_csk_accept(struct sock *sk, int flags, int *err, bool kern)
 	struct request_sock *req;
 	struct sock *newsk;
 	int error;
+	long amt;
 
 	lock_sock(sk);
 
@@ -476,6 +477,10 @@ struct sock *inet_csk_accept(struct sock *sk, int flags, int *err, bool kern)
 		spin_unlock_bh(&queue->fastopenq.lock);
 	}
 	mem_cgroup_sk_alloc(newsk);
+	amt = sk_memory_allocated(newsk);
+	if (amt && newsk->sk_memcg)
+		mem_cgroup_charge_skmem(newsk->sk_memcg, amt);
+
 out:
 	release_sock(sk);
 	if (req)
-- 
2.14.3