[PATCH 0/5] A few cleanup and bugfix patches for sbitmap

[PATCH 0/5] A few cleanup and bugfix patches for sbitmap

[Kemeng Shi]

Hi, this series contain a patch to fix protential wakup lost in
__sbq_wake_up and some random cleanup patches to remove unnecessary
check and repeat code.
Thanks.

Kemeng Shi (5):
  sbitmap: don't consume nr for inactive waitqueue to avoid lost wakeups
  sbitmap: remove unnecessary calculation of alloc_hint in
    __sbitmap_get_shallow
  sbitmap: remove redundant check in __sbitmap_queue_get_batch
  sbitmap: rewrite sbitmap_find_bit_in_index to reduce repeat code
  sbitmap: add sbitmap_find_bit to remove repeat code in
    __sbitmap_get/__sbitmap_get_shallow

 lib/sbitmap.c | 111 +++++++++++++++++++++++---------------------------
 1 file changed, 51 insertions(+), 60 deletions(-)

-- 
2.30.0

[PATCH 1/5] sbitmap: don't consume nr for inactive waitqueue to avoid lost wakeups

[Kemeng Shi]

If we decremented queue without waiters, we should not decremente freed
bits number "nr", or all "nr" could be consumed in a empty queue and no
wakeup will be called.
Currently, for case "wait_cnt > 0", "nr" will not be decremented if we
decremented queue without watiers and retry is returned to avoid lost
wakeups. However for case "wait_cnt == 0", "nr" will be decremented
unconditionally and maybe decremented to zero. Although retry is
returned by active state of queue, it's not actually executed for "nr"
is zero.

Fix this by only decrementing "nr" for active queue when "wait_cnt ==
0". After this fix, "nr" will always be non-zero when we decremented
inactive queue for case "wait_cnt == 0", so the need to retry could
be returned by "nr" and active state of waitqueue returned for the same
purpose is not needed.

Signed-off-by: Kemeng Shi <shikemeng@huawei.com>
---
 lib/sbitmap.c | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/lib/sbitmap.c b/lib/sbitmap.c
index 7280ae8ca88c..e40759bcf821 100644
--- a/lib/sbitmap.c
+++ b/lib/sbitmap.c
@@ -604,7 +604,6 @@ static bool __sbq_wake_up(struct sbitmap_queue *sbq, int *nr)
 	struct sbq_wait_state *ws;
 	unsigned int wake_batch;
 	int wait_cnt, cur, sub;
-	bool ret;
 
 	if (*nr <= 0)
 		return false;
@@ -632,15 +631,15 @@ static bool __sbq_wake_up(struct sbitmap_queue *sbq, int *nr)
 	if (wait_cnt > 0)
 		return !waitqueue_active(&ws->wait);
 
-	*nr -= sub;
-
 	/*
 	 * When wait_cnt == 0, we have to be particularly careful as we are
 	 * responsible to reset wait_cnt regardless whether we've actually
-	 * woken up anybody. But in case we didn't wakeup anybody, we still
-	 * need to retry.
+	 * woken up anybody. But in case we didn't wakeup anybody, we should
+	 * not consume nr and need to retry to avoid lost wakeups.
 	 */
