[PATCH bpf-next v3 0/3] selftests/bpf: libarena: Add ordered-walk and lockless-peek support to the rbtree

[PATCH bpf-next v3 0/3] selftests/bpf: libarena: Add ordered-walk and lockless-peek support to the rbtree

[Changwoo Min]

A sched_ext scheduler can build a task queue on libarena's rbtree, keyed
by vtime, in place of a vtime-ordered DSQ. Such a queue needs three things
the tree lacks. This series adds them:

- Serve equal keys in insertion order, as vtime-ordered DSQs do (1/3).
- Add rb_next(), so that a consumer can walk past tasks it cannot run
  in key order (2/3).
- Cache the leftmost node and the node count, so that other CPUs can
  peek at the head and the length without the writer's lock, like
  scx_bpf_dsq_peek() and scx_bpf_dsq_nr_queued(). Add rb_init() and
  rb_fini() for a tree embedded in another object (3/3).

Each patch adds selftests, which run with:

  $> ./test_progs -t 'libarena/test_rbtree*'

Changes in v3:
- Make the lockless readers static inline. A plain static function in
  a header fails -Wunused-function under -Werror in a file that does
  not call it (bot+bpf-ci, 3/3).

Changes in v2:
- Drop the comment on the insert helper and keep the name
  rbnode_least_upper_bound() (Emil, 1/3).
- Move rb_next() before the caching patch and into rbtree.bpf.c with the
  other traversal functions (Alexei, Emil, 2/3).
- Hand the head to its successor at the top of rb_node_remove(), before
  an RB_ALLOC tree frees the node, so that a lockless reader never sees a
  freed head. This drops the wrapper, was_first and the rename to
  rbnode_remove() (Alexei, 3/3).
- Drop the cache line alignment of the lockless fields, which showed no
  measurable benefit (Alexei, 3/3).
- Rename the lockless readers to rb_first_lockless(),
  rb_first_key_lockless() and rb_nr_nodes_lockless(), and document that
  a caller that dereferences the head must keep the node valid itself
  (Emil, 3/3).
- Drop __always_inline from static functions (Emil, 3/3).
- Add Reviewed-by tags from Alexei and Emil.
- Rebase on bpf-next.

v2: https://lore.kernel.org/bpf/20261111073416.203153-1-changwoo@igalia.com/
v1: https://lore.kernel.org/bpf/20261028060046.549097-1-changwoo@igalia.com/

Changwoo Min (3):
  selftests/bpf: libarena: Order equal rbtree keys FIFO
  selftests/bpf: libarena: Add rb_next()
  selftests/bpf: libarena: Cache the rbtree's leftmost node and count

 .../bpf/libarena/include/libarena/rbtree.h    |  57 +++-
 .../bpf/libarena/selftests/test_rbtree.bpf.c  | 303 ++++++++++++++++++
 .../selftests/bpf/libarena/src/rbtree.bpf.c   | 120 +++++--
 3 files changed, 457 insertions(+), 23 deletions(-)


base-commit: 0fb18ce8a2e9de830433923cf65838933d615359
-- 
2.56.0

[PATCH bpf-next v3 1/3] selftests/bpf: libarena: Order equal rbtree keys FIFO

[Changwoo Min]

An RB_DUPLICATE tree links a new key to the left of an equal one, so
rb_least() and rb_pop() return the newest of several equal keys. A
queue of tasks keyed by vtime should instead serve equal keys in
arrival order, as sched_ext's vtime-ordered DSQs do.

Send equal keys right on insert, as the kernel's rb_add() does. A new
node then lands after every node with an equal key. Rotations and
removals preserve the in-order sequence, so equal keys stay in
insertion order without a tie-break field.

Add a selftest that interleaves equal keys while the tree grows and
checks the pop order.

