[PATCH net-next 02/27] sock: optimise sock_def_write_space send refcounting

[Pavel Begunkov <asml.silence@gmail.com>]

sock_def_write_space() is extensively used by UDP and there is some
room for optimisation. When sock_wfree() needs to do ->sk_write_space(),
it modifies ->sk_wmem_alloc in two steps. First, it puts all but one
refs and calls ->sk_write_space(), and then puts down remaining 1.
That's needed because the callback relies on ->sk_wmem_alloc being
subbed but something should hold the socket alive.

The idea behind this patch is to take advantage of SOCK_RCU_FREE and
ensure the socket is not freed by wrapping ->sk_write_space() in an RCU
section. Then we can remove one extra refcount atomic.

Note: not all callbacks might be RCU prepared, so we carve out a
sock_def_write_space() specific path.

Signed-off-by: Pavel Begunkov <asml.silence@gmail.com>
---
 net/core/sock.c | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/net/core/sock.c b/net/core/sock.c
index f5766d6e27cb..9389bb602c64 100644
--- a/net/core/sock.c
+++ b/net/core/sock.c
@@ -144,6 +144,8 @@
 static DEFINE_MUTEX(proto_list_mutex);
 static LIST_HEAD(proto_list);
 
+static void sock_def_write_space(struct sock *sk);
+
 /**
  * sk_ns_capable - General socket capability test
  * @sk: Socket to use a capability on or through
@@ -2300,8 +2302,20 @@ void sock_wfree(struct sk_buff *skb)
 {
 	struct sock *sk = skb->sk;
 	unsigned int len = skb->truesize;
+	bool free;
 
 	if (!sock_flag(sk, SOCK_USE_WRITE_QUEUE)) {
+		if (sock_flag(sk, SOCK_RCU_FREE) &&
+		    sk->sk_write_space == sock_def_write_space) {
+			rcu_read_lock();
+			free = refcount_sub_and_test(len, &sk->sk_wmem_alloc);
+			sock_def_write_space(sk);
+			rcu_read_unlock();
+			if (unlikely(free))
+				__sk_free(sk);
+			return;
+		}
+
 		/*
 		 * Keep a reference on sk_wmem_alloc, this will be released
 		 * after sk_write_space() call
-- 
2.35.1