-	ret = !waitqueue_active(&ws->wait);
+	if (waitqueue_active(&ws->wait))
+		*nr -= sub;
+
 	wake_batch = READ_ONCE(sbq->wake_batch);
 
 	/*
@@ -669,7 +668,7 @@ static bool __sbq_wake_up(struct sbitmap_queue *sbq, int *nr)
 	sbq_index_atomic_inc(&sbq->wake_index);
 	atomic_set(&ws->wait_cnt, wake_batch);
 
-	return ret || *nr;
+	return *nr;
 }
 
 void sbitmap_queue_wake_up(struct sbitmap_queue *sbq, int nr)
-- 
2.30.0

Re: [PATCH 1/5] sbitmap: don't consume nr for inactive waitqueue to avoid lost wakeups

[Kemeng Shi]

on 12/1/2022 12:54 PM, Kemeng Shi wrote:
> If we decremented queue without waiters, we should not decremente freed
> bits number "nr", or all "nr" could be consumed in a empty queue and no
> wakeup will be called.
> Currently, for case "wait_cnt > 0", "nr" will not be decremented if we
> decremented queue without watiers and retry is returned to avoid lost
> wakeups. However for case "wait_cnt == 0", "nr" will be decremented
> unconditionally and maybe decremented to zero. Although retry is
> returned by active state of queue, it's not actually executed for "nr"
> is zero.
> 
> Fix this by only decrementing "nr" for active queue when "wait_cnt ==
> 0". After this fix, "nr" will always be non-zero when we decremented
> inactive queue for case "wait_cnt == 0", so the need to retry could
> be returned by "nr" and active state of waitqueue returned for the same
> purpose is not needed.
> 
> Signed-off-by: Kemeng Shi <shikemeng@huawei.com>
> ---
>  lib/sbitmap.c | 13 ++++++-------
>  1 file changed, 6 insertions(+), 7 deletions(-)
> 
> diff --git a/lib/sbitmap.c b/lib/sbitmap.c
> index 7280ae8ca88c..e40759bcf821 100644
> --- a/lib/sbitmap.c
> +++ b/lib/sbitmap.c
> @@ -604,7 +604,6 @@ static bool __sbq_wake_up(struct sbitmap_queue *sbq, int *nr)
>  	struct sbq_wait_state *ws;
>  	unsigned int wake_batch;
>  	int wait_cnt, cur, sub;
> -	bool ret;
>  
>  	if (*nr <= 0)
>  		return false;
> @@ -632,15 +631,15 @@ static bool __sbq_wake_up(struct sbitmap_queue *sbq, int *nr)
>  	if (wait_cnt > 0)
>  		return !waitqueue_active(&ws->wait);
>  
> -	*nr -= sub;
> -
>  	/*
>  	 * When wait_cnt == 0, we have to be particularly careful as we are
>  	 * responsible to reset wait_cnt regardless whether we've actually
> -	 * woken up anybody. But in case we didn't wakeup anybody, we still
> -	 * need to retry.
> +	 * woken up anybody. But in case we didn't wakeup anybody, we should
> +	 * not consume nr and need to retry to avoid lost wakeups.
>  	 */
> -	ret = !waitqueue_active(&ws->wait);
There is a warnning reported by checkpatch.pl which is "WARNING:waitqueue_active
without comment" but I don't know why.
> +	if (waitqueue_active(&ws->wait))
> +		*nr -= sub;
> +
>  	wake_batch = READ_ONCE(sbq->wake_batch);
>  
>  	/*
> @@ -669,7 +668,7 @@ static bool __sbq_wake_up(struct sbitmap_queue *sbq, int *nr)
>  	sbq_index_atomic_inc(&sbq->wake_index);
>  	atomic_set(&ws->wait_cnt, wake_batch);
>  
> -	return ret || *nr;
> +	return *nr;
>  }
>  
>  void sbitmap_queue_wake_up(struct sbitmap_queue *sbq, int nr)
> 
Besides, there are some git config problems for my huaweicloud email, I will send
patchset with huaweicloud email when I fix them.

Thanks.
-- 
Best wishes
Kemeng Shi

Re: [PATCH 1/5] sbitmap: don't consume nr for inactive waitqueue to avoid lost wakeups

[Jens Axboe]

On 12/1/22 12:21?AM, Kemeng Shi wrote:
> 
> 
> on 12/1/2022 12:54 PM, Kemeng Shi wrote:
>> If we decremented queue without waiters, we should not decremente freed
>> bits number "nr", or all "nr" could be consumed in a empty queue and no
>> wakeup will be called.
>> Currently, for case "wait_cnt > 0", "nr" will not be decremented if we
>> decremented queue without watiers and retry is returned to avoid lost
>> wakeups. However for case "wait_cnt == 0", "nr" will be decremented
>> unconditionally and maybe decremented to zero. Although retry is
>> returned by active state of queue, it's not actually executed for "nr"
>> is zero.
>>
>> Fix this by only decrementing "nr" for active queue when "wait_cnt ==
>> 0". After this fix, "nr" will always be non-zero when we decremented
>> inactive queue for case "wait_cnt == 0", so the need to retry could
>> be returned by "nr" and active state of waitqueue returned for the same
>> purpose is not needed.
>>
>> Signed-off-by: Kemeng Shi <shikemeng@huawei.com>
>> ---
>>  lib/sbitmap.c | 13 ++++++-------
>>  1 file changed, 6 insertions(+), 7 deletions(-)
>>
>> diff --git a/lib/sbitmap.c b/lib/sbitmap.c
>> index 7280ae8ca88c..e40759bcf821 100644
>> --- a/lib/sbitmap.c
>> +++ b/lib/sbitmap.c
>> @@ -604,7 +604,6 @@ static bool __sbq_wake_up(struct sbitmap_queue *sbq, int *nr)
>>  	struct sbq_wait_state *ws;
>>  	unsigned int wake_batch;
>>  	int wait_cnt, cur, sub;
>> -	bool ret;
>>  
>>  	if (*nr <= 0)
>>  		return false;
>> @@ -632,15 +631,15 @@ static bool __sbq_wake_up(struct sbitmap_queue *sbq, int *nr)
>>  	if (wait_cnt > 0)
>>  		return !waitqueue_active(&ws->wait);
>>  
>> -	*nr -= sub;
>> -
>>  	/*
>>  	 * When wait_cnt == 0, we have to be particularly careful as we are
>>  	 * responsible to reset wait_cnt regardless whether we've actually
>> -	 * woken up anybody. But in case we didn't wakeup anybody, we still
>> -	 * need to retry.
>> +	 * woken up anybody. But in case we didn't wakeup anybody, we should
>> +	 * not consume nr and need to retry to avoid lost wakeups.
>>  	 */
>> -	ret = !waitqueue_active(&ws->wait);
> There is a warnning reported by checkpatch.pl which is
> "WARNING:waitqueue_active without comment" but I don't know why.

