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

[Changwoo Min <changwoo@igalia.com>]

On 11/2/26 5:05 PM, Emil Tsalapatis wrote:
> On Wed Oct 28, 2026 at 6:00 AM UTC, Changwoo Min wrote:
>> 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. rb_least() and rb_pop() use the
>> cached head. rb_first(), rb_first_key() and rb_nr_nodes() read them
>> without the lock. The head pointer, published with release ordering,
>> decides emptiness, and the key is read through it. Both fields sit on
>> their own cache line.
>>
>> 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>
>> ---
>>   .../bpf/libarena/include/libarena/rbtree.h    |  53 +++++
>>   .../bpf/libarena/selftests/test_rbtree.bpf.c  | 209 ++++++++++++++++++
>>   .../selftests/bpf/libarena/src/rbtree.bpf.c   | 122 ++++++++--
>>   3 files changed, 365 insertions(+), 19 deletions(-)
>>
>> diff --git a/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h b/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h
>> index 00e63b12d82e..a8dd8c2e7596 100644
>> --- a/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h
>> +++ b/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h
>> @@ -4,6 +4,9 @@
>>   
>>   #define RB_MAXLVL_PRINT (16)
>>   
>> +/* Cache line size assumed for keeping the lockless fields apart. */
>> +#define RB_CACHELINE_SIZE (64)
>> +
>>   struct rbnode;
>>   
>>   struct rbnode {
>> @@ -56,14 +59,34 @@ enum rbtree_insert_mode {
>>   };
>>   
>>   struct rbtree {
>> +	/* Writer fields, on a cache line of their own. */
>>   	struct rbnode __arena *root;
>>   	enum rbtree_alloc alloc;
>>   	enum rbtree_insert_mode insert;
>> +
>> +	/*
>> +	 * The leftmost node and the node count on their own cache line.
>> +	 * The writer keeps them current, and readers may read them without
>> +	 * the writer's lock (see rb_first() and friends).
>> +	 */
>> +	struct {
>> +		struct rbnode __arena *first;
>> +		uint64_t nr;
>> +	} __attribute__((aligned(RB_CACHELINE_SIZE)));
>>   };
>>   
>>   #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);
> 
> Unless we are constantly creating/destroying trees, I don't think this is necessary
> - allocations are in the 100ns range. Is there a reason to require being able to embed
> the tree into other data structures?

In my use case for a task queue, allocation cost doesn't matter; the
life cycle of an rbtree (a task queue) is the same as a sched-ext
scheduler. The main reason is to embed the struct to avoid pointer
dereferencing (which probably causes another cache miss) in the critical
path.

I measured it with a sched-ext side task queue built on top of the
rbtree (using scx_lavd replacing DSQs) on a 2-socket, 192-CPU EPYC
machine (no hyperthreading). I compared two versions: 1) an rbtree is
embedded in a task queue (embedded), and 2)  an rbtree is separately
allocated so it is accessed through a pointer in the task queue
(pointer). I ran schbench with full load (-m 4 -t 48). Each benchmark
ran 10 times, alternating the two versions:

                wakeup p99 (us)
                embedded    pointer
   run 1        78.0        82.7   (+6.0%)
   run 2        77.7        81.9   (+5.4%)
   run 3        79.4        85.8   (+8.1%)

While there are no statistically meaningful differences in throughput
and median latency, p99 latency consistently shows 5-8% overhead across
runs. So I would like to keep rb_init() and rb_fini().

> 
>> +
>>   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);
>> @@ -80,4 +103,34 @@ 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() returns may have
>> + * left the tree, or been freed in an RB_ALLOC tree, by the time the
>> + * caller reads it. The count is a separate read and need not agree
>> + * with the head.
>> + */
> 
> Can we somehow communicate in the function name that it has nonstandard semantics?

What about adding *_lockless() suffix? So they will be
rb_first_lockless(), rb_first_key_lockless()
and rb_nr_nodes_lockless().

> Also sorry if we've discussed this before, but what is the exact use of getting
> the node address back speculatively if we cannot guarantee the validity of the
> pointer at all? Doesn't rb_first_key cover the case where we locklessly peek
> the key?

