[RFC PATCH 1/6] lib/sbitmap: add ranged allocation, bounded wakeup relay, and ranged weight

[Keith Busch <kbusch@meta.com>]

From: Keith Busch <kbusch@kernel.org>

Add some helpers that let a caller carve a shared bitmap into per-user
windows.

  * sbitmap_queue_get_range() allocates a free bit only within requested range.
  * sbitmap_queue_wake_up_relay() hands a wakeup to another waiter when the
    woken one cannot use the freed bit (it fell outside that waiter's window).
  * sbitmap_weight_range() counts set (and not cleared) bits in a sub-range.

No callers yet. The intention is for blk-mq users for shared tag
fairness without requiring per-io atomic accounting operations.

Signed-off-by: Keith Busch <kbusch@kernel.org>
---
 include/linux/sbitmap.h |  53 +++++++++++
 lib/sbitmap.c           | 200 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 253 insertions(+)

diff --git a/include/linux/sbitmap.h b/include/linux/sbitmap.h
index cc7ad189caa5c..9fceda5d9de51 100644
--- a/include/linux/sbitmap.h
+++ b/include/linux/sbitmap.h
@@ -86,6 +86,7 @@ struct sbitmap {
 
 #define SBQ_WAIT_QUEUES 8
 #define SBQ_WAKE_BATCH 8
+#define SBQ_RELAY_BUDGET SBQ_WAKE_BATCH
 
 /**
  * struct sbq_wait_state - Wait queue in a &struct sbitmap_queue.
@@ -149,6 +150,14 @@ struct sbitmap_queue {
 	 * @wakeup_cnt: Number of thread wake ups issued.
 	 */
 	atomic_t wakeup_cnt;
+
+	/**
+	 * @relay_credits: Remaining range-relay wakeups allowed before the next
+	 * genuine completion. Refilled by sbitmap_queue_wake_up() and consumed
+	 * by sbitmap_queue_wake_up_relay(); bounds the relay chain so it cannot
+	 * cycle when no waiter can use a freed bit.
+	 */
+	atomic_t relay_credits;
 };
 
 /**
@@ -375,6 +384,18 @@ void sbitmap_show(struct sbitmap *sb, struct seq_file *m);
  */
 unsigned int sbitmap_weight(const struct sbitmap *sb);
 
+/**
+ * sbitmap_weight_range() - Return how many set and not cleared bits in a
+ * sub-range of a &struct sbitmap.
+ * @sb: Bitmap to check.
+ * @off: First bit of the range.
+ * @len: Number of bits in the range.
+ *
+ * Return: How many bits in [@off, @off + @len) are set and not cleared.
+ */
+unsigned int sbitmap_weight_range(const struct sbitmap *sb, unsigned int off,
+				  unsigned int len);
+
 /**
  * sbitmap_bitmap_show() - Write a hex dump of a &struct sbitmap to a &struct
  * seq_file.
@@ -472,6 +493,22 @@ unsigned long __sbitmap_queue_get_batch(struct sbitmap_queue *sbq, int nr_tags,
 int sbitmap_queue_get_shallow(struct sbitmap_queue *sbq,
 			      unsigned int shallow_depth);
 
+/**
+ * sbitmap_queue_get_range() - Try to allocate a free bit from a restricted
+ * range of a &struct sbitmap_queue, with preemption already disabled.
+ * @sbq: Bitmap queue to allocate from.
+ * @min: First bit (inclusive) the caller is allowed to allocate.
+ * @max: Last bit (exclusive) the caller is allowed to allocate.
+ *
+ * Allocation is confined to the [@min, @max) window. This lets several users
+ * share one bitmap while each is bounded to its own (possibly overlapping)
+ * window, providing fairness without a separate per-user accounting counter.
+ *
+ * Return: Non-negative allocated bit number if successful, -1 otherwise.
+ */
+int sbitmap_queue_get_range(struct sbitmap_queue *sbq, unsigned int min,
+			    unsigned int max);
+
 /**
  * sbitmap_queue_get() - Try to allocate a free bit from a &struct
  * sbitmap_queue.
@@ -575,6 +612,22 @@ void sbitmap_queue_wake_all(struct sbitmap_queue *sbq);
  */
 void sbitmap_queue_wake_up(struct sbitmap_queue *sbq, int nr);
 
+/**
+ * sbitmap_queue_wake_up_relay() - Relay a wakeup to the next waiter.
+ * @sbq: Bitmap queue to wake up.
+ * @cur: The relaying waiter's own wait queue, skipped so it does not wake
+ *       itself (it is about to sleep again); may be NULL.
+ *
+ * Wake a single waiter on the next active wait queue without touching the
+ * wake_batch accounting. Intended for range-limited allocation: a waiter that
+ * was woken but could not use the freed bit (it fell outside the waiter's
+ * allowed range) calls this to pass the wakeup along to another waiter that
+ * may be able to use it. The relay is bounded by a credit budget refilled
+ * only by genuine completions, so it cannot cycle indefinitely.
+ */
+void sbitmap_queue_wake_up_relay(struct sbitmap_queue *sbq,
+				 struct sbq_wait_state *cur);
+
 /**
  * sbitmap_queue_show() - Dump &struct sbitmap_queue information to a &struct
  * seq_file.
diff --git a/lib/sbitmap.c b/lib/sbitmap.c
index 4d188d05db153..3264085b4d533 100644
--- a/lib/sbitmap.c
+++ b/lib/sbitmap.c
@@ -338,6 +338,130 @@ static int sbitmap_get_shallow(struct sbitmap *sb, unsigned long shallow_depth)
 	return nr;
 }
 
+/*
+ * Find and set a free bit in the [low, high) sub-range of a single word.
+ * Wraps back to @low once so a free bit below @hint is still found.
+ */
+static int __sbitmap_get_word_range(unsigned long *word, unsigned int low,
+				    unsigned int high, unsigned int hint)
+{
+	bool wrap;
+	int nr;
+
+	if (hint < low || hint >= high)
+		hint = low;
+	/* don't wrap if already starting from the bottom of the range */
+	wrap = hint > low;
+
+	while (1) {
+		nr = find_next_zero_bit(word, high, hint);
+		if (unlikely(nr >= high)) {
+			if (wrap) {
+				hint = low;
+				wrap = false;
+				continue;
+			}
+			return -1;
+		}
+		if (!test_and_set_bit_lock(nr, word))
+			break;
+		hint = nr + 1;
+		if (hint >= high) {
+			if (!wrap)
+				return -1;
+			hint = low;
+			wrap = false;
+		}
+	}
+
+	return nr;
+}
+
+static int sbitmap_find_bit_in_word_range(struct sbitmap_word *map,
+					  unsigned int low, unsigned int high,
+					  unsigned int hint)
+{
+	int nr;
+
+	do {
+		nr = __sbitmap_get_word_range(&map->word, low, high, hint);
+		if (nr != -1)
+			break;
+		if (!sbitmap_deferred_clear(map, high, low, false))
+			break;
+	} while (1);
+
+	return nr;
+}
+
+/*
+ * Allocate a bit restricted to the [min, max) range of the bitmap. Used to
+ * carve the shared depth into (possibly overlapping) per-user windows without
+ * a separate accounting counter: a user simply cannot allocate outside its
+ * window.
+ */
+static int __sbitmap_get_range(struct sbitmap *sb, unsigned int min,
+			       unsigned int max, unsigned int hint)
+{
+	unsigned int min_index, max_index, index;
+	unsigned int nwords, i;
+	int nr = -1;
+
+	if (max > sb->depth)
+		max = sb->depth;
+	if (min >= max)
+		return -1;
+
+	min_index = SB_NR_TO_INDEX(sb, min);
+	max_index = SB_NR_TO_INDEX(sb, max - 1);
+	nwords = max_index - min_index + 1;
+
+	if (hint < min || hint >= max)
+		hint = min;
+	index = SB_NR_TO_INDEX(sb, hint);
+
+	for (i = 0; i < nwords; i++) {
+		unsigned int word_base = index << sb->shift;
+		unsigned int word_depth = __map_depth(sb, index);
+		unsigned int low, high, h;
+
+		low = (index == min_index) ? SB_NR_TO_BIT(sb, min) : 0;
+		high = (index == max_index) ? SB_NR_TO_BIT(sb, max - 1) + 1 :
+					      word_depth;
+		h = (hint >= word_base && hint < word_base + word_depth) ?
+			SB_NR_TO_BIT(sb, hint) : low;
+
+		nr = sbitmap_find_bit_in_word_range(&sb->map[index], low,
+						    high, h);
+		if (nr != -1) {
+			nr += word_base;
+			break;
+		}
+
+		if (++index > max_index)
+			index = min_index;
+	}
+
+	return nr;
+}
+
+static int sbitmap_get_range(struct sbitmap *sb, unsigned int min,
+			     unsigned int max)
+{
+	int nr;
+	unsigned int hint, depth;
+
+	if (WARN_ON_ONCE(unlikely(!sb->alloc_hint)))
+		return -1;
+
+	depth = READ_ONCE(sb->depth);
+	hint = update_alloc_hint_before_get(sb, depth);
+	nr = __sbitmap_get_range(sb, min, max, hint);
+	update_alloc_hint_after_get(sb, depth, hint, nr);
+
+	return nr;
+}
+
 bool sbitmap_any_bit_set(const struct sbitmap *sb)
 {
 	unsigned int i;
@@ -377,6 +501,38 @@ unsigned int sbitmap_weight(const struct sbitmap *sb)
 }
 EXPORT_SYMBOL_GPL(sbitmap_weight);
 
+unsigned int sbitmap_weight_range(const struct sbitmap *sb, unsigned int off,
+				  unsigned int len)
+{
+	unsigned int weight = 0;
+	unsigned int end = off + len;
+	unsigned int index;
+
+	if (end > sb->depth)
+		end = sb->depth;
+
+	for (index = SB_NR_TO_INDEX(sb, off); index < sb->map_nr; index++) {
+		const struct sbitmap_word *word = &sb->map[index];
+		unsigned int base = index << sb->shift;
+		unsigned int wbits = __map_depth(sb, index);
+		unsigned int lo, hi;
+		unsigned long val;
+
+		if (base >= end)
+			break;
+		val = READ_ONCE(word->word) & ~READ_ONCE(word->cleared);
+		lo = (off > base) ? off - base : 0;
+		hi = (end - base < wbits) ? end - base : wbits;
+		if (lo)
+			val &= ~((1UL << lo) - 1);
+		if (hi < BITS_PER_LONG)
+			val &= (1UL << hi) - 1;
+		weight += hweight_long(val);
+	}
+	return weight;
+}
+EXPORT_SYMBOL_GPL(sbitmap_weight_range);
+
 void sbitmap_show(struct sbitmap *sb, struct seq_file *m)
 {
 	seq_printf(m, "depth=%u\n", sb->depth);
@@ -462,6 +618,7 @@ int sbitmap_queue_init_node(struct sbitmap_queue *sbq, unsigned int depth,
 	atomic_set(&sbq->ws_active, 0);
 	atomic_set(&sbq->completion_cnt, 0);
 	atomic_set(&sbq->wakeup_cnt, 0);
+	atomic_set(&sbq->relay_credits, 0);
 
 	sbq->ws = kzalloc_node(SBQ_WAIT_QUEUES * sizeof(*sbq->ws), flags, node);
 	if (!sbq->ws) {
@@ -573,6 +730,13 @@ int sbitmap_queue_get_shallow(struct sbitmap_queue *sbq,
 }
 EXPORT_SYMBOL_GPL(sbitmap_queue_get_shallow);
 
+int sbitmap_queue_get_range(struct sbitmap_queue *sbq, unsigned int min,
+			    unsigned int max)
+{
+	return sbitmap_get_range(&sbq->sb, min, max);
+}
+EXPORT_SYMBOL_GPL(sbitmap_queue_get_range);
+
 void sbitmap_queue_min_shallow_depth(struct sbitmap_queue *sbq,
 				     unsigned int min_shallow_depth)
 {
@@ -629,10 +793,46 @@ void sbitmap_queue_wake_up(struct sbitmap_queue *sbq, int nr)
 	} while (!atomic_try_cmpxchg(&sbq->wakeup_cnt,
 				     &wakeups, wakeups + wake_batch));
 
+	atomic_set(&sbq->relay_credits, SBQ_RELAY_BUDGET);
 	__sbitmap_queue_wake_up(sbq, wake_batch);
 }
 EXPORT_SYMBOL_GPL(sbitmap_queue_wake_up);
 
+/*
+ * Wake a single waiter on the next active wait queue, bypassing the
+ * wake_batch accounting (no bit was freed). This "relays" a wakeup that the
+ * current waiter consumed but could not use - e.g. the freed bit fell outside
+ * its allowed allocation range - so a waiter that *can* use a free bit gets a
+ * chance to run. Relaying continues, one hop per failed waiter, until a usable
+ * waiter is found or the chain credits drain.
+ */
+void sbitmap_queue_wake_up_relay(struct sbitmap_queue *sbq,
+				 struct sbq_wait_state *cur)
+{
+	int wake_index, i;
+
+	if (!atomic_read(&sbq->ws_active))
+		return;
+
+	if (atomic_dec_if_positive(&sbq->relay_credits) < 0)
+		return;
+
+	wake_index = atomic_read(&sbq->wake_index);
+	for (i = 0; i < SBQ_WAIT_QUEUES; i++) {
+		struct sbq_wait_state *ws = &sbq->ws[wake_index];
+
+		wake_index = sbq_index_inc(wake_index);
+		if (ws != cur && waitqueue_active(&ws->wait)) {
+			wake_up_nr(&ws->wait, 1);
+			break;
+		}
+	}
+
+	if (wake_index != atomic_read(&sbq->wake_index))
+		atomic_set(&sbq->wake_index, wake_index);
+}
+EXPORT_SYMBOL_GPL(sbitmap_queue_wake_up_relay);
+
 static inline void sbitmap_update_cpu_hint(struct sbitmap *sb, int cpu, int tag)
 {
 	if (likely(!sb->round_robin && tag < sb->depth))
-- 
2.52.0