Most likely because waitqueue_active() could be racy, so a comment is
warranted on why it's safe rather than using wq_has_sleeper().

-- 
Jens Axboe

Re: [PATCH 1/5] sbitmap: don't consume nr for inactive waitqueue to avoid lost wakeups

[Kemeng Shi]

on 12/2/2022 8:58 AM, Jens Axboe wrote:
> On 12/1/22 12:21?AM, Kemeng Shi wrote:
>>
>>
>> on 12/1/2022 12:54 PM, Kemeng Shi wrote:
>>> If we decremented queue without waiters, we should not decremente freed
>>> bits number "nr", or all "nr" could be consumed in a empty queue and no
>>> wakeup will be called.
>>> Currently, for case "wait_cnt > 0", "nr" will not be decremented if we
>>> decremented queue without watiers and retry is returned to avoid lost
>>> wakeups. However for case "wait_cnt == 0", "nr" will be decremented
>>> unconditionally and maybe decremented to zero. Although retry is
>>> returned by active state of queue, it's not actually executed for "nr"
>>> is zero.
>>>
>>> Fix this by only decrementing "nr" for active queue when "wait_cnt ==
>>> 0". After this fix, "nr" will always be non-zero when we decremented
>>> inactive queue for case "wait_cnt == 0", so the need to retry could
>>> be returned by "nr" and active state of waitqueue returned for the same
>>> purpose is not needed.
>>>
>>> Signed-off-by: Kemeng Shi <shikemeng@huawei.com>
>>> ---
>>>  lib/sbitmap.c | 13 ++++++-------
>>>  1 file changed, 6 insertions(+), 7 deletions(-)
>>>
>>> diff --git a/lib/sbitmap.c b/lib/sbitmap.c
>>> index 7280ae8ca88c..e40759bcf821 100644
>>> --- a/lib/sbitmap.c
>>> +++ b/lib/sbitmap.c
>>> @@ -604,7 +604,6 @@ static bool __sbq_wake_up(struct sbitmap_queue *sbq, int *nr)
>>>  	struct sbq_wait_state *ws;
>>>  	unsigned int wake_batch;
>>>  	int wait_cnt, cur, sub;
>>> -	bool ret;
>>>  
>>>  	if (*nr <= 0)
>>>  		return false;
>>> @@ -632,15 +631,15 @@ static bool __sbq_wake_up(struct sbitmap_queue *sbq, int *nr)
>>>  	if (wait_cnt > 0)
>>>  		return !waitqueue_active(&ws->wait);
>>>  
>>> -	*nr -= sub;
>>> -
>>>  	/*
>>>  	 * When wait_cnt == 0, we have to be particularly careful as we are
>>>  	 * responsible to reset wait_cnt regardless whether we've actually
>>> -	 * woken up anybody. But in case we didn't wakeup anybody, we still
>>> -	 * need to retry.
>>> +	 * woken up anybody. But in case we didn't wakeup anybody, we should
>>> +	 * not consume nr and need to retry to avoid lost wakeups.
>>>  	 */
>>> -	ret = !waitqueue_active(&ws->wait);
>> There is a warnning reported by checkpatch.pl which is
>> "WARNING:waitqueue_active without comment" but I don't know why.
> 
> Most likely because waitqueue_active() could be racy, so a comment is
> warranted on why it's safe rather than using wq_has_sleeper().
Thanks for explanation, so the patch seems fine as comment is present
already though it doesn't mention sting "waitqueue_active" directly.
No bother anymore, this patch will be dropped as the fixed code is
stale.
Thanks again.