Signed-off-by: Changwoo Min <changwoo@igalia.com>
Reviewed-by: Alexei Starovoitov <ast@kernel.org>
Reviewed-by: Emil Tsalapatis <emil@etsalapatis.com>
---
 .../bpf/libarena/include/libarena/rbtree.h    |  6 +--
 .../bpf/libarena/selftests/test_rbtree.bpf.c  | 37 +++++++++++++++++++
 .../selftests/bpf/libarena/src/rbtree.bpf.c   |  6 +--
 3 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h b/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h
index 486428911d96..00e63b12d82e 100644
--- a/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h
+++ b/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h
@@ -45,9 +45,9 @@ enum rbtree_alloc {
  * the existing node.
  *
  * RB_DUPLICATE: Allow nodes with identical keys in the rbtree.
- * Finding/popping/removing a key acts on any of the nodes
- * with the appropriate key - there is no ordering by time
- * of insertion.
+ * Nodes with identical keys are kept in insertion order, so
+ * rb_least() and rb_pop() return the oldest of them. Finding
+ * or removing by key acts on any of them.
  */
 enum rbtree_insert_mode {
 	RB_DEFAULT,
diff --git a/tools/testing/selftests/bpf/libarena/selftests/test_rbtree.bpf.c b/tools/testing/selftests/bpf/libarena/selftests/test_rbtree.bpf.c
index 856c484a009a..1e2a0fac1e4f 100644
--- a/tools/testing/selftests/bpf/libarena/selftests/test_rbtree.bpf.c
+++ b/tools/testing/selftests/bpf/libarena/selftests/test_rbtree.bpf.c
@@ -248,6 +248,43 @@ __weak int test_rbtree_duplicate(void)
 	return rb_destroy(rbtree);
 }
 
+#define FIFO_NR_NODES	64
+#define FIFO_NR_KEYS	4
+
+SEC("syscall")
+__weak int test_rbtree_duplicate_fifo(void)
+{
+	const u64 per_key = FIFO_NR_NODES / FIFO_NR_KEYS;
+	struct rbtree __arena *rbtree;
+	u64 key, value;
+	int ret, i;
+
+	rbtree = rb_create(RB_ALLOC, RB_DUPLICATE);
+	if (!rbtree)
+		return 1;
+
+	for (i = 0; i < FIFO_NR_NODES && can_loop; i++) {
+		ret = rb_insert(rbtree, i % FIFO_NR_KEYS, i);
+		if (ret)
+			return 2;
+	}
+
+	for (i = 0; i < FIFO_NR_NODES && can_loop; i++) {
+		ret = rb_pop(rbtree, &key, &value);
+		if (ret)
+			return 3;
+
+		/* All of key 0 in insertion order, then all of key 1, ... */
+		if (key != i / per_key)
+			return 4;
+
+		if (value != key + FIFO_NR_KEYS * (i % per_key))
+			return 5;
+	}
+
+	return rb_destroy(rbtree);
+}
+
 static inline int
 clean_up_noalloc_tree(struct rbtree __arena *rbtree)
 {
diff --git a/tools/testing/selftests/bpf/libarena/src/rbtree.bpf.c b/tools/testing/selftests/bpf/libarena/src/rbtree.bpf.c
index 7f0f6dc3e17d..831a1ef1c2a1 100644
--- a/tools/testing/selftests/bpf/libarena/src/rbtree.bpf.c
+++ b/tools/testing/selftests/bpf/libarena/src/rbtree.bpf.c
@@ -171,7 +171,7 @@ struct rbnode __arena *rbnode_least_upper_bound(struct rbnode __arena *subtree,
 		return NULL;
 
 	while (can_loop) {
-		dir = (key <= node->key) ? 0 : 1;
+		dir = (key < node->key) ? 0 : 1;
 
 		if (!node->child[dir])
 			break;
@@ -278,8 +278,8 @@ int rb_node_insert(struct rbtree __arena *rbtree,
 	}
 
 	node->parent = parent;
-	/* Also works if key == parent->key. */
-	if (key <= parent->key)
+	/* An equal key goes right, after the existing ones. */
+	if (key < parent->key)
 		parent->left = node;
 	else
 		parent->right = node;
-- 
2.56.0

[PATCH bpf-next v3 2/3] selftests/bpf: libarena: Add rb_next()

[Changwoo Min]

Add rb_next(), the in-order successor through parent pointers.

Add a selftest that walks a tree with repeated keys and checks the walk
against the order in which the head leaves the tree.

Signed-off-by: Changwoo Min <changwoo@igalia.com>
Reviewed-by: Alexei Starovoitov <ast@kernel.org>
Reviewed-by: Emil Tsalapatis <emil@etsalapatis.com>
---
 .../bpf/libarena/include/libarena/rbtree.h    |  1 +
 .../bpf/libarena/selftests/test_rbtree.bpf.c  | 57 +++++++++++++++++++
 .../selftests/bpf/libarena/src/rbtree.bpf.c   | 20 +++++++
 3 files changed, 78 insertions(+)

diff --git a/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h b/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h
index 00e63b12d82e..788e25fcd7fd 100644
--- a/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h
+++ b/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h
@@ -71,6 +71,7 @@ int rb_find(struct rbtree __arena *rbtree, u64 key, u64 *value);
 int rb_print(struct rbtree __arena *rbtree);
 int rb_least(struct rbtree __arena *rbtree, u64 *key, u64 *value);
 int rb_pop(struct rbtree __arena *rbtree, u64 *key, u64 *value);
+struct rbnode __arena *rb_next(struct rbnode __arena *node);
 
 int rb_insert_node(struct rbtree __arena *rbtree, struct rbnode __arena *node);
 int rb_remove_node(struct rbtree __arena *rbtree, struct rbnode __arena *node);
diff --git a/tools/testing/selftests/bpf/libarena/selftests/test_rbtree.bpf.c b/tools/testing/selftests/bpf/libarena/selftests/test_rbtree.bpf.c
index 1e2a0fac1e4f..a92b00f799f8 100644
--- a/tools/testing/selftests/bpf/libarena/selftests/test_rbtree.bpf.c
+++ b/tools/testing/selftests/bpf/libarena/selftests/test_rbtree.bpf.c
@@ -285,6 +285,63 @@ __weak int test_rbtree_duplicate_fifo(void)
 	return rb_destroy(rbtree);
 }
 
+#define NEXT_NR_NODES	64
+
+static struct rbnode __arena *next_nodes[NEXT_NR_NODES];
+
+SEC("syscall")
+__weak int test_rbtree_next(void)
+{
+	struct rbtree __arena *rbtree;
+	struct rbnode __arena *node;
+	u64 seed = 0xcafe, value;
+	int i;
+
+	rbtree = rb_create(RB_NOALLOC, RB_DUPLICATE);
+	if (!rbtree)
+		return 1;
+
+	/* Pseudo-random keys with repeats, so the walk crosses equal keys. */
+	for (i = 0; i < NEXT_NR_NODES && can_loop; i++) {
+		/* LCG using Knuth's MMIX constants for determinism. */
+		seed = seed * 6364136223846793005ULL + 1442695040888963407ULL;
+		node = rb_node_alloc((seed >> 33) % 16, i);
+		if (!node)
+			return 2;
+
+		if (rb_insert_node(rbtree, node))
+			return 3;
+	}
+
+	/* Start the walk at the leftmost node. */
+	node = rbtree->root;
+	while (node->left && can_loop)
+		node = node->left;
+
+	for (i = 0; i < NEXT_NR_NODES && can_loop; i++) {
+		if (!node)
+			return 4;
+
+		next_nodes[i] = node;
+		node = rb_next(node);
+	}
+
+	if (node)
+		return 5;
+
+	for (i = 0; i < NEXT_NR_NODES && can_loop; i++) {
+		if (rb_least(rbtree, NULL, &value) || value != next_nodes[i]->value)
+			return 6;
+
+		if (rb_remove_node(rbtree, next_nodes[i]))
+			return 7;
+
+		rb_node_free(next_nodes[i]);
+	}
+
+	return rb_destroy(rbtree);
+}
+
 static inline int
 clean_up_noalloc_tree(struct rbtree __arena *rbtree)
 {
diff --git a/tools/testing/selftests/bpf/libarena/src/rbtree.bpf.c b/tools/testing/selftests/bpf/libarena/src/rbtree.bpf.c
index 831a1ef1c2a1..bf39cb2aa751 100644
--- a/tools/testing/selftests/bpf/libarena/src/rbtree.bpf.c
+++ b/tools/testing/selftests/bpf/libarena/src/rbtree.bpf.c
@@ -413,6 +413,26 @@ __weak int rb_least(struct rbtree __arena *rbtree, u64 *key, u64 *value)
 	return 0;
 }
 
+/*
+ * Return the in-order successor of @node, or NULL if @node is the last one.
+ * Walking from the leftmost node with it visits the tree in key order, equal
+ * keys in insertion order. Call it under the writer's lock.
+ */
+__weak
+struct rbnode __arena *rb_next(struct rbnode __arena *node)
+{
+	struct rbnode __arena *parent;
+
+	if (node->right)
+		return rbnode_least(node->right);
+
+	/* Climb until we leave a left subtree; its parent comes next. */
+	while ((parent = node->parent) && node == parent->right && can_loop)
+		node = parent;
+
+	return parent;
+}
+
 
 /*
  * If we are referencing ourselves, a and b have a parent-child relation,
-- 
2.56.0

[PATCH bpf-next v3 3/3] selftests/bpf: libarena: Cache the rbtree's leftmost node and count

[Changwoo Min]

A run queue built on the tree needs lockless peeks at its head and
length, like scx_bpf_dsq_peek() and scx_bpf_dsq_nr_queued() for
sched_ext DSQs.

Keep the leftmost node and the node count in struct rbtree. An insert
changes the head only when it links left of it, and removing the head
hands it to its in-order successor before the tree changes. rb_least()
and rb_pop() use the cached head. rb_first_lockless(),
rb_first_key_lockless() and rb_nr_nodes_lockless() read them without
the lock. The head pointer, published with release ordering, decides
emptiness, and the key is read through it.

Add rb_init() for a tree embedded in another object, such as a queue
that keeps its lock next to the tree, and rb_fini() to empty it, since
rb_destroy() also frees the tree's memory.

Add selftests for the head and count across inserts, removals and
updates, and for an embedded tree.

Signed-off-by: Changwoo Min <changwoo@igalia.com>
Reviewed-by: Alexei Starovoitov <ast@kernel.org>
Reviewed-by: Emil Tsalapatis <emil@etsalapatis.com>
---
 .../bpf/libarena/include/libarena/rbtree.h    |  50 +++++
 .../bpf/libarena/selftests/test_rbtree.bpf.c  | 209 ++++++++++++++++++
 .../selftests/bpf/libarena/src/rbtree.bpf.c   |  94 ++++++--
 3 files changed, 336 insertions(+), 17 deletions(-)

diff --git a/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h b/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h
index 788e25fcd7fd..7b8afa38474b 100644
--- a/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h
+++ b/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h
@@ -59,11 +59,28 @@ struct rbtree {
 	struct rbnode __arena *root;
 	enum rbtree_alloc alloc;
 	enum rbtree_insert_mode insert;
+
+	/*
+	 * The leftmost node and the node count. The writer keeps them current,
+	 * and readers may read them without the writer's lock (see
+	 * rb_first_lockless() and friends).
+	 */
+	struct rbnode __arena *first;
+	uint64_t nr;
 };
 
 #ifdef __BPF__
 struct rbtree __arena *rb_create(enum rbtree_alloc alloc, enum rbtree_insert_mode insert);
 
+/*
+ * Initialize and finalize a tree embedded in another arena object.
+ * rb_fini() frees the nodes of an RB_ALLOC tree and fails with -EBUSY on a
+ * non-empty RB_NOALLOC tree. It leaves the tree's memory to its owner.
+ */
+int rb_init(struct rbtree __arena *rbtree, enum rbtree_alloc alloc,
+	    enum rbtree_insert_mode insert);
+int rb_fini(struct rbtree __arena *rbtree);
+
 int rb_destroy(struct rbtree __arena *rbtree);
 int rb_insert(struct rbtree __arena *rbtree, u64 key, u64 value);
 int rb_remove(struct rbtree __arena *rbtree, u64 key);
@@ -81,4 +98,37 @@ void rb_node_free(struct rbnode __arena *rbnode);
 
 int rb_integrity_check(struct rbtree __arena *rbtree);
 
+/*
+ * Lockless reads of the head and the node count. They need not be called
+ * under the writer's lock, so they are advisory: the tree may change before
+ * the caller acts on the result. The node rb_first_lockless() returns may
+ * have left the tree, or been freed in an RB_ALLOC tree, by the time the
+ * caller reads it. A caller that dereferences it must keep the node's memory
+ * valid itself, for example by embedding the node in an object that is freed
+ * only after an RCU grace period and reading it inside an RCU read-side
+ * critical section. The count is a separate read and need not agree with the
+ * head.
+ */
+static inline struct rbnode __arena *rb_first_lockless(struct rbtree __arena *rbtree)
+{
+	return READ_ONCE(rbtree->first);
+}
+
+/* -ENOENT if the tree is empty, 0 with the leftmost key in @key otherwise. */
+static inline int rb_first_key_lockless(struct rbtree __arena *rbtree, u64 *key)
+{
+	struct rbnode __arena *first = READ_ONCE(rbtree->first);
+
+	if (!first)
+		return -ENOENT;
+
+	*key = READ_ONCE(first->key);
+	return 0;
+}
+
+static inline u64 rb_nr_nodes_lockless(struct rbtree __arena *rbtree)
+{
+	return READ_ONCE(rbtree->nr);
+}
+
 #endif /* __BPF__ */
diff --git a/tools/testing/selftests/bpf/libarena/selftests/test_rbtree.bpf.c b/tools/testing/selftests/bpf/libarena/selftests/test_rbtree.bpf.c
index a92b00f799f8..4da6f293e286 100644
--- a/tools/testing/selftests/bpf/libarena/selftests/test_rbtree.bpf.c
+++ b/tools/testing/selftests/bpf/libarena/selftests/test_rbtree.bpf.c
@@ -285,6 +285,146 @@ __weak int test_rbtree_duplicate_fifo(void)
 	return rb_destroy(rbtree);
 }
 
+/*
+ * Check the lockless head and count against a walk to the leftmost node,
+ * which does not use the cached head.
+ */
+__weak int rbtree_check_first(struct rbtree __arena *rbtree, u64 nr)
+{
+	struct rbnode __arena *least = rbtree->root;
+	u64 key;
+
+	while (least && least->left && can_loop)
+		least = least->left;
+
+	if (rb_first_lockless(rbtree) != least)
+		return 1;
+
+	if (rb_nr_nodes_lockless(rbtree) != nr)
+		return 2;
+
+	if (!least)
+		return rb_first_key_lockless(rbtree, &key) == -ENOENT ? 0 : 3;
+
+	if (rb_first_key_lockless(rbtree, &key) || key != least->key)
+		return 4;
+
+	return 0;
+}
+
+#define FIRST_NR_NODES	64
+
+static struct rbnode __arena *first_nodes[FIRST_NR_NODES];
+
+SEC("syscall")
+__weak int test_rbtree_first(void)
+{
+	struct rbtree __arena *rbtree;
+	struct rbnode __arena *node;
+	u64 seed = 12345, nr = 0;
+	int ret, i;
+
+	rbtree = rb_create(RB_NOALLOC, RB_DUPLICATE);
+	if (!rbtree)
+		return 1;
+
+	ret = rbtree_check_first(rbtree, 0);
+	if (ret)
+		return 10 + ret;
+
+	/* Pseudo-random keys with repeats, so the head moves and gets ties. */
+	for (i = 0; i < FIRST_NR_NODES && can_loop; i++) {
+		seed = seed * 6364136223846793005ULL + 1442695040888963407ULL;
+		node = rb_node_alloc((seed >> 33) % 32, i);
+		if (!node)
+			return 2;
+
+		first_nodes[i] = node;
+
+		ret = rb_insert_node(rbtree, node);
+		if (ret)
+			return 3;
+
+		ret = rbtree_check_first(rbtree, ++nr);
+		if (ret)
+			return 20 + ret;
+	}
+
+	/* Remove every other node in insertion order, head or not. */
+	for (i = 0; i < FIRST_NR_NODES && can_loop; i += 2) {
+		ret = rb_remove_node(rbtree, first_nodes[i]);
+		if (ret)
+			return 4;
+
+		rb_node_free(first_nodes[i]);
+
+		ret = rbtree_check_first(rbtree, --nr);
+		if (ret)
+			return 30 + ret;
+	}
+
+	/* Remove the head of the remaining half, emptying the tree. */
+	for (i = 0; i < FIRST_NR_NODES / 2 && can_loop; i++) {
+		node = rb_first_lockless(rbtree);
+		if (!node)
+			return 5;
+
+		ret = rb_remove_node(rbtree, node);
+		if (ret)
+			return 6;
+
+		rb_node_free(node);
+
+		ret = rbtree_check_first(rbtree, --nr);
+		if (ret)
+			return 40 + ret;
+	}
+
+	return rb_destroy(rbtree);
+}
+
+SEC("syscall")
+__weak int test_rbtree_first_update(void)
+{
+	struct rbtree __arena *rbtree;
+	struct rbnode __arena *node;
+	int ret;
+
+	rbtree = rb_create(RB_ALLOC, RB_UPDATE);
+	if (!rbtree)
+		return 1;
+
+	if (rb_insert(rbtree, 5, 1))
+		return 2;
+
+	if (rb_insert(rbtree, 5, 2))
+		return 3;
+
+	node = rb_first_lockless(rbtree);
+	if (!node || node->key != 5 || node->value != 2)
+		return 4;
+
+	ret = rbtree_check_first(rbtree, 1);
+	if (ret)
+		return 10 + ret;
+
+	if (rb_insert(rbtree, 3, 3))
+		return 5;
+
+	ret = rbtree_check_first(rbtree, 2);
+	if (ret)
+		return 20 + ret;
+
+	if (rb_remove(rbtree, 3) || rb_remove(rbtree, 5))
+		return 6;
+
+	ret = rbtree_check_first(rbtree, 0);
+	if (ret)
+		return 30 + ret;
+
+	return rb_destroy(rbtree);
+}
+
 #define NEXT_NR_NODES	64
 
 static struct rbnode __arena *next_nodes[NEXT_NR_NODES];
@@ -342,6 +482,75 @@ __weak int test_rbtree_next(void)
 	return rb_destroy(rbtree);
 }
 
