[PATCH bpf-next v3 02/10] rhashtable: Add rhashtable_walk_enter_from()

[Mykyta Yatsenko <mykyta.yatsenko5@gmail.com>]

From: Mykyta Yatsenko <yatsenko@meta.com>

BPF resizable hashmap needs efficient iteration resume for
get_next_key and seq_file iterators. rhashtable_walk_enter()
always starts from bucket 0, forcing linear skip of already-seen
elements.

Add rhashtable_walk_enter_from() that looks up the key's bucket
and positions the walker there, so walk_next returns the successor
directly. If a resize moved the key to the future table, the
walker resets to the beginning.

Refactor __rhashtable_lookup into __rhashtable_lookup_one to reuse
the single-table lookup in both the two-table search and the new
enter_from positioning.

Signed-off-by: Mykyta Yatsenko <yatsenko@meta.com>
---
 include/linux/rhashtable.h |  31 +++++++++--
 lib/rhashtable.c           |  41 ++++++++++++++
 lib/test_rhashtable.c      | 131 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 198 insertions(+), 5 deletions(-)

diff --git a/include/linux/rhashtable.h b/include/linux/rhashtable.h
index 0480509a6339..6ac547d6c619 100644
--- a/include/linux/rhashtable.h
+++ b/include/linux/rhashtable.h
@@ -246,6 +246,11 @@ static inline void rhashtable_walk_start(struct rhashtable_iter *iter)
 	(void)rhashtable_walk_start_check(iter);
 }
 
+void rhashtable_walk_enter_from(struct rhashtable *ht,
+				struct rhashtable_iter *iter,
+				const void *key,
+				const struct rhashtable_params params);
+
 void *rhashtable_walk_next(struct rhashtable_iter *iter);
 void *rhashtable_walk_peek(struct rhashtable_iter *iter);
 void rhashtable_walk_stop(struct rhashtable_iter *iter) __releases_shared(RCU);
@@ -606,8 +611,8 @@ static inline int rhashtable_compare(struct rhashtable_compare_arg *arg,
 }
 
 /* Internal function, do not use. */
