[PATCH v8 1/5] tracing: Add lock-free stackmap for stack trace deduplication

[Li Pengfei <ljdlns1987@gmail.com>]

From: Pengfei Li <lipengfei28@xiaomi.com>

A kernel stack costs 80-160 bytes per ftrace event, and the same stacks
recur. Add a preallocated hash map that stores each distinct stack once
and returns a 64-bit cookie, made of the map generation and the table
slot, that the stack recording path can record instead of the stack.

Lookup and insert are lock-free and may run in NMI context, so the map
depends on ARCH_HAVE_NMI_SAFE_CMPXCHG like tracing_map. Reset clears the
map in place and bumps the generation so that old cookies stop matching.
The stack_map file that resolves cookies is created by the next patch.

Assisted-by: LLM
Signed-off-by: Pengfei Li <lipengfei28@xiaomi.com>
---
 kernel/trace/Kconfig          |  13 ++
 kernel/trace/Makefile         |   1 +
 kernel/trace/trace.h          |   3 +
 kernel/trace/trace_stackmap.c | 382 ++++++++++++++++++++++++++++++++++
 kernel/trace/trace_stackmap.h |  22 ++
 5 files changed, 421 insertions(+)
 create mode 100644 kernel/trace/trace_stackmap.c
 create mode 100644 kernel/trace/trace_stackmap.h

diff --git a/kernel/trace/Kconfig b/kernel/trace/Kconfig
index 084f34dc6c9f..a83c53c02eea 100644
--- a/kernel/trace/Kconfig
+++ b/kernel/trace/Kconfig
@@ -412,6 +412,19 @@ config STACK_TRACER
 
 	  Say N if unsure.
 
+config FTRACE_STACKMAP
+	bool "Ftrace stack map deduplication"
+	depends on TRACING
+	depends on STACKTRACE
+	depends on ARCH_HAVE_NMI_SAFE_CMPXCHG
+	help
+	  Store each distinct kernel stack recorded by ftrace once, in a
+	  table allocated when options/stackmap is enabled, and record a
+	  64-bit cookie in the ring buffer instead of the full stack.
+	  See Documentation/trace/ftrace-stackmap.rst.
+
+	  If unsure, say N.
+
 config TRACE_PREEMPT_TOGGLE
 	bool
 	help
diff --git a/kernel/trace/Makefile b/kernel/trace/Makefile
index f934ff586bd4..7710ec2659e9 100644
--- a/kernel/trace/Makefile
+++ b/kernel/trace/Makefile
@@ -86,6 +86,7 @@ obj-$(CONFIG_HWLAT_TRACER) += trace_hwlat.o
 obj-$(CONFIG_OSNOISE_TRACER) += trace_osnoise.o
 obj-$(CONFIG_NOP_TRACER) += trace_nop.o
 obj-$(CONFIG_STACK_TRACER) += trace_stack.o
+obj-$(CONFIG_FTRACE_STACKMAP) += trace_stackmap.o
 obj-$(CONFIG_MMIOTRACE) += trace_mmiotrace.o
 obj-$(CONFIG_FUNCTION_GRAPH_TRACER) += trace_functions_graph.o
 obj-$(CONFIG_TRACE_BRANCH_PROFILING) += trace_branch.o
