[PATCH bpf-next 2/6] bpf, net: rework cookie generator as per-cpu one

[Daniel Borkmann <daniel@xxxxxxxxxxxxx>]

With its use in BPF the cookie generator can be called very frequently
in particular when used out of cgroup v2 hooks (e.g. connect / sendmsg)
and attached to the root cgroup, for example, when used in v1/v2 mixed
environments. In particular when there's a high churn on sockets in the
system there can be many parallel requests to the bpf_get_socket_cookie()
and bpf_get_netns_cookie() helpers which then cause contention on the
shared atomic counter. As similarly done in f991bd2e1421 ("fs: introduce
a per-cpu last_ino allocator"), add a small helper library that both can
then use for the 64 bit counters.

Signed-off-by: Daniel Borkmann <daniel@xxxxxxxxxxxxx>
---
 include/linux/cookie.h   | 41 ++++++++++++++++++++++++++++++++++++++++
 net/core/net_namespace.c |  5 +++--
 net/core/sock_diag.c     |  7 ++++---
 3 files changed, 48 insertions(+), 5 deletions(-)
 create mode 100644 include/linux/cookie.h

diff --git a/include/linux/cookie.h b/include/linux/cookie.h
new file mode 100644
index 000000000000..2488203dc004
--- /dev/null
+++ b/include/linux/cookie.h
@@ -0,0 +1,41 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+#ifndef __LINUX_COOKIE_H
+#define __LINUX_COOKIE_H
+
+#include <linux/atomic.h>
+#include <linux/percpu.h>
+
+struct gen_cookie {
+	u64 __percpu	*local_last;
+	atomic64_t	 shared_last ____cacheline_aligned_in_smp;
+};
+
+#define COOKIE_LOCAL_BATCH	4096
+
+#define DEFINE_COOKIE(name)					\
+	static DEFINE_PER_CPU(u64, __##name);			\
+	static struct gen_cookie name = {			\
+		.local_last	= &__##name,			\
+		.shared_last	= ATOMIC64_INIT(0),		\
+	}
+
+static inline u64 gen_cookie_next(struct gen_cookie *gc)
+{
+	u64 *local_last = &get_cpu_var(*gc->local_last);
+	u64 val = *local_last;
+
+	if (__is_defined(CONFIG_SMP) &&
+	    unlikely((val & (COOKIE_LOCAL_BATCH - 1)) == 0)) {
+		s64 next = atomic64_add_return(COOKIE_LOCAL_BATCH,
+					       &gc->shared_last);
+		val = next - COOKIE_LOCAL_BATCH;
+	}
+	val++;
+	if (unlikely(!val))
+		val++;
+	*local_last = val;
+	put_cpu_var(local_last);
+	return val;
+}
+
+#endif /* __LINUX_COOKIE_H */
diff --git a/net/core/net_namespace.c b/net/core/net_namespace.c
index dcd61aca343e..cf29ed8950d1 100644
--- a/net/core/net_namespace.c
+++ b/net/core/net_namespace.c
@@ -19,6 +19,7 @@
 #include <linux/net_namespace.h>
 #include <linux/sched/task.h>
 #include <linux/uidgid.h>
+#include <linux/cookie.h>
 
 #include <net/sock.h>
 #include <net/netlink.h>
@@ -69,7 +70,7 @@ EXPORT_SYMBOL_GPL(pernet_ops_rwsem);
 
 static unsigned int max_gen_ptrs = INITIAL_NET_GEN_PTRS;
 
-static atomic64_t cookie_gen;
+DEFINE_COOKIE(net_cookie);
 
 u64 net_gen_cookie(struct net *net)
 {
@@ -78,7 +79,7 @@ u64 net_gen_cookie(struct net *net)
 
 		if (res)
 			return res;
-		res = atomic64_inc_return(&cookie_gen);
+		res = gen_cookie_next(&net_cookie);
 		atomic64_cmpxchg(&net->net_cookie, 0, res);
 	}
 }
diff --git a/net/core/sock_diag.c b/net/core/sock_diag.c
index c13ffbd33d8d..4a4180e8dd35 100644
--- a/net/core/sock_diag.c
+++ b/net/core/sock_diag.c
@@ -11,7 +11,7 @@
 #include <linux/tcp.h>
 #include <linux/workqueue.h>
 #include <linux/nospec.h>
-
+#include <linux/cookie.h>
 #include <linux/inet_diag.h>
 #include <linux/sock_diag.h>
 
@@ -19,7 +19,8 @@ static const struct sock_diag_handler *sock_diag_handlers[AF_MAX];
 static int (*inet_rcv_compat)(struct sk_buff *skb, struct nlmsghdr *nlh);
 static DEFINE_MUTEX(sock_diag_table_mutex);
 static struct workqueue_struct *broadcast_wq;
-static atomic64_t cookie_gen;
+
+DEFINE_COOKIE(sock_cookie);
 
 u64 sock_gen_cookie(struct sock *sk)
 {
@@ -28,7 +29,7 @@ u64 sock_gen_cookie(struct sock *sk)
 
 		if (res)
 			return res;
-		res = atomic64_inc_return(&cookie_gen);
+		res = gen_cookie_next(&sock_cookie);
 		atomic64_cmpxchg(&sk->sk_cookie, 0, res);
 	}
 }
-- 
2.21.0