-static __always_inline struct rhash_head *__rhashtable_lookup(
-	struct rhashtable *ht, const void *key,
+static __always_inline struct rhash_head *__rhashtable_lookup_one(
+	struct rhashtable *ht, struct bucket_table *tbl, const void *key,
 	const struct rhashtable_params params,
 	const enum rht_lookup_freq freq)
 	__must_hold_shared(RCU)
@@ -617,13 +622,10 @@ static __always_inline struct rhash_head *__rhashtable_lookup(
 		.key = key,
 	};
 	struct rhash_lock_head __rcu *const *bkt;
-	struct bucket_table *tbl;
 	struct rhash_head *he;
 	unsigned int hash;
 
 	BUILD_BUG_ON(!__builtin_constant_p(freq));
-	tbl = rht_dereference_rcu(ht->tbl, ht);
-restart:
 	hash = rht_key_hashfn(ht, tbl, key, params);
 	bkt = rht_bucket(tbl, hash);
 	do {
@@ -639,6 +641,25 @@ static __always_inline struct rhash_head *__rhashtable_lookup(
 		 */
 	} while (he != RHT_NULLS_MARKER(bkt));
 
+	return NULL;
+}
+
+/* Internal function, do not use. */
+static __always_inline struct rhash_head *__rhashtable_lookup(
+	struct rhashtable *ht, const void *key,
+	const struct rhashtable_params params,
+	const enum rht_lookup_freq freq)
+	__must_hold_shared(RCU)
+{
+	struct bucket_table *tbl;
+	struct rhash_head *he;
+
+	tbl = rht_dereference_rcu(ht->tbl, ht);
+restart:
+	he = __rhashtable_lookup_one(ht, tbl, key, params, freq);
+	if (he)
+		return he;
+
 	/* Ensure we see any new tables. */
 	smp_rmb();
 
diff --git a/lib/rhashtable.c b/lib/rhashtable.c
index 6074ed5f66f3..2a6098edf737 100644
--- a/lib/rhashtable.c
+++ b/lib/rhashtable.c
@@ -692,6 +692,47 @@ void rhashtable_walk_enter(struct rhashtable *ht, struct rhashtable_iter *iter)
 }
 EXPORT_SYMBOL_GPL(rhashtable_walk_enter);
 
+/**
+ * rhashtable_walk_enter_from - Initialise a walk starting at a key's bucket
+ * @ht:		Table to walk over
+ * @iter:	Hash table iterator
+ * @key:	Key whose bucket to start from
+ * @params:	Hash table parameters
+ *
+ * Like rhashtable_walk_enter(), but positions the iterator at the bucket
+ * containing @key.  If @key is not present in the main table, the iterator is
+ * positioned at the beginning.
+ *
+ * Same constraints as rhashtable_walk_enter() apply.
+ */
+void rhashtable_walk_enter_from(struct rhashtable *ht,
+				struct rhashtable_iter *iter,
+				const void *key,
+				const struct rhashtable_params params)
+				__must_hold(RCU)
+{
+	struct bucket_table *tbl;
+	struct rhash_head *he;
+
+	rhashtable_walk_enter(ht, iter);
+
+	if (!key)
+		return;
+
+	tbl = iter->walker.tbl;
+	if (!tbl)
+		return;
+
+	he = __rhashtable_lookup_one(ht, tbl, key, params,
+				     RHT_LOOKUP_NORMAL);
+	if (!he)
+		return;
+
+	iter->slot = rht_key_hashfn(ht, tbl, key, params);
+	iter->p = he;
+}
+EXPORT_SYMBOL_GPL(rhashtable_walk_enter_from);
+
 /**
  * rhashtable_walk_exit - Free an iterator
  * @iter:	Hash table Iterator
diff --git a/lib/test_rhashtable.c b/lib/test_rhashtable.c
index 0b33559a910b..0c4f489bf655 100644
--- a/lib/test_rhashtable.c
+++ b/lib/test_rhashtable.c
@@ -23,6 +23,7 @@
 #include <linux/random.h>
 #include <linux/vmalloc.h>
 #include <linux/wait.h>
+#include <linux/cleanup.h>
 
 #define MAX_ENTRIES	1000000
 #define TEST_INSERT_FAIL INT_MAX
@@ -679,6 +680,133 @@ static int threadfunc(void *data)
 	return err;
 }
 
+static int __init test_walk_enter_from(void)
+{
+	struct rhashtable ht;
+	struct test_obj objs[4];
+	struct rhashtable_iter iter;
+	struct test_obj *obj;
+	struct test_obj_val key = { };
+	int err, i, count;
+	int len_freq[4] = { 0, 0, 0, 0 };
+
+	err = rhashtable_init(&ht, &test_rht_params);
+	if (err)
+		return err;
+
+	/* Insert 4 elements with keys 0, 2, 4, 6 */
+	for (i = 0; i < 4; i++) {
+		objs[i].value.id = i * 2;
+		objs[i].value.tid = 0;
+		err = rhashtable_insert_fast(&ht, &objs[i].node,
+					     test_rht_params);
+		if (err) {
+			pr_warn("walk_enter_from: insert %d failed: %d\n",
+				i, err);
+			goto out;
+		}
+	}
+
+	/*
+	 * Test 1: walk_enter_from positions at key, walk_next returns
+	 * the successor (not the key itself).
+	 */
+	for (i = 0; i < 4; i++) {
+		key.id = i * 2;
+		scoped_guard(rcu) {
+			rhashtable_walk_enter_from(&ht, &iter, &key,
+						   test_rht_params);
+			rhashtable_walk_start(&iter);
+		}
+
+		do {
+			obj = rhashtable_walk_next(&iter);
+		} while (PTR_ERR(obj) == -EAGAIN);
+
+		/* Successor must not be the key itself */
+		if (IS_ERR(obj) || (obj && obj->value.id == i * 2)) {
+			pr_warn("walk_enter_from: returned err %ld or key %d instead of successor\n",
+				PTR_ERR(obj), i * 2);
+			err = -EINVAL;
+		}
+
+		rhashtable_walk_stop(&iter);
+		rhashtable_walk_exit(&iter);
+
+		if (err)
+			goto out;
+	}
+
+	/* Test 2: walk_enter_from with non-existent key starts from the beginning */
+	key.id = 99;
+	scoped_guard(rcu) {
+		rhashtable_walk_enter_from(&ht, &iter, &key, test_rht_params);
+		rhashtable_walk_start(&iter);
+	}
+
+	do {
+		obj = rhashtable_walk_next(&iter);
+	} while (PTR_ERR(obj) == -EAGAIN);
+
+	if (IS_ERR(obj)) {
+		pr_warn("walk_enter_from: returned err %ld\n", PTR_ERR(obj));
+		err = -EINVAL;
+	}
+
+	/* Should still return some element (iteration from bucket start) */
+	rhashtable_walk_stop(&iter);
+	rhashtable_walk_exit(&iter);
+
+	if (err)
+		goto out;
+
+	/* Test 3: verify walk_enter_from + walk_next can iterate remaining elements */
+	for (i = 0; i < 4; i++) {
+		key.id = i * 2;
+		count = 0;
+
+		scoped_guard(rcu) {
+			rhashtable_walk_enter_from(&ht, &iter, &key, test_rht_params);
+			rhashtable_walk_start(&iter);
+		}
+
+		/* Count the number of elements after this one */
+		do {
+			obj = rhashtable_walk_next(&iter);
+			count++;
+		} while (obj && (!IS_ERR(obj) || (PTR_ERR(obj) == -EAGAIN)));
+
+		rhashtable_walk_stop(&iter);
+		rhashtable_walk_exit(&iter);
+		if (count > 4) {
+			pr_warn("walk_enter_from: walked %d elements instead of max 4\n", count);
+			goto out;
+		}
+		/*
+		 * Collect the lengths of every walked path,
+		 * we should have each length walked one time
+		 */
+		len_freq[count - 1]++;
+	}
+
+	for (i = 0; i < 4; i++) {
+		/*
+		 * Should see at least some elements after key 0.
+		 * Exact count depends on hash distribution.
+		 */
+		if (len_freq[i] != 1) {
+			pr_warn("walk_enter_from: duplicate tail length\n");
+			err = -EINVAL;
+			goto out;
+		}
+	}
+
+	pr_info("walk_enter_from: all tests passed\n");
+out:
+	rhashtable_destroy(&ht);
+	return err;
+}
+
 static int __init test_rht_init(void)
 {
 	unsigned int entries;
@@ -738,6 +866,9 @@ static int __init test_rht_init(void)
 
 	test_insert_duplicates_run();
 
+	pr_info("Testing walk_enter_from: %s\n",
+		test_walk_enter_from() == 0 ? "pass" : "FAIL");
+
 	if (!tcount)
 		return 0;
 

-- 
2.52.0