codetag: rcu-ify idx to codetag lookupmemalloc-prof-v7

This substantially reworks how we store and iterate over codetags in
multiple modules.

We now provide a stable index for each codetag, which doesn't change as
modules are loaded and unloaded - meaning on module load, we find a free
range of indices for the new module.

Index-to-codetag lookup uses an eytzinger tree, and we use RCU for the
eytzinger tree for lockless lookup. The eytzinger tree has one entry per
module, where entries store the starting index for each module's range.

After eytzinger lookup we use eytzinger_to_inorder(), as the
codetag_modules are still a flat array.

With the new stable indices, iteration becomes simpler; the iterator
just stores the current index and peek() finds the first codetag >= the
current index.

Signed-off-by: Kent Overstreet <kent.overstreet@linux.dev>

3 files changed, 246 insertions, 162 deletions

diff --git a/include/linux/codetag.h b/include/linux/codetag.h
index c2a579ccd455..6aa8cf22d88b 100644
--- a/include/linux/codetag.h
+++ b/include/linux/codetag.h
@@ -7,7 +7,7 @@
 
 #include <linux/types.h>
 
-struct codetag_iterator;
+struct codetag_iter;
 struct codetag_type;
 struct codetag_module;
 struct seq_buf;
@@ -39,11 +39,10 @@ struct codetag_type_desc {
 			      struct codetag_module *cmod);
 };
 
