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

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@iogearbox.net>
---
 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;