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

[Changwoo Min <changwoo@igalia.com>]

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>
---
 .../bpf/libarena/include/libarena/rbtree.h    |  6 +--
 .../bpf/libarena/selftests/test_rbtree.bpf.c  | 37 +++++++++++++++++++
 .../selftests/bpf/libarena/src/rbtree.bpf.c   | 16 +++++---
 3 files changed, 51 insertions(+), 8 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..e6c9430f734b 100644
--- a/tools/testing/selftests/bpf/libarena/src/rbtree.bpf.c
+++ b/tools/testing/selftests/bpf/libarena/src/rbtree.bpf.c
@@ -161,8 +161,14 @@ struct rbnode __arena *rbnode_find(struct rbnode __arena *subtree, u64 key)
 	return node;
 }
 
+/*
+ * Find the node to link @key under. Smaller keys go left and everything else
+ * right, so a new key lands after every node with an equal key and equal keys
+ * stay in insertion order. Rotations and removals preserve the in-order
+ * sequence, so they keep that order too.
+ */
 static
-struct rbnode __arena *rbnode_least_upper_bound(struct rbnode __arena *subtree, uint64_t key)
+struct rbnode __arena *rbnode_insert_parent(struct rbnode __arena *subtree, uint64_t key)
 {
 	struct rbnode __arena *node = subtree;
 	int dir;
@@ -171,7 +177,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;
@@ -255,7 +261,7 @@ int rb_node_insert(struct rbtree __arena *rbtree,
 	if (rbtree->insert != RB_DUPLICATE)
 		parent = rbnode_find(parent, key);
 	else
-		parent = rbnode_least_upper_bound(parent, key);
+		parent = rbnode_insert_parent(parent, key);
 
 	if (key == parent->key && rbtree->insert != RB_DUPLICATE) {
 		if (rbtree->insert == RB_UPDATE) {
@@ -278,8 +284,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