The original sbitmap performed a cleanup operation before acquiring a bit,
but this was sometimes unnecessary overhead. This patch reduced the
frequency of cleanup operations, executing them only when necessary, and
abstracted a helper function that can acquire multiple free bits in a
single operation.

Signed-off-by: hexue <xue01.he@samsung.com>
---
 lib/sbitmap.c | 71 ++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 47 insertions(+), 24 deletions(-)

diff --git a/lib/sbitmap.c b/lib/sbitmap.c
index 4d188d05db15..98fd27a896f1 100644
--- a/lib/sbitmap.c
+++ b/lib/sbitmap.c
@@ -208,6 +208,43 @@ static int sbitmap_find_bit_in_word(struct sbitmap_word *map,
 	return nr;
 }
 
+static unsigned long sbitmap_find_bits_in_word(struct sbitmap_word *map,
+					unsigned int map_depth, int nr_tags, unsigned long *nr)
+{
+	unsigned long val, bit_nr, get_mask;
+
+	while (1) {
+		val = READ_ONCE(map->word);
+		if (val == (1UL << (map_depth - 1)) - 1) {
+			if (!sbitmap_deferred_clear(map, 0, 0, 0))
+				return 0;
+			continue;
+		}
+
+		bit_nr = find_first_zero_bit(&val, map_depth);
+
+		/* Ensure that the lengths of get_mask and val are consistent
+		 * to avoid NULL pointer dereference
+		 */
+		if (bit_nr + nr_tags <= map_depth)
+			break;
+
+		if (!sbitmap_deferred_clear(map, 0, 0, 0))
+			return 0;
+	};
+
+	atomic_long_t *ptr = (atomic_long_t *) &map->word;
+
+	get_mask = ((1UL << nr_tags) - 1) << bit_nr;
+	while (!atomic_long_try_cmpxchg(ptr, &val,
+					  get_mask | val))
+		;
+	get_mask = (get_mask & ~val) >> bit_nr;
+
+	*nr = bit_nr;
+	return get_mask;
+}
+
 static unsigned int __map_depth_with_shallow(const struct sbitmap *sb,
 					     int index,
 					     unsigned int shallow_depth)
@@ -517,7 +554,7 @@ unsigned long __sbitmap_queue_get_batch(struct sbitmap_queue *sbq, int nr_tags,
 {
 	struct sbitmap *sb = &sbq->sb;
 	unsigned int hint, depth;
-	unsigned long index, nr;
+	unsigned long index;
 	int i;
 
 	if (unlikely(sb->round_robin))
@@ -530,32 +567,18 @@ unsigned long __sbitmap_queue_get_batch(struct sbitmap_queue *sbq, int nr_tags,
 
 	for (i = 0; i < sb->map_nr; i++) {
 		struct sbitmap_word *map = &sb->map[index];
-		unsigned long get_mask;
+		unsigned long get_mask, nr;
 		unsigned int map_depth = __map_depth(sb, index);
-		unsigned long val;
 
-		sbitmap_deferred_clear(map, 0, 0, 0);
-		val = READ_ONCE(map->word);
-		if (val == (1UL << (map_depth - 1)) - 1)
-			goto next;
-
-		nr = find_first_zero_bit(&val, map_depth);
-		if (nr + nr_tags <= map_depth) {
-			atomic_long_t *ptr = (atomic_long_t *) &map->word;
-
-			get_mask = ((1UL << nr_tags) - 1) << nr;
-			while (!atomic_long_try_cmpxchg(ptr, &val,
-							  get_mask | val))
-				;
-			get_mask = (get_mask & ~val) >> nr;
-			if (get_mask) {
-				*offset = nr + (index << sb->shift);
-				update_alloc_hint_after_get(sb, depth, hint,
-							*offset + nr_tags - 1);
-				return get_mask;
-			}
+		get_mask = sbitmap_find_bits_in_word(map, map_depth, nr_tags, &nr);
+
+		if (get_mask) {
+			*offset = nr + (index << sb->shift);
+			update_alloc_hint_after_get(sb, depth, hint,
+						*offset + nr_tags - 1);
+			return get_mask;
 		}
-next:
+
 		/* Jump to next index. */
 		if (++index >= sb->map_nr)
 			index = 0;
-- 
2.34.1