-- 
Best wishes
Kemeng Shi

Re: [PATCH 1/5] sbitmap: don't consume nr for inactive waitqueue to avoid lost wakeups

[Gabriel Krisman Bertazi]

Kemeng Shi <shikemeng@huawei.com> writes:

> If we decremented queue without waiters, we should not decremente freed
> bits number "nr", or all "nr" could be consumed in a empty queue and no
> wakeup will be called.
> Currently, for case "wait_cnt > 0", "nr" will not be decremented if we
> decremented queue without watiers and retry is returned to avoid lost
> wakeups. However for case "wait_cnt == 0", "nr" will be decremented
> unconditionally and maybe decremented to zero. Although retry is
> returned by active state of queue, it's not actually executed for "nr"
> is zero.
>

Hi Kemeng,

Fwiw, I sent a patch rewriting this algorithm which is now merged in
axboe/for-next.  It drops the per-waitqueue wait_cnt entirely.  You can
find it here:

https://lore.kernel.org/lkml/20221110153533.go5qs3psm75h27mx@quack3/T/

Thanks!


-- 
Gabriel Krisman Bertazi

Re: [PATCH 1/5] sbitmap: don't consume nr for inactive waitqueue to avoid lost wakeups

[Kemeng Shi]

on 12/1/2022 9:32 PM, Gabriel Krisman Bertazi wrote:
> Kemeng Shi <shikemeng@huawei.com> writes:
> 
>> If we decremented queue without waiters, we should not decremente freed
>> bits number "nr", or all "nr" could be consumed in a empty queue and no
>> wakeup will be called.
>> Currently, for case "wait_cnt > 0", "nr" will not be decremented if we
>> decremented queue without watiers and retry is returned to avoid lost
>> wakeups. However for case "wait_cnt == 0", "nr" will be decremented
>> unconditionally and maybe decremented to zero. Although retry is
>> returned by active state of queue, it's not actually executed for "nr"
>> is zero.
>>
> 
> Hi Kemeng,
> 
> Fwiw, I sent a patch rewriting this algorithm which is now merged in
> axboe/for-next.  It drops the per-waitqueue wait_cnt entirely.  You can
> find it here:
> 
> https://lore.kernel.org/lkml/20221110153533.go5qs3psm75h27mx@quack3/T/
> 
> Thanks!
Hi Gabriel,
Thanks for remind me of this, I will recheck my patches in the
axboe/for-next branch.

-- 
Best wishes
Kemeng Shi

[PATCH 2/5] sbitmap: remove unnecessary calculation of alloc_hint in __sbitmap_get_shallow

[Kemeng Shi]

We calculate allow_hint in next word as following:

/* low sb->shift bit of alloc_hint will be 0 after this shift */
alloc_hint = index << sb->shift;