+struct rb_embed {
+	u64 pad;
+	struct rbtree tree;
+};
+
+SEC("syscall")
+__weak int test_rbtree_embedded(void)
+{
+	struct rb_embed __arena *embed;
+	struct rbtree __arena *rbtree;
+	struct rbnode __arena *node;
+	int ret, i;
+
+	embed = arena_malloc(sizeof(*embed));
+	if (!embed)
+		return 1;
+
+	rbtree = &embed->tree;
+
+	/* Poison the fields rb_init() must set. */
+	WRITE_ONCE(rbtree->root, (struct rbnode __arena *)~0ULL);
+	WRITE_ONCE(rbtree->first, (struct rbnode __arena *)~0ULL);
+	WRITE_ONCE(rbtree->nr, ~0ULL);
+
+	if (rb_init(rbtree, RB_NOALLOC, RB_DUPLICATE))
+		return 2;
+
+	ret = rbtree_check_first(rbtree, 0);
+	if (ret)
+		return 10 + ret;
+
+	for (i = 0; i < 4 && can_loop; i++) {
+		node = rb_node_alloc(4 - i, i);
+		if (!node)
+			return 3;
+
+		if (rb_insert_node(rbtree, node))
+			return 4;
+
+		ret = rbtree_check_first(rbtree, i + 1);
+		if (ret)
+			return 20 + ret;
+	}
+
+	if (rb_fini(rbtree) != -EBUSY)
+		return 7;
+
+	for (i = 0; i < 4 && can_loop; i++) {
+		node = rb_first_lockless(rbtree);
+		if (!node || node->key != i + 1)
+			return 5;
+
+		if (rb_remove_node(rbtree, node))
+			return 6;
+
+		rb_node_free(node);
+
+		ret = rbtree_check_first(rbtree, 3 - i);
+		if (ret)
+			return 30 + ret;
+	}
+
+	if (rb_fini(rbtree))
+		return 8;
+
+	arena_free(embed);
+	return 0;
+}
+
 static inline int
 clean_up_noalloc_tree(struct rbtree __arena *rbtree)
 {
diff --git a/tools/testing/selftests/bpf/libarena/src/rbtree.bpf.c b/tools/testing/selftests/bpf/libarena/src/rbtree.bpf.c
index bf39cb2aa751..90fa0036e9d2 100644
--- a/tools/testing/selftests/bpf/libarena/src/rbtree.bpf.c
+++ b/tools/testing/selftests/bpf/libarena/src/rbtree.bpf.c
@@ -15,14 +15,12 @@ static int rbnode_replace(struct rbtree __arena *rbtree,
 			  struct rbnode __arena *existing,
 			  struct rbnode __arena *replacement);
 
-struct rbtree __arena *rb_create(enum rbtree_alloc alloc,
-				 enum rbtree_insert_mode insert)
+__weak
+int rb_init(struct rbtree __arena *rbtree, enum rbtree_alloc alloc,
+	    enum rbtree_insert_mode insert)
 {
-	struct rbtree __arena *rbtree;
-
-	rbtree = arena_malloc(sizeof(*rbtree));
 	if (unlikely(!rbtree))
-		return NULL;
+		return -EINVAL;
 
 	/*
 	 * RB_UPDATE overwrites existing values in the nodes, but RB_NOALLOC
@@ -32,19 +30,47 @@ struct rbtree __arena *rb_create(enum rbtree_alloc alloc,
 	 */
 	if (alloc == RB_NOALLOC && insert == RB_UPDATE) {
 		arena_stderr("WARNING: Cannot combine RB_NOALLOC and RB_UPDATE");
+		return -EINVAL;
+	}
+
+	WRITE_ONCE(rbtree->root, NULL);
+	WRITE_ONCE(rbtree->alloc, alloc);
+	WRITE_ONCE(rbtree->insert, insert);
+	WRITE_ONCE(rbtree->first, NULL);
+	WRITE_ONCE(rbtree->nr, 0);
+
+	return 0;
+}
+
+struct rbtree __arena *rb_create(enum rbtree_alloc alloc,
+				 enum rbtree_insert_mode insert)
+{
+	struct rbtree __arena *rbtree;
+
+	rbtree = arena_malloc(sizeof(*rbtree));
+	if (unlikely(!rbtree))
+		return NULL;
+
+	if (rb_init(rbtree, alloc, insert)) {
 		arena_free(rbtree);
 		return NULL;
 	}
 
-	rbtree->alloc = alloc;
-	rbtree->insert = insert;
-	rbtree->root = NULL;
-
 	return rbtree;
 }
 
+static void rb_set_first(struct rbtree __arena *rbtree,
+			 struct rbnode __arena *node)
+{
+	/*
+	 * Publish the new leftmost node with release ordering, so a lockless
+	 * reader that sees it also sees the node's fields.
+	 */
+	smp_store_release(&rbtree->first, node);
+}
+
 __weak
-int rb_destroy(struct rbtree __arena *rbtree)
+int rb_fini(struct rbtree __arena *rbtree)
 {
 	int ret = 0;
 
@@ -63,11 +89,11 @@ int rb_destroy(struct rbtree __arena *rbtree)
 		 * the now freed nodes. Throw an error instead.
 		 */
 		if (rbtree->root) {
-			arena_stderr("WARNING: Destroying RB_NOALLOC tree with > 0 nodes");
+			arena_stderr("WARNING: Finalizing RB_NOALLOC tree with > 0 nodes");
 			return -EBUSY;
 		}
 
-		goto out;
+		return 0;
 	}
 
 	while (rbtree->root && can_loop) {
@@ -76,11 +102,22 @@ int rb_destroy(struct rbtree __arena *rbtree)
 			break;
 	}
 
-out:
-	arena_free(rbtree);
 	return ret;
 }
 