Some use cases require more than the key. In the case of an rbtree-based
task queue, vtime is the key of the rbtree. A BPF scheduler (like
scx_lavd) peeks at the head of the task queue to determine task size
(i.e., task load) to make a load-balancing decision. The task load is
stored in the task context, where an rbnode is embedded. In this case,
the rbnode stays valid because a task context is freed only after an RCU
grace period, and the caller of the lockless peek is in the RCU read-
critical section. Basically, the caller is responsible for safely
reading the pointer. I will comment this clearly in v2. What do you
think?

> 
>> +static __always_inline struct rbnode __arena *rb_first(struct rbtree __arena *rbtree)
>> +{
>> +	return READ_ONCE(rbtree->first);
>> +}
>> +
>> +/* -ENOENT if the tree is empty, 0 with the leftmost key in @key otherwise. */
>> +static __always_inline int rb_first_key(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 __always_inline u64 rb_nr_nodes(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 1e2a0fac1e4f..57c4395c046a 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,215 @@ __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(rbtree) != least)
>> +		return 1;
>> +
>> +	if (rb_nr_nodes(rbtree) != nr)
>> +		return 2;
>> +
>> +	if (!least)
>> +		return rb_first_key(rbtree, &key) == -ENOENT ? 0 : 3;
>> +
>> +	if (rb_first_key(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(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(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);
>> +}
>> +
>> +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(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 e6c9430f734b..2622c78302af 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,51 @@ 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;
>> +}
>> +
>> +/*
>> + * arena_malloc() aligns a block to its power-of-two size, so the tree's
>> + * lockless fields get a cache line of their own.
>> + */
>> +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 __always_inline 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 +93,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 +106,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. */
>> @@ -255,6 +296,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;
>>   	}
>>   
>> @@ -273,6 +316,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);
>>   
>> @@ -290,6 +337,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)
>> @@ -410,7 +465,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)
>> @@ -578,8 +633,8 @@ static inline bool rbnode_has_red_children(struct rbnode __arena *node)
>>   }
>>   
>>   static
>> -int rb_node_remove(struct rbtree __arena *rbtree,
>> -		   struct rbnode __arena *node)
>> +int rbnode_remove(struct rbtree __arena *rbtree,
>> +		  struct rbnode __arena *node)
>>   {
>>   	struct rbnode __arena *parent, *sibling, *close_nephew, *distant_nephew;
>>   	bool free = (rbtree->alloc == RB_ALLOC);
>> @@ -765,6 +820,35 @@ int rb_node_remove(struct rbtree __arena *rbtree,
>>   	return 0;
>>   }
>>   
>> +/*
>> + * Remove @node and keep the head and the node count current. The head has no
>> + * left child, so if it is the node going away, its successor is the leftmost
>> + * node of its right subtree, or else its parent. Take it before the tree
>> + * changes; the removal and its fixups preserve the in-order sequence.
>> + */
>> +static __always_inline
> 
> We don't need the __always_inline for static functions, they are there to help
> with verifier limits but upstream libarena tracks bpf-next which can verify
> pretty much anything. I'll re-add those as needed when importing to Github
> in the compat branch.

Sure. I will remove it in v2.

Regards,
Changwoo Min>
>> +int rb_node_remove(struct rbtree __arena *rbtree,
>> +		   struct rbnode __arena *node)
>> +{
>> +	struct rbnode __arena *next = NULL;
>> +	bool was_first;
>> +	int ret;
>> +
>> +	was_first = node == rbtree->first;
>> +	if (was_first)
>> +		next = node->right ? rbnode_least(node->right) : node->parent;
>> +
>> +	ret = rbnode_remove(rbtree, node);
>> +	if (ret)
>> +		return ret;
>> +
>> +	if (was_first)
>> +		rb_set_first(rbtree, next);
>> +	WRITE_ONCE(rbtree->nr, rbtree->nr - 1);
>> +
>> +	return 0;
>> +}
>> +
>>   __weak
>>   int rb_remove_node(struct rbtree __arena *rbtree,
>>   		   struct rbnode __arena *node)
>> @@ -813,7 +897,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;
>>   
>