diff --git a/kernel/trace/trace.h b/kernel/trace/trace.h
index 80fe152af1dd..7b2cdc0dc253 100644
--- a/kernel/trace/trace.h
+++ b/kernel/trace/trace.h
@@ -453,6 +453,9 @@ struct trace_array {
 	struct cond_snapshot	*cond_snapshot;
 #endif
 	struct trace_func_repeats	__percpu *last_func_repeats;
+#ifdef CONFIG_FTRACE_STACKMAP
+	struct ftrace_stackmap __rcu	*stackmap;
+#endif
 	/*
 	 * On boot up, the ring buffer is set to the minimum size, so that
 	 * we do not waste memory on systems that are not using tracing.
diff --git a/kernel/trace/trace_stackmap.c b/kernel/trace/trace_stackmap.c
new file mode 100644
index 000000000000..e5ffd1850ebf
--- /dev/null
+++ b/kernel/trace/trace_stackmap.c
@@ -0,0 +1,382 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Stack trace deduplication for ftrace: each distinct stack is stored once,
+ * and the ring buffer records a cookie (map generation << 32 | slot).
+ */
+
+#include <linux/kernel.h>
+#include <linux/slab.h>
+#include <linux/jhash.h>
+#include <linux/seq_file.h>
+#include <linux/vmalloc.h>
+#include <linux/atomic.h>
+#include <linux/random.h>
+#include <linux/rcupdate.h>
+#include <linux/refcount.h>
+
+#include "trace.h"
+#include "trace_stackmap.h"
+
+#define FTRACE_STACKMAP_BITS		14
+#define FTRACE_STACKMAP_MAX_PROBE	64
+
+struct stackmap_elt {
+	u32		nr;
+	atomic_t	ref_count;
+	unsigned long	ips[FTRACE_STACKMAP_MAX_DEPTH];
+};
+
+struct stackmap_entry {
+	u32			key;
+	struct stackmap_elt	*val;
+};
+
+struct ftrace_stackmap {
+	unsigned int		map_bits;
+	unsigned int		map_size;
+	unsigned int		max_elts;
+	u32			hash_seed;
+	atomic_t		next_elt;
+	struct stackmap_entry	*entries;
+	struct stackmap_elt	*elts;
+	atomic_t		resetting;
+	u32			generation;
+	refcount_t		ref;
+};
+
+/* Serializes tracefs readers against reset. */
+static DEFINE_MUTEX(stackmap_lock);
+
+static u32 stackmap_next_gen = 1;
+
+static struct ftrace_stackmap *stackmap_locked(struct trace_array *tr)
+{
+	return rcu_dereference_protected(tr->stackmap,
+					 lockdep_is_held(&trace_types_lock));
+}
+
+static struct stackmap_elt *stackmap_load_elt(struct stackmap_entry *entry)
+{
+	/* Pairs with smp_store_release() in ftrace_stackmap_get_cookie(). */
+	return smp_load_acquire(&entry->val);
+}
+
+static struct stackmap_elt *stackmap_get_elt(struct ftrace_stackmap *smap)
+{
+	int idx;
+
+	idx = atomic_fetch_add_unless(&smap->next_elt, 1, smap->max_elts);
+	return idx < smap->max_elts ? &smap->elts[idx] : NULL;
+}
+
+static void ftrace_stackmap_destroy(struct ftrace_stackmap *smap)
+{
+	vfree(smap->elts);
+	vfree(smap->entries);
+	kfree(smap);
+}
+
+static struct ftrace_stackmap *ftrace_stackmap_create(void)
+{
+	struct ftrace_stackmap *smap;
+
+	smap = kzalloc_obj(*smap, GFP_KERNEL);
+	if (!smap)
+		return ERR_PTR(-ENOMEM);
+
+	smap->map_bits = FTRACE_STACKMAP_BITS;
+	smap->max_elts = 1U << FTRACE_STACKMAP_BITS;
+	smap->map_size = 1U << (FTRACE_STACKMAP_BITS + 1);
+
+	smap->entries = __vcalloc(smap->map_size, sizeof(*smap->entries),
+				  GFP_KERNEL | __GFP_RETRY_MAYFAIL | __GFP_NOWARN);
+	if (!smap->entries)
+		goto fail;
+
+	smap->elts = __vcalloc(smap->max_elts, sizeof(*smap->elts),
+			       GFP_KERNEL | __GFP_RETRY_MAYFAIL | __GFP_NOWARN);
+	if (!smap->elts)
+		goto fail;
+
+	smap->hash_seed = get_random_u32();
+	smap->generation = stackmap_next_gen;
+	refcount_set(&smap->ref, 1);
+	return smap;
+
+fail:
+	ftrace_stackmap_destroy(smap);
+	return ERR_PTR(-ENOMEM);
+}
+
+int ftrace_stackmap_alloc(struct trace_array *tr)
+{
+	struct ftrace_stackmap *smap;
+
+	if (stackmap_locked(tr))
+		return 0;
+
+	smap = ftrace_stackmap_create();
+	if (IS_ERR(smap))
+		return PTR_ERR(smap);
+
+	rcu_assign_pointer(tr->stackmap, smap);
+	return 0;
+}
+
+static struct ftrace_stackmap *stackmap_get(struct trace_array *tr)
+{
+	struct ftrace_stackmap *smap;
+
+	guard(mutex)(&trace_types_lock);
+	smap = stackmap_locked(tr);
+	if (smap)
+		refcount_inc(&smap->ref);
+	return smap;
+}
+
+static void stackmap_put(struct ftrace_stackmap *smap)
+{
+	if (!smap)
+		return;
+	if (refcount_dec_and_test(&smap->ref))
+		ftrace_stackmap_destroy(smap);
+}
+
+static void stackmap_reset(struct trace_array *tr)
+{
+	struct ftrace_stackmap *smap = stackmap_locked(tr);
+
+	if (!smap)
+		return;
+
+	/* Keep new callers out, then wait for the ones already inside. */
+	atomic_set(&smap->resetting, 1);
+	synchronize_rcu();
+
+	scoped_guard(mutex, &stackmap_lock) {
+		memset(smap->entries, 0, sizeof(*smap->entries) * smap->map_size);
+		memset(smap->elts, 0, sizeof(*smap->elts) * smap->max_elts);
+		atomic_set(&smap->next_elt, 0);
+		smap->generation++;
+	}
+
+	/* Pairs with atomic_read_acquire() in ftrace_stackmap_get_cookie(). */
+	atomic_set_release(&smap->resetting, 0);
+}
+
+bool ftrace_stackmap_get_cookie(struct ftrace_stackmap *smap,
+				unsigned long *ips, unsigned int nr_entries,
+				u64 *cookie)
+{
+	u32 key_hash, idx, test_key, trace_len;
+	struct stackmap_entry *entry;
+	struct stackmap_elt *val;
+	int probes = 0;
+
+	/* Pairs with atomic_set_release() in stackmap_reset(). */
+	if (atomic_read_acquire(&smap->resetting))
+		return false;
+
+	trace_len = nr_entries * sizeof(unsigned long);
+	key_hash = jhash2((const u32 *)ips, trace_len / sizeof(u32),
+			  smap->hash_seed);
+	/* Zero marks a free slot. */
+	if (key_hash == 0)
+		key_hash = 1;
+
+	idx = key_hash >> (32 - (smap->map_bits + 1));
+
+	while (probes < FTRACE_STACKMAP_MAX_PROBE) {
+		idx &= smap->map_size - 1;
+		entry = &smap->entries[idx];
+		test_key = READ_ONCE(entry->key);
+
+		if (test_key == key_hash) {
+			val = stackmap_load_elt(entry);
+			/* Insert in progress; from NMI it may be the interrupted one. */
+			if (!val)
+				return false;
+			if (READ_ONCE(val->nr) == nr_entries &&
+			    memcmp(val->ips, ips, trace_len) == 0) {
+				atomic_add_unless(&val->ref_count, 1, INT_MAX);
+				goto found;
+			}
+		} else if (!test_key) {
+			if (atomic_read(&smap->next_elt) >= smap->max_elts)
+				break;
+			/* Lost the race; the slot may now hold this stack. */
+			if (cmpxchg(&entry->key, 0, key_hash) != 0)
+				continue;
+
+			val = stackmap_get_elt(smap);
+			if (!val)
+				break;
+
+			val->nr = nr_entries;
+			atomic_set(&val->ref_count, 1);
+			memcpy(val->ips, ips, trace_len);
+			/* Pairs with stackmap_load_elt(). */
+			smp_store_release(&entry->val, val);
+			goto found;
+		}
+
+		idx++;
+		probes++;
+	}
+
+	return false;
+found:
+	*cookie = ((u64)smap->generation << FTRACE_STACKMAP_COOKIE_GEN_SHIFT) | idx;
+	return true;
+}
+
+struct stackmap_seq_private {
+	struct trace_array	*tr;
+	struct ftrace_stackmap	*smap;
+};
+
+static struct stackmap_entry *stackmap_seq_seek(struct ftrace_stackmap *smap,
+						loff_t *pos)
+{
+	loff_t i;
+
+	for (i = *pos; i < smap->map_size; i++) {
+		if (READ_ONCE(smap->entries[i].key) &&
+		    stackmap_load_elt(&smap->entries[i])) {
+			*pos = i;
+			return &smap->entries[i];
+		}
+	}
+	/* Advance past the end or seq_read() repeats the last entry. */
+	*pos = smap->map_size;
+	return NULL;
+}
+
+static void *stackmap_seq_start(struct seq_file *m, loff_t *pos)
+{
+	struct stackmap_seq_private *priv = m->private;
+
+	mutex_lock(&stackmap_lock);
+	if (!priv->smap)
+		return NULL;
+	return stackmap_seq_seek(priv->smap, pos);
+}
+
+static void *stackmap_seq_next(struct seq_file *m, void *v, loff_t *pos)
+{
+	struct stackmap_seq_private *priv = m->private;
+
+	(*pos)++;
+	return stackmap_seq_seek(priv->smap, pos);
+}
+
+static void stackmap_seq_stop(struct seq_file *m, void *v)
+{
+	/* seq_file calls stop() even when start() returned NULL. */
+	mutex_unlock(&stackmap_lock);
+}
+
+static int stackmap_seq_show(struct seq_file *m, void *v)
+{
+	struct stackmap_entry *entry = v;
+	struct stackmap_seq_private *priv = m->private;
+	struct stackmap_elt *elt;
+	u32 idx = entry - priv->smap->entries;
+	u32 i, nr;
+
+	elt = stackmap_load_elt(entry);
+	if (!elt)
+		return 0;
+
+	nr = READ_ONCE(elt->nr);
+	if (nr > FTRACE_STACKMAP_MAX_DEPTH)
+		nr = FTRACE_STACKMAP_MAX_DEPTH;
+
+	seq_printf(m, "stack_id %u gen %u ref %u depth %u\n",
+		   idx, priv->smap->generation, atomic_read(&elt->ref_count), nr);
+	for (i = 0; i < nr; i++) {
+		unsigned long ip = elt->ips[i];
+
+		if (ip == FTRACE_TRAMPOLINE_MARKER)
+			seq_puts(m, " => [FTRACE TRAMPOLINE]\n");
+		else
+			seq_printf(m, " => %pS\n", (void *)ip);
+	}
+	return 0;
+}
+
+static const struct seq_operations stackmap_seq_ops = {
+	.start	= stackmap_seq_start,
+	.next	= stackmap_seq_next,
+	.stop	= stackmap_seq_stop,
+	.show	= stackmap_seq_show,
+};
+
+static int stackmap_open(struct inode *inode, struct file *file)
+{
+	struct trace_array *tr = inode->i_private;
+	struct stackmap_seq_private *priv;
+	struct seq_file *m;
+	int ret;
+
+	ret = tracing_check_open_get_tr(tr);
+	if (ret)
+		return ret;
+
+	ret = seq_open_private(file, &stackmap_seq_ops,
+			       sizeof(struct stackmap_seq_private));
+	if (ret) {
+		trace_array_put(tr);
+		return ret;
+	}
+
+	m = file->private_data;
+	priv = m->private;
+	priv->tr = tr;
+	if (file->f_mode & FMODE_READ)
+		priv->smap = stackmap_get(tr);
+	return 0;
+}
+
+static int stackmap_release(struct inode *inode, struct file *file)
+{
+	struct seq_file *m = file->private_data;
+	struct stackmap_seq_private *priv = m->private;
+	struct ftrace_stackmap *smap = priv->smap;
+	struct trace_array *tr = priv->tr;
+	int ret;
+
+	ret = seq_release_private(inode, file);
+	stackmap_put(smap);
+	trace_array_put(tr);
+	return ret;
+}
+
+static ssize_t stackmap_write(struct file *file, const char __user *ubuf,
+			      size_t count, loff_t *ppos)
+{
+	struct seq_file *m = file->private_data;
+	struct stackmap_seq_private *priv = m->private;
+	char buf[8];
+
+	if (count >= sizeof(buf))
+		return -EINVAL;
+	if (copy_from_user(buf, ubuf, count))
+		return -EFAULT;
+	buf[count] = '\0';
+
+	if (sysfs_streq(buf, "0")) {
+		guard(mutex)(&trace_types_lock);
+		stackmap_reset(priv->tr);
+		return count;
+	}
+	return -EINVAL;
+}
+
+const struct file_operations ftrace_stackmap_fops = {
+	.open		= stackmap_open,
+	.read		= seq_read,
+	.write		= stackmap_write,
+	.llseek		= seq_lseek,
+	.release	= stackmap_release,
+};
diff --git a/kernel/trace/trace_stackmap.h b/kernel/trace/trace_stackmap.h
new file mode 100644
index 000000000000..6699988da390
--- /dev/null
+++ b/kernel/trace/trace_stackmap.h
@@ -0,0 +1,22 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+#ifndef _TRACE_STACKMAP_H
+#define _TRACE_STACKMAP_H
+
+#include <linux/types.h>
+
+#define FTRACE_STACKMAP_MAX_DEPTH	64
+
+/* A cookie holds the generation in its high half and the slot below. */
+#define FTRACE_STACKMAP_COOKIE_GEN_SHIFT	32
+
+struct trace_array;
+struct ftrace_stackmap;
+
+int ftrace_stackmap_alloc(struct trace_array *tr);
+bool ftrace_stackmap_get_cookie(struct ftrace_stackmap *smap,
+				unsigned long *ips, unsigned int nr_entries,
+				u64 *cookie);
+
+extern const struct file_operations ftrace_stackmap_fops;
+
+#endif /* _TRACE_STACKMAP_H */
-- 
2.34.1