/* get low sb->shift bit of alloc_hit */
SB_NR_TO_BIT(sb, alloc_hint)

So alloc_hit in next word will always be zero. Simpfy alloc_hit calculation
in __sbitmap_get_shallow according to the alloc_hit calculation in
__sbitmap_get.

Signed-off-by: Kemeng Shi <shikemeng@huawei.com>
---
 lib/sbitmap.c | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/lib/sbitmap.c b/lib/sbitmap.c
index e40759bcf821..791edf5ea4ca 100644
--- a/lib/sbitmap.c
+++ b/lib/sbitmap.c
@@ -243,6 +243,7 @@ static int __sbitmap_get_shallow(struct sbitmap *sb,
 	int nr = -1;
 
 	index = SB_NR_TO_INDEX(sb, alloc_hint);
+	alloc_hint = SB_NR_TO_BIT(sb, alloc_hint);
 
 	for (i = 0; i < sb->map_nr; i++) {
 again:
@@ -250,7 +251,7 @@ static int __sbitmap_get_shallow(struct sbitmap *sb,
 					min_t(unsigned int,
 					      __map_depth(sb, index),
 					      shallow_depth),
-					SB_NR_TO_BIT(sb, alloc_hint), true);
+					alloc_hint, true);
 		if (nr != -1) {
 			nr += index << sb->shift;
 			break;
@@ -260,13 +261,9 @@ static int __sbitmap_get_shallow(struct sbitmap *sb,
 			goto again;
 
 		/* Jump to next index. */
-		index++;
-		alloc_hint = index << sb->shift;
-
-		if (index >= sb->map_nr) {
+		alloc_hint = 0;
+		if (++index >= sb->map_nr)
 			index = 0;
-			alloc_hint = 0;
-		}
 	}
 
 	return nr;
-- 
2.30.0

[PATCH 3/5] sbitmap: remove redundant check in __sbitmap_queue_get_batch

[Kemeng Shi]

Commit fbb564a557809 ("lib/sbitmap: Fix invalid loop in
__sbitmap_queue_get_batch()") mentioned that "Checking free bits when
setting the target bits. Otherwise, it may reuse the busying bits."
This commit add check to make sure all masked bits in word before
cmpxchg is zero. Then the existing check after cmpxchg to check any
zero bit is existing in masked bits in word is redundant.

Actually, old value of word before cmpxchg is stored in val and we
will filter out busy bits in val by "(get_mask & ~val)" after cmpxchg.
So we will not reuse busy bits methioned in commit fbb564a557809
("lib/sbitmap: Fix invalid loop in __sbitmap_queue_get_batch()"). Revert
new-added check to remove redundant check.

Signed-off-by: Kemeng Shi <shikemeng@huawei.com>
---
 lib/sbitmap.c | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/lib/sbitmap.c b/lib/sbitmap.c
index 791edf5ea4ca..70b62b1e40a1 100644
--- a/lib/sbitmap.c
+++ b/lib/sbitmap.c
@@ -534,11 +534,9 @@ unsigned long __sbitmap_queue_get_batch(struct sbitmap_queue *sbq, int nr_tags,
 
 			get_mask = ((1UL << nr_tags) - 1) << nr;
 			val = READ_ONCE(map->word);
-			do {
-				if ((val & ~get_mask) != val)
-					goto next;
-			} while (!atomic_long_try_cmpxchg(ptr, &val,
-							  get_mask | val));
+			while (!atomic_long_try_cmpxchg(ptr, &val,
+							  get_mask | val))
+				;
 			get_mask = (get_mask & ~val) >> nr;
 			if (get_mask) {
 				*offset = nr + (index << sb->shift);
-- 
2.30.0

[PATCH 4/5] sbitmap: rewrite sbitmap_find_bit_in_index to reduce repeat code

[Kemeng Shi]

Rewrite sbitmap_find_bit_in_index as following:
1. Rename sbitmap_find_bit_in_index to sbitmap_find_bit_in_word
2. Accept "struct sbitmap_word *" directly instead of accepting
"struct sbitmap *" and "int index" to get "struct sbitmap_word *".
3. Accept depth/shallow_depth and wrap for __sbitmap_get_word from caller
to support need of both __sbitmap_get_shallow and __sbitmap_get.

With helper function sbitmap_find_bit_in_word, we can remove repeat
code in __sbitmap_get_shallow to find bit considring deferred clear.

Signed-off-by: Kemeng Shi <shikemeng@huawei.com>
---
 lib/sbitmap.c | 29 ++++++++++++++---------------
 1 file changed, 14 insertions(+), 15 deletions(-)

diff --git a/lib/sbitmap.c b/lib/sbitmap.c
index 70b62b1e40a1..b6a2cdb9bdaf 100644
--- a/lib/sbitmap.c
+++ b/lib/sbitmap.c
@@ -167,15 +167,16 @@ static int __sbitmap_get_word(unsigned long *word, unsigned long depth,
 	return nr;
 }
 
-static int sbitmap_find_bit_in_index(struct sbitmap *sb, int index,
-				     unsigned int alloc_hint)
+static int sbitmap_find_bit_in_word(struct sbitmap_word *map,
+				    unsigned int depth,
+				    unsigned int alloc_hint,
+				    bool wrap)
 {
-	struct sbitmap_word *map = &sb->map[index];
 	int nr;
 
 	do {
-		nr = __sbitmap_get_word(&map->word, __map_depth(sb, index),
-					alloc_hint, !sb->round_robin);
+		nr = __sbitmap_get_word(&map->word, depth,
+					alloc_hint, wrap);
 		if (nr != -1)
 			break;
 		if (!sbitmap_deferred_clear(map))
@@ -203,7 +204,8 @@ static int __sbitmap_get(struct sbitmap *sb, unsigned int alloc_hint)
 		alloc_hint = 0;
 
 	for (i = 0; i < sb->map_nr; i++) {
-		nr = sbitmap_find_bit_in_index(sb, index, alloc_hint);
+		nr = sbitmap_find_bit_in_word(&sb->map[index], __map_depth(sb, index),
+					      alloc_hint, !sb->round_robin);
 		if (nr != -1) {
 			nr += index << sb->shift;
 			break;
@@ -246,20 +248,17 @@ static int __sbitmap_get_shallow(struct sbitmap *sb,
 	alloc_hint = SB_NR_TO_BIT(sb, alloc_hint);
 
 	for (i = 0; i < sb->map_nr; i++) {
-again:
-		nr = __sbitmap_get_word(&sb->map[index].word,
-					min_t(unsigned int,
-					      __map_depth(sb, index),
-					      shallow_depth),
-					alloc_hint, true);
+		nr = sbitmap_find_bit_in_word(&sb->map[index],
+					      min_t(unsigned int,
+						    __map_depth(sb, index),
+						    shallow_depth),
+					      alloc_hint, true);
+
 		if (nr != -1) {
 			nr += index << sb->shift;
 			break;
 		}
 
-		if (sbitmap_deferred_clear(&sb->map[index]))
-			goto again;
-
 		/* Jump to next index. */
 		alloc_hint = 0;
 		if (++index >= sb->map_nr)
-- 
2.30.0

[PATCH 5/5] sbitmap: add sbitmap_find_bit to remove repeat code in __sbitmap_get/__sbitmap_get_shallow

[Kemeng Shi]

There are three differences between __sbitmap_get and
__sbitmap_get_shallow when searching free bit:
1. __sbitmap_get_shallow limit number of bit to search per word.
__sbitmap_get has no such limit.
2. __sbitmap_get_shallow always searches with wrap set. __sbitmap_get set
wrap according to round_robin.
3. __sbitmap_get_shallow always searches from first bit in first word.
__sbitmap_get searches from first bit when round_robin is not set
otherwise searches from SB_NR_TO_BIT(sb, alloc_hint).

Add helper function sbitmap_find_bit function to do common search while
accept "limit depth per word", "wrap flag" and "first bit to
search" from caller to support the need of both __sbitmap_get and
__sbitmap_get_shallow.

Signed-off-by: Kemeng Shi <shikemeng@huawei.com>
---
 lib/sbitmap.c | 72 +++++++++++++++++++++++++--------------------------
 1 file changed, 35 insertions(+), 37 deletions(-)

diff --git a/lib/sbitmap.c b/lib/sbitmap.c
index b6a2cdb9bdaf..93e7db484b96 100644
--- a/lib/sbitmap.c
+++ b/lib/sbitmap.c
@@ -186,26 +186,22 @@ static int sbitmap_find_bit_in_word(struct sbitmap_word *map,
 	return nr;
 }
 
-static int __sbitmap_get(struct sbitmap *sb, unsigned int alloc_hint)
+static int sbitmap_find_bit(struct sbitmap *sb,
+			    unsigned int depth,
+			    unsigned int index,
+			    unsigned int alloc_hint,
+			    bool wrap)
 {
-	unsigned int i, index;
+	unsigned int i;
 	int nr = -1;
 
-	index = SB_NR_TO_INDEX(sb, alloc_hint);
-
-	/*
-	 * Unless we're doing round robin tag allocation, just use the
-	 * alloc_hint to find the right word index. No point in looping
-	 * twice in find_next_zero_bit() for that case.
-	 */
-	if (sb->round_robin)
-		alloc_hint = SB_NR_TO_BIT(sb, alloc_hint);
-	else
-		alloc_hint = 0;
-
 	for (i = 0; i < sb->map_nr; i++) {
-		nr = sbitmap_find_bit_in_word(&sb->map[index], __map_depth(sb, index),
-					      alloc_hint, !sb->round_robin);
+		nr = sbitmap_find_bit_in_word(&sb->map[index],
+					      min_t(unsigned int,
+						    __map_depth(sb, index),
+						    depth),
+					      alloc_hint, wrap);
+
 		if (nr != -1) {
 			nr += index << sb->shift;
 			break;
@@ -215,11 +211,32 @@ static int __sbitmap_get(struct sbitmap *sb, unsigned int alloc_hint)
 		alloc_hint = 0;
 		if (++index >= sb->map_nr)
 			index = 0;
+
 	}
 
 	return nr;
 }
 
+static int __sbitmap_get(struct sbitmap *sb, unsigned int alloc_hint)
+{
+	unsigned int index;
+
+	index = SB_NR_TO_INDEX(sb, alloc_hint);
+
+	/*
+	 * Unless we're doing round robin tag allocation, just use the
+	 * alloc_hint to find the right word index. No point in looping
+	 * twice in find_next_zero_bit() for that case.
+	 */
+	if (sb->round_robin)
+		alloc_hint = SB_NR_TO_BIT(sb, alloc_hint);
+	else
+		alloc_hint = 0;
+
+	return sbitmap_find_bit(sb, UINT_MAX, index, alloc_hint,
+				!sb->round_robin);
+}
+
 int sbitmap_get(struct sbitmap *sb)
 {
 	int nr;
@@ -241,31 +258,12 @@ static int __sbitmap_get_shallow(struct sbitmap *sb,
 				 unsigned int alloc_hint,
 				 unsigned long shallow_depth)
 {
-	unsigned int i, index;
-	int nr = -1;
+	unsigned int index;
 
 	index = SB_NR_TO_INDEX(sb, alloc_hint);
 	alloc_hint = SB_NR_TO_BIT(sb, alloc_hint);
 
-	for (i = 0; i < sb->map_nr; i++) {
-		nr = sbitmap_find_bit_in_word(&sb->map[index],
-					      min_t(unsigned int,
-						    __map_depth(sb, index),
-						    shallow_depth),
-					      alloc_hint, true);
-
-		if (nr != -1) {
-			nr += index << sb->shift;
-			break;
-		}
-
-		/* Jump to next index. */
-		alloc_hint = 0;
-		if (++index >= sb->map_nr)
-			index = 0;
-	}
-
-	return nr;
+	return sbitmap_find_bit(sb, shallow_depth, index, alloc_hint, true);
 }
 
 int sbitmap_get_shallow(struct sbitmap *sb, unsigned long shallow_depth)
-- 
2.30.0