+__weak
+int rb_destroy(struct rbtree __arena *rbtree)
+{
+	int ret;
+
+	ret = rb_fini(rbtree);
+	if (ret)
+		return ret;
+
+	arena_free(rbtree);
+	return 0;
+}
+
 static inline int rbnode_dir(struct rbnode __arena *node)
 {
 	/* Arbitrarily choose a direction for the root. */
@@ -249,6 +286,8 @@ int rb_node_insert(struct rbtree __arena *rbtree,
 
 	if (!parent) {
 		rbtree->root = node;
+		rb_set_first(rbtree, node);
+		WRITE_ONCE(rbtree->nr, 1);
 		return 0;
 	}
 
@@ -267,6 +306,10 @@ int rb_node_insert(struct rbtree __arena *rbtree,
 			if (ret)
 				return ret;
 
+			/* Same key and place, so only the head's identity changes. */
+			if (rbtree->first == parent)
+				rb_set_first(rbtree, node);
+
 			if (rbtree->alloc == RB_ALLOC)
 				rb_node_free(parent);
 
@@ -284,6 +327,14 @@ int rb_node_insert(struct rbtree __arena *rbtree,
 	else
 		parent->right = node;
 
+	/*
+	 * The node is the new head only if it went left of the old one. The
+	 * rotations below never change which node is leftmost.
+	 */
+	if (parent == rbtree->first && parent->left == node)
+		rb_set_first(rbtree, node);
+	WRITE_ONCE(rbtree->nr, rbtree->nr + 1);
+
 	while (can_loop) {
 		parent = node->parent;
 		if (!parent)
@@ -404,7 +455,7 @@ __weak int rb_least(struct rbtree __arena *rbtree, u64 *key, u64 *value)
 	if (!rbtree->root)
 		return -ENOENT;
 
-	least = rbnode_least(rbtree->root);
+	least = rbtree->first;
 	if (key)
 		*key = least->key;
 	if (value)
@@ -601,6 +652,15 @@ int rb_node_remove(struct rbtree __arena *rbtree,
 	bool is_red;
 	int dir;
 
+	/*
+	 * Hand the head to its successor before the tree changes. The removal
+	 * frees @node in an RB_ALLOC tree, so the head must not point to it
+	 * by then.
+	 */
+	if (rbtree->first == node)
+		rb_set_first(rbtree, rb_next(node));
+	WRITE_ONCE(rbtree->nr, rbtree->nr - 1);
+
 	/* Both children present, replace with next largest key. */
 	if (node->left && node->right) {
 		/*
@@ -827,7 +887,7 @@ int rb_pop(struct rbtree __arena *rbtree, u64 *key, u64 *value)
 	if (rbtree->alloc != RB_ALLOC)
 		return -EINVAL;
 
-	node = rbnode_least(rbtree->root);
+	node = rbtree->first;
 	if (unlikely(!node))
 		return -ENOENT;
 
-- 
2.56.0