-struct codetag_iterator {
+struct codetag_iter {
 	struct codetag_type *cttype;
 	struct codetag_module *cmod;
-	unsigned long mod_id;
-	struct codetag *ct;
+	unsigned idx;
 };
 
 #ifdef MODULE
@@ -60,10 +59,27 @@ struct codetag_iterator {
 	.flags		= 0,				\
 }
 
+struct codetag *idx_to_codetag(struct codetag_type *cttype, unsigned idx);
+
 void codetag_lock_module_list(struct codetag_type *cttype, bool lock);
 bool codetag_trylock_module_list(struct codetag_type *cttype);
-struct codetag_iterator codetag_get_ct_iter(struct codetag_type *cttype);
-struct codetag *codetag_next_ct(struct codetag_iterator *iter);
+
+static inline struct codetag_iter codetag_iter_init(struct codetag_type *cttype, unsigned idx)
+{
+	return (struct codetag_iter) { .cttype = cttype, .idx = idx };
+}
+
+static inline void codetag_iter_advance(struct codetag_iter *iter)
+{
+	iter->idx++;
+}
+
+struct codetag *codetag_iter_peek(struct codetag_iter *);
+
+#define for_each_codetag(_cttype, _iter, _ct)				\
+	for (struct codetag_iter _iter = codetag_iter_init(_cttype, 0);	\
+	     (_ct = codetag_iter_peek(&_iter));				\
+	     codetag_iter_advance(&_iter))
 
 void codetag_to_text(struct seq_buf *out, struct codetag *ct);
 
diff --git a/lib/alloc_tag.c b/lib/alloc_tag.c
index 531dbe2f5456..66259854702a 100644
--- a/lib/alloc_tag.c
+++ b/lib/alloc_tag.c
@@ -18,42 +18,36 @@ DEFINE_STATIC_KEY_MAYBE(CONFIG_MEM_ALLOC_PROFILING_ENABLED_BY_DEFAULT,
 
 static void *allocinfo_start(struct seq_file *m, loff_t *pos)
 {
-	struct codetag_iterator *iter;
-	struct codetag *ct;
-	loff_t node = *pos;
-
-	iter = kzalloc(sizeof(*iter), GFP_KERNEL);
-	m->private = iter;
+	struct codetag_iter *iter = kzalloc(sizeof(*iter), GFP_KERNEL);
 	if (!iter)
 		return NULL;
 
+	m->private = iter;
+
 	codetag_lock_module_list(alloc_tag_cttype, true);
-	*iter = codetag_get_ct_iter(alloc_tag_cttype);
-	while ((ct = codetag_next_ct(iter)) != NULL && node)
-		node--;
+	*iter = codetag_iter_init(alloc_tag_cttype, *pos);
 
-	return ct ? iter : NULL;
+	return codetag_iter_peek(iter) ? iter : NULL;
 }
 
 static void *allocinfo_next(struct seq_file *m, void *arg, loff_t *pos)
 {
-	struct codetag_iterator *iter = (struct codetag_iterator *)arg;
-	struct codetag *ct = codetag_next_ct(iter);
+	struct codetag_iter *iter = arg;
 
-	(*pos)++;
-	if (!ct)
-		return NULL;
+	codetag_iter_advance(iter);
+	*pos = iter->idx;
 
-	return iter;
+	return codetag_iter_peek(iter) ? iter : NULL;
 }
 
 static void allocinfo_stop(struct seq_file *m, void *arg)
 {
-	struct codetag_iterator *iter = (struct codetag_iterator *)m->private;
+	struct codetag_iter *iter = m->private;
 
 	if (iter) {
 		codetag_lock_module_list(alloc_tag_cttype, false);
 		kfree(iter);
+		m->private = NULL;
 	}
 }
 
@@ -71,13 +65,13 @@ static void alloc_tag_to_text(struct seq_buf *out, struct codetag *ct)
 
 static int allocinfo_show(struct seq_file *m, void *arg)
 {
-	struct codetag_iterator *iter = (struct codetag_iterator *)arg;
+	struct codetag_iter *iter = (struct codetag_iter *)arg;
 	char *bufp;
 	size_t n = seq_get_buf(m, &bufp);
 	struct seq_buf buf;
 
 	seq_buf_init(&buf, bufp, n);
-	alloc_tag_to_text(&buf, iter->ct);
+	alloc_tag_to_text(&buf, codetag_iter_peek(iter));
 	seq_commit(m, seq_buf_used(&buf));
 	return 0;
 }
@@ -91,9 +85,7 @@ static const struct seq_operations allocinfo_seq_op = {
 
 size_t alloc_tag_top_users(struct codetag_bytes *tags, size_t count, bool can_sleep)
 {
-	struct codetag_iterator iter;
 	struct codetag *ct;
-	struct codetag_bytes n;
 	unsigned int i, nr = 0;
 
 	if (can_sleep)
@@ -101,12 +93,9 @@ size_t alloc_tag_top_users(struct codetag_bytes *tags, size_t count, bool can_sl
 	else if (!codetag_trylock_module_list(alloc_tag_cttype))
 		return 0;
 
-	iter = codetag_get_ct_iter(alloc_tag_cttype);
-	while ((ct = codetag_next_ct(&iter))) {
+	for_each_codetag(alloc_tag_cttype, iter, ct) {
 		struct alloc_tag_counters counter = alloc_tag_read(ct_to_alloc_tag(ct));
-
-		n.ct	= ct;
-		n.bytes = counter.bytes;
+		struct codetag_bytes n = { .ct = ct, .bytes = counter.bytes };
 
 		for (i = 0; i < nr; i++)
 			if (n.bytes > tags[i].bytes)
@@ -121,7 +110,6 @@ size_t alloc_tag_top_users(struct codetag_bytes *tags, size_t count, bool can_sl
 			tags[i] = n;
 		}
 	}
-
 	codetag_lock_module_list(alloc_tag_cttype, false);
 
 	return nr;
@@ -135,18 +123,15 @@ static void __init procfs_init(void)
 static bool alloc_tag_module_unload(struct codetag_type *cttype,
 				    struct codetag_module *cmod)
 {
-	struct codetag_iterator iter = codetag_get_ct_iter(cttype);
-	struct alloc_tag_counters counter;
 	bool module_unused = true;
-	struct alloc_tag *tag;
 	struct codetag *ct;
 
-	for (ct = codetag_next_ct(&iter); ct; ct = codetag_next_ct(&iter)) {
+	for_each_codetag(cttype, iter, ct) {
 		if (iter.cmod != cmod)
 			continue;
 
-		tag = ct_to_alloc_tag(ct);
-		counter = alloc_tag_read(tag);
+		struct alloc_tag *tag = ct_to_alloc_tag(ct);
+		struct alloc_tag_counters counter = alloc_tag_read(tag);
 
 		if (WARN(counter.bytes,
 			 "%s:%u module %s func:%s has %llu allocated at module unload",
diff --git a/lib/codetag.c b/lib/codetag.c
index 5ace625f2328..8819e5469b04 100644
--- a/lib/codetag.c
+++ b/lib/codetag.c
@@ -1,28 +1,68 @@
 // SPDX-License-Identifier: GPL-2.0-only
 #include <linux/codetag.h>
-#include <linux/idr.h>
+#include <linux/darray.h>
+#include <linux/eytzinger.h>
 #include <linux/kallsyms.h>
 #include <linux/module.h>
 #include <linux/seq_buf.h>
 #include <linux/slab.h>
 #include <linux/vmalloc.h>
 
-struct codetag_type {
-	struct list_head link;
-	unsigned int count;
-	struct idr mod_idr;
-	struct rw_semaphore mod_lock; /* protects mod_idr */
-	struct codetag_type_desc desc;
-};
+#ifdef DEBUG
+#define EBUG_ON(cond)	BUG_ON(cond)
+#else
+#define EBUG_ON(cond)	do {} while (0)
+#endif
 
-struct codetag_range {
-	struct codetag *start;
-	struct codetag *stop;
+struct codetag_eytz_entry {
+	u16				idx;
 };
 
 struct codetag_module {
-	struct module *mod;
-	struct codetag_range range;
+	unsigned			idx;
+	unsigned			nr;
+	struct codetag			*start;
+	struct module			*mod;
+};
+
+struct codetag_modules {
+	struct rcu_head			rcu;
+	u16				nr_modules;
+	u16				eytz_extra;
+	struct codetag_eytz_entry	e[];
+	/*
+	 * Additionally, we have an array of @nr_modules codetag_modules after
+	 * e[nr_modules]
+	 */
+};
+
+static inline size_t codetag_modules_bytes(unsigned nr_modules)
+{
+	return sizeof(struct codetag_modules) +
+		sizeof(struct codetag_eytz_entry) * nr_modules +
+		sizeof(struct codetag_module) * nr_modules;
+}
+
+static inline struct codetag_module *codetag_mods_array(struct codetag_modules *m)
+{
+	return (void *) &m->e[m->nr_modules];
+}
+
+static void codetag_modules_init_eytz(struct codetag_modules *mods)
+{
+	mods->eytz_extra = eytzinger0_extra(mods->nr_modules);
+
+	for (unsigned i = 0; i < mods->nr_modules; i++)
+		mods->e[__inorder_to_eytzinger0(i, mods->nr_modules, mods->eytz_extra)].idx =
+			codetag_mods_array(mods)[i].idx;
+}
+
+struct codetag_type {
+	struct list_head		link;
+	unsigned int			count;
+	struct rw_semaphore		mod_lock; /* protects mods */
+	struct codetag_modules __rcu	*mods_rcu;
+	struct codetag_type_desc	desc;
 };
 
 static DEFINE_MUTEX(codetag_lock);
@@ -41,88 +81,91 @@ bool codetag_trylock_module_list(struct codetag_type *cttype)
 	return down_read_trylock(&cttype->mod_lock) != 0;
 }
 
-struct codetag_iterator codetag_get_ct_iter(struct codetag_type *cttype)
+#define cmp_int(l, r)		((l > r) - (l < r))
+
+static inline int codetag_eytz_entry_cmp(const void *_l, const void *_r)
 {
-	struct codetag_iterator iter = {
-		.cttype = cttype,
-		.cmod = NULL,
-		.mod_id = 0,
-		.ct = NULL,
-	};
+	const struct codetag_eytz_entry *l = _l;
+	const struct codetag_eytz_entry *r = _r;
 
-	return iter;
+	return cmp_int(l->idx, r->idx);
 }
 
-static inline struct codetag *get_first_module_ct(struct codetag_module *cmod)
+__always_inline
+static inline struct codetag_module *codetag_idx_to_cmod(struct codetag_modules *mods, unsigned idx)
 {
-	return cmod->range.start < cmod->range.stop ? cmod->range.start : NULL;
+	struct codetag_eytz_entry search = { .idx = idx };
+	unsigned e = eytzinger0_find_le(mods->e,
+					mods->nr_modules,
+					sizeof(mods->e[0]),
+					codetag_eytz_entry_cmp,
+					&search);
+	if (e >= mods->nr_modules)
+		return NULL;
+
+	struct codetag_module *mod = codetag_mods_array(mods) +
+		__eytzinger0_to_inorder(e, mods->nr_modules, mods->eytz_extra);
+
+	EBUG_ON(mods->e[e].idx != mod->idx);
+	return mod;
 }
 
-static inline
-struct codetag *get_next_module_ct(struct codetag_iterator *iter)
+static struct codetag *__idx_to_codetag(struct codetag_type *cttype,
+					struct codetag_module *cmod,
+					unsigned idx)
 {
-	struct codetag *res = (struct codetag *)
-			((char *)iter->ct + iter->cttype->desc.tag_size);
+	EBUG_ON(idx < cmod->idx);
+	EBUG_ON(idx >= cmod->idx + cmod->nr);
 
-	return res < iter->cmod->range.stop ? res : NULL;
+	return (void *) cmod->start + (idx - cmod->idx) * cttype->desc.tag_size;
 }
 
-struct codetag *codetag_next_ct(struct codetag_iterator *iter)
+/* @idx must point to a valid codetag, not a gap */
+struct codetag *idx_to_codetag(struct codetag_type *cttype, unsigned idx)
 {
-	struct codetag_type *cttype = iter->cttype;
-	struct codetag_module *cmod;
-	struct codetag *ct;
+	rcu_read_lock();
+	struct codetag_modules *mods = rcu_dereference(cttype->mods_rcu);
+	struct codetag *ct = __idx_to_codetag(cttype, codetag_idx_to_cmod(mods, idx), idx);
+	rcu_read_unlock();
+	return ct;
+}
 
-	lockdep_assert_held(&cttype->mod_lock);
+/* finds the first valid codetag at idx >= iter->idx, or returns NULL */
+static struct codetag *__codetag_iter_peek(struct codetag_iter *iter)
+{
+	struct codetag_type *cttype = iter->cttype;
+	struct codetag_modules *mods = rcu_dereference(cttype->mods_rcu);
 
-	if (unlikely(idr_is_empty(&cttype->mod_idr)))
+	iter->cmod = codetag_idx_to_cmod(mods, iter->idx);
+	if (!iter->cmod)
 		return NULL;
 
-	ct = NULL;
-	while (true) {
-		cmod = idr_find(&cttype->mod_idr, iter->mod_id);
-
-		/* If module was removed move to the next one */
-		if (!cmod)
-			cmod = idr_get_next_ul(&cttype->mod_idr,
-					       &iter->mod_id);
-
-		/* Exit if no more modules */
-		if (!cmod)
-			break;
-
-		if (cmod != iter->cmod) {
-			iter->cmod = cmod;
-			ct = get_first_module_ct(cmod);
-		} else
-			ct = get_next_module_ct(iter);
+	if (iter->cmod->idx + iter->cmod->nr <= iter->idx) {
+		iter->cmod++;
+		if (iter->cmod == codetag_mods_array(mods) + mods->nr_modules)
+			return NULL;
+	}
 
-		if (ct)
-			break;
+	iter->idx = max(iter->idx, iter->cmod->idx);
 
-		iter->mod_id++;
-	}
+	return __idx_to_codetag(cttype, iter->cmod, iter->idx);
+}
 
-	iter->ct = ct;
+struct codetag *codetag_iter_peek(struct codetag_iter *iter)
+{
+	rcu_read_lock();
+	struct codetag *ct = __codetag_iter_peek(iter);
+	rcu_read_unlock();
 	return ct;
+
 }
 
 void codetag_to_text(struct seq_buf *out, struct codetag *ct)
 {
+	seq_buf_printf(out, "%s:%u", ct->filename, ct->lineno);
 	if (ct->modname)
-		seq_buf_printf(out, "%s:%u [%s] func:%s",
-			       ct->filename, ct->lineno,
-			       ct->modname, ct->function);
-	else
-		seq_buf_printf(out, "%s:%u func:%s",
-			       ct->filename, ct->lineno, ct->function);
-}
-
-static inline size_t range_size(const struct codetag_type *cttype,
-				const struct codetag_range *range)
-{
-	return ((char *)range->stop - (char *)range->start) /
-			cttype->desc.tag_size;
+		seq_buf_printf(out, " [%s]", ct->modname);
+	seq_buf_printf(out, " func:%s", ct->function);
 }
 
 #ifdef CONFIG_MODULES
@@ -146,59 +189,73 @@ static void *get_symbol(struct module *mod, const char *prefix, const char *name
 	return ret;
 }
 
-static struct codetag_range get_section_range(struct module *mod,
-					      const char *section)
+static int __codetag_module_init(struct codetag_type *cttype, struct module *mod)
 {
-	return (struct codetag_range) {
-		get_symbol(mod, "__start_", section),
-		get_symbol(mod, "__stop_", section),
-	};
-}
+	struct codetag *start	= get_symbol(mod, "__start_", cttype->desc.section);
+	struct codetag *stop	= get_symbol(mod, "__stop_", cttype->desc.section);
 
-static int codetag_module_init(struct codetag_type *cttype, struct module *mod)
-{
-	struct codetag_range range;
-	struct codetag_module *cmod;
-	int err;
+	BUG_ON(start > stop);
 
-	range = get_section_range(mod, cttype->desc.section);
-	if (!range.start || !range.stop) {
+	if (!start || !stop) {
 		pr_warn("Failed to load code tags of type %s from the module %s\n",
-			cttype->desc.section,
-			mod ? mod->name : "(built-in)");
+			cttype->desc.section, mod ? mod->name : "(built-in)");
 		return -EINVAL;
 	}
 
+	struct codetag_module cmod = {
+		.nr	= ((void *) stop - (void *) start) / cttype->desc.tag_size,
+		.start	= start,
+		.mod	= mod,
+	};
+
 	/* Ignore empty ranges */
-	if (range.start == range.stop)
+	if (!cmod.nr)
 		return 0;
 
-	BUG_ON(range.start > range.stop);
+	struct codetag_modules *old_mods =
+		rcu_dereference_protected(cttype->mods_rcu, lockdep_is_held(&cttype->mod_lock));
+	unsigned old_nr_modules = old_mods ? old_mods->nr_modules : 0;
+	unsigned new_nr_modules = old_nr_modules + 1;
+	struct codetag_modules *new_mods = kzalloc(codetag_modules_bytes(new_nr_modules), GFP_KERNEL);
 
-	cmod = kmalloc(sizeof(*cmod), GFP_KERNEL);
-	if (unlikely(!cmod))
+	if (!new_mods)
 		return -ENOMEM;
 
-	cmod->mod = mod;
-	cmod->range = range;
+	new_mods->nr_modules = new_nr_modules;
+	struct codetag_module *mod_a = codetag_mods_array(new_mods);
 
-	down_write(&cttype->mod_lock);
-	err = idr_alloc(&cttype->mod_idr, cmod, 0, 0, GFP_KERNEL);
-	if (err >= 0) {
-		cttype->count += range_size(cttype, &range);
-		if (cttype->desc.module_load)
-			cttype->desc.module_load(cttype, cmod);
-	}
-	up_write(&cttype->mod_lock);
+	if (old_mods)
+		memcpy(mod_a, codetag_mods_array(old_mods),
+		       sizeof(struct codetag_module) * old_mods->nr_modules);
 
-	if (err < 0) {
-		kfree(cmod);
-		return err;
+
+	for (unsigned i = 0; i < old_nr_modules; i++) {
+		if (cmod.idx + cmod.nr <= mod_a[i].idx) {
+			array_insert_item(mod_a, old_nr_modules, i, cmod);
+			goto insert_done;
+		}
+
+		cmod.idx = mod_a[i].idx + mod_a[i].nr;
 	}
 
+	mod_a[old_nr_modules] = cmod;
+insert_done:
+	codetag_modules_init_eytz(new_mods);
+
+	rcu_assign_pointer(cttype->mods_rcu, new_mods);
+	kfree_rcu(old_mods, rcu);
 	return 0;
 }
 
+static int codetag_module_init(struct codetag_type *cttype, struct module *mod)
+{
+	down_write(&cttype->mod_lock);
+	int ret = __codetag_module_init(cttype, mod);
+	up_write(&cttype->mod_lock);
+
+	return ret;
+}
+
 void codetag_load_module(struct module *mod)
 {
 	struct codetag_type *cttype;
@@ -212,6 +269,50 @@ void codetag_load_module(struct module *mod)
 	mutex_unlock(&codetag_lock);
 }
 
+static bool cttype_unload_module(struct codetag_type *cttype, struct module *mod)
+{
+	bool unload_ok = true;
+
+	struct codetag_modules *new_mods = NULL;
+	struct codetag_modules *old_mods =
+		rcu_dereference_protected(cttype->mods_rcu, lockdep_is_held(&cttype->mod_lock));
+	struct codetag_module *mod_a = codetag_mods_array(old_mods);
+
+	unsigned pos;
+	for (pos = 0; pos < old_mods->nr_modules; pos++)
+		if (mod_a[pos].mod == mod)
+			goto found;
+	return true;
+found:
+	if (cttype->desc.module_unload &&
+	    !cttype->desc.module_unload(cttype, &mod_a[pos]))
+		unload_ok = false;
+	cttype->count -= mod_a[pos].nr;
+
+	unsigned new_nr = old_mods->nr_modules - 1;
+	if (!new_nr)
+		goto out;
+
+	new_mods = kzalloc(codetag_modules_bytes(old_mods->nr_modules), GFP_KERNEL);
+	if (!new_mods)
+		return false;
+
+	new_mods->nr_modules = new_nr;
+	mod_a = codetag_mods_array(new_mods);
+
+	memcpy(mod_a, codetag_mods_array(old_mods),
+	       sizeof(struct codetag_module) * old_mods->nr_modules);
+	memmove(&mod_a[pos],
+		&mod_a[pos + 1],
+		sizeof(mod_a[0]) * new_nr - pos);
+
+	codetag_modules_init_eytz(new_mods);
+out:
+	rcu_assign_pointer(cttype->mods_rcu, new_mods);
+	kfree_rcu(old_mods, rcu);
+	return unload_ok;
+}
+
 bool codetag_unload_module(struct module *mod)
 {
 	struct codetag_type *cttype;
@@ -222,26 +323,9 @@ bool codetag_unload_module(struct module *mod)
 
 	mutex_lock(&codetag_lock);
 	list_for_each_entry(cttype, &codetag_types, link) {
-		struct codetag_module *found = NULL;
-		struct codetag_module *cmod;
-		unsigned long mod_id, tmp;
-
 		down_write(&cttype->mod_lock);
-		idr_for_each_entry_ul(&cttype->mod_idr, cmod, tmp, mod_id) {
-			if (cmod->mod && cmod->mod == mod) {
-				found = cmod;
-				break;
-			}
-		}
-		if (found) {
-			if (cttype->desc.module_unload)
-				if (!cttype->desc.module_unload(cttype, cmod))
-					unload_ok = false;
-
-			cttype->count -= range_size(cttype, &cmod->range);
-			idr_remove(&cttype->mod_idr, mod_id);
-			kfree(cmod);
-		}
+		if (!cttype_unload_module(cttype, mod))
+			unload_ok = false;
 		up_write(&cttype->mod_lock);
 	}
 	mutex_unlock(&codetag_lock);
@@ -266,7 +350,6 @@ codetag_register_type(const struct codetag_type_desc *desc)
 		return ERR_PTR(-ENOMEM);
 
 	cttype->desc = *desc;
-	idr_init(&cttype->mod_idr);
 	init_rwsem(&cttype->mod_lock);
 
 	err = codetag_module_init(cttype, NULL);