dm-bufio: avoid redundant buffer_tree lookups

dm-bufio's map from block number to buffer is organized as a hash table
of red-black trees.  It does far more lookups in this hash table than
necessary: typically one lookup to lock the tree, one lookup to search
the tree, and one lookup to unlock the tree.  Only one of those lookups
is needed.  Optimize it to do only the minimum number of lookups.

This improves performance.   It also reduces the object code size,
considering that the redundant hash table lookups were being inlined.
For example, the size of the text section of dm-bufio.o decreases from
15599 to 15070 bytes with gcc 15 and x86_64, or from 20652 to 20244
bytes with clang 21 and arm64.

Signed-off-by: Eric Biggers <ebiggers@kernel.org>
Signed-off-by: Mikulas Patocka <mpatocka@redhat.com>

Author: Eric Biggers

drivers/md/dm-bufio.c

@@ -401,36 +401,51 @@ static inline unsigned int cache_index(sector_t block, unsigned int num_locks)
 	return dm_hash_locks_index(block, num_locks);
 }
 
-static inline void cache_read_lock(struct dm_buffer_cache *bc, sector_t block)
+/* Get the buffer tree in the cache for the given block.  Doesn't lock it. */
+static inline struct buffer_tree *cache_get_tree(struct dm_buffer_cache *bc,
+						 sector_t block)
+{
+	return &bc->trees[cache_index(block, bc->num_locks)];
+}
+
+/* Lock the given buffer tree in the cache for reading. */
+static inline void cache_read_lock(struct dm_buffer_cache *bc,
+				   struct buffer_tree *tree)
 {
 	if (static_branch_unlikely(&no_sleep_enabled) && bc->no_sleep)
-		read_lock_bh(&bc->trees[cache_index(block, bc->num_locks)].u.spinlock);
+		read_lock_bh(&tree->u.spinlock);
 	else
-		down_read(&bc->trees[cache_index(block, bc->num_locks)].u.lock);
+		down_read(&tree->u.lock);
 }
 
-static inline void cache_read_unlock(struct dm_buffer_cache *bc, sector_t block)
+/* Unlock the given buffer tree in the cache for reading. */
+static inline void cache_read_unlock(struct dm_buffer_cache *bc,
+				     struct buffer_tree *tree)
 {
 	if (static_branch_unlikely(&no_sleep_enabled) && bc->no_sleep)
-		read_unlock_bh(&bc->trees[cache_index(block, bc->num_locks)].u.spinlock);
+		read_unlock_bh(&tree->u.spinlock);
 	else
-		up_read(&bc->trees[cache_index(block, bc->num_locks)].u.lock);
+		up_read(&tree->u.lock);
 }
 
-static inline void cache_write_lock(struct dm_buffer_cache *bc, sector_t block)
+/* Lock the given buffer tree in the cache for writing. */
+static inline void cache_write_lock(struct dm_buffer_cache *bc,
+				    struct buffer_tree *tree)
 {
 	if (static_branch_unlikely(&no_sleep_enabled) && bc->no_sleep)
-		write_lock_bh(&bc->trees[cache_index(block, bc->num_locks)].u.spinlock);
+		write_lock_bh(&tree->u.spinlock);
 	else
-		down_write(&bc->trees[cache_index(block, bc->num_locks)].u.lock);
+		down_write(&tree->u.lock);
 }
 
-static inline void cache_write_unlock(struct dm_buffer_cache *bc, sector_t block)
+/* Unlock the given buffer tree in the cache for writing. */
+static inline void cache_write_unlock(struct dm_buffer_cache *bc,
+				      struct buffer_tree *tree)
 {
 	if (static_branch_unlikely(&no_sleep_enabled) && bc->no_sleep)
-		write_unlock_bh(&bc->trees[cache_index(block, bc->num_locks)].u.spinlock);
+		write_unlock_bh(&tree->u.spinlock);
 	else
-		up_write(&bc->trees[cache_index(block, bc->num_locks)].u.lock);
+		up_write(&tree->u.lock);
 }
 
 /*
@@ -602,17 +617,19 @@ static void __cache_inc_buffer(struct dm_buffer *b)
 	WRITE_ONCE(b->last_accessed, jiffies);
 }
 
-static struct dm_buffer *cache_get(struct dm_buffer_cache *bc, sector_t block)
+static struct dm_buffer *cache_get(struct dm_buffer_cache *bc,
+				   struct buffer_tree *tree, sector_t block)
 {
 	struct dm_buffer *b;
 
-	cache_read_lock(bc, block);
-	b = __cache_get(&bc->trees[cache_index(block, bc->num_locks)].root, block);
+	/* Assuming tree == cache_get_tree(bc, block) */
+	cache_read_lock(bc, tree);
+	b = __cache_get(&tree->root, block);
 	if (b) {
 		lru_reference(&b->lru);
 		__cache_inc_buffer(b);
 	}
-	cache_read_unlock(bc, block);
+	cache_read_unlock(bc, tree);
 
 	return b;
 }
@@ -663,7 +680,7 @@ static struct dm_buffer *__cache_evict(struct dm_buffer_cache *bc, int list_mode
 
 	b = le_to_buffer(le);
 	/* __evict_pred will have locked the appropriate tree. */
-	rb_erase(&b->node, &bc->trees[cache_index(b->block, bc->num_locks)].root);
+	rb_erase(&b->node, &cache_get_tree(bc, b->block)->root);
 
 	return b;
 }
@@ -686,15 +703,17 @@ static struct dm_buffer *cache_evict(struct dm_buffer_cache *bc, int list_mode,
 /*
  * Mark a buffer as clean or dirty. Not threadsafe.
  */
-static void cache_mark(struct dm_buffer_cache *bc, struct dm_buffer *b, int list_mode)
+static void cache_mark(struct dm_buffer_cache *bc, struct buffer_tree *tree,
+		       struct dm_buffer *b, int list_mode)
 {
-	cache_write_lock(bc, b->block);
+	/* Assuming tree == cache_get_tree(bc, b->block) */
+	cache_write_lock(bc, tree);
 	if (list_mode != b->list_mode) {
 		lru_remove(&bc->lru[b->list_mode], &b->lru);
 		b->list_mode = list_mode;
 		lru_insert(&bc->lru[b->list_mode], &b->lru);
 	}
-	cache_write_unlock(bc, b->block);
+	cache_write_unlock(bc, tree);
 }
 
 /*--------------*/
@@ -820,19 +839,21 @@ static bool __cache_insert(struct rb_root *root, struct dm_buffer *b)
 	return true;
 }
 
-static bool cache_insert(struct dm_buffer_cache *bc, struct dm_buffer *b)
+static bool cache_insert(struct dm_buffer_cache *bc, struct buffer_tree *tree,
+			 struct dm_buffer *b)
 {
 	bool r;
 
 	if (WARN_ON_ONCE(b->list_mode >= LIST_SIZE))
 		return false;
 
-	cache_write_lock(bc, b->block);
+	/* Assuming tree == cache_get_tree(bc, b->block) */
+	cache_write_lock(bc, tree);
 	BUG_ON(atomic_read(&b->hold_count) != 1);
-	r = __cache_insert(&bc->trees[cache_index(b->block, bc->num_locks)].root, b);
+	r = __cache_insert(&tree->root, b);
 	if (r)
 		lru_insert(&bc->lru[b->list_mode], &b->lru);
-	cache_write_unlock(bc, b->block);
+	cache_write_unlock(bc, tree);
 
 	return r;
 }
@@ -845,21 +866,23 @@ static bool cache_insert(struct dm_buffer_cache *bc, struct dm_buffer *b)
  *
  * Not threadsafe.
  */
-static bool cache_remove(struct dm_buffer_cache *bc, struct dm_buffer *b)
+static bool cache_remove(struct dm_buffer_cache *bc, struct buffer_tree *tree,
+			 struct dm_buffer *b)
 {
 	bool r;
 
-	cache_write_lock(bc, b->block);
+	/* Assuming tree == cache_get_tree(bc, b->block) */
+	cache_write_lock(bc, tree);
 
 	if (atomic_read(&b->hold_count) != 1) {
 		r = false;
 	} else {
 		r = true;
-		rb_erase(&b->node, &bc->trees[cache_index(b->block, bc->num_locks)].root);
+		rb_erase(&b->node, &tree->root);
 		lru_remove(&bc->lru[b->list_mode], &b->lru);
 	}
 
-	cache_write_unlock(bc, b->block);
+	cache_write_unlock(bc, tree);
 
 	return r;
 }
@@ -1725,14 +1748,16 @@ static void __check_watermark(struct dm_bufio_client *c,
  *--------------------------------------------------------------
  */
 
-static void cache_put_and_wake(struct dm_bufio_client *c, struct dm_buffer *b)
+static void cache_put_and_wake(struct dm_bufio_client *c,
+			       struct buffer_tree *tree, struct dm_buffer *b)
 {
 	bool wake;
 
-	cache_read_lock(&c->cache, b->block);
+	/* Assuming tree == cache_get_tree(&c->cache, b->block) */
+	cache_read_lock(&c->cache, tree);
 	BUG_ON(!atomic_read(&b->hold_count));
 	wake = atomic_dec_and_test(&b->hold_count);
-	cache_read_unlock(&c->cache, b->block);
+	cache_read_unlock(&c->cache, tree);
 
 	/*
 	 * Relying on waitqueue_active() is racey, but we sleep
@@ -1746,7 +1771,8 @@ static void cache_put_and_wake(struct dm_bufio_client *c, struct dm_buffer *b)
  * This assumes you have already checked the cache to see if the buffer
  * is already present (it will recheck after dropping the lock for allocation).
  */
-static struct dm_buffer *__bufio_new(struct dm_bufio_client *c, sector_t block,
+static struct dm_buffer *__bufio_new(struct dm_bufio_client *c,
+				     struct buffer_tree *tree, sector_t block,
 				     enum new_flag nf, int *need_submit,
 				     struct list_head *write_list)
 {
@@ -1766,7 +1792,7 @@ static struct dm_buffer *__bufio_new(struct dm_bufio_client *c, sector_t block,
 	 * We've had a period where the mutex was unlocked, so need to
 	 * recheck the buffer tree.
 	 */
-	b = cache_get(&c->cache, block);
+	b = cache_get(&c->cache, tree, block);
 	if (b) {
 		__free_buffer_wake(new_b);
 		goto found_buffer;
@@ -1794,13 +1820,13 @@ static struct dm_buffer *__bufio_new(struct dm_bufio_client *c, sector_t block,
 	 * is set.  Otherwise another thread could get it and use
 	 * it before it had been read.
 	 */
-	cache_insert(&c->cache, b);
+	cache_insert(&c->cache, tree, b);
 
 	return b;
 
 found_buffer:
 	if (nf == NF_PREFETCH) {
-		cache_put_and_wake(c, b);
+		cache_put_and_wake(c, tree, b);
 		return NULL;
 	}
 
@@ -1812,7 +1838,7 @@ static struct dm_buffer *__bufio_new(struct dm_bufio_client *c, sector_t block,
 	 * the same buffer, it would deadlock if we waited.
 	 */
 	if (nf == NF_GET && unlikely(test_bit_acquire(B_READING, &b->state))) {
-		cache_put_and_wake(c, b);
+		cache_put_and_wake(c, tree, b);
 		return NULL;
 	}
 
@@ -1846,6 +1872,7 @@ static void *new_read(struct dm_bufio_client *c, sector_t block,
 		      enum new_flag nf, struct dm_buffer **bp,
 		      unsigned short ioprio)
 {
+	struct buffer_tree *tree;
 	int need_submit = 0;
 	struct dm_buffer *b;
 
@@ -1857,10 +1884,11 @@ static void *new_read(struct dm_bufio_client *c, sector_t block,
 	 * Fast path, hopefully the block is already in the cache.  No need
 	 * to get the client lock for this.
 	 */
-	b = cache_get(&c->cache, block);
+	tree = cache_get_tree(&c->cache, block);
+	b = cache_get(&c->cache, tree, block);
 	if (b) {
 		if (nf == NF_PREFETCH) {
-			cache_put_and_wake(c, b);
+			cache_put_and_wake(c, tree, b);
 			return NULL;
 		}
 
@@ -1872,7 +1900,7 @@ static void *new_read(struct dm_bufio_client *c, sector_t block,
 		 * the same buffer, it would deadlock if we waited.
 		 */
 		if (nf == NF_GET && unlikely(test_bit_acquire(B_READING, &b->state))) {
-			cache_put_and_wake(c, b);
+			cache_put_and_wake(c, tree, b);
 			return NULL;
 		}
 	}
@@ -1882,7 +1910,7 @@ static void *new_read(struct dm_bufio_client *c, sector_t block,
 			return NULL;
 
 		dm_bufio_lock(c);
-		b = __bufio_new(c, block, nf, &need_submit, &write_list);
+		b = __bufio_new(c, tree, block, nf, &need_submit, &write_list);
 		dm_bufio_unlock(c);
 	}
 
@@ -1969,18 +1997,20 @@ static void __dm_bufio_prefetch(struct dm_bufio_client *c,
 	blk_start_plug(&plug);
 
 	for (; n_blocks--; block++) {
-		int need_submit;
+		struct buffer_tree *tree;
 		struct dm_buffer *b;
+		int need_submit;
 
-		b = cache_get(&c->cache, block);
+		tree = cache_get_tree(&c->cache, block);
+		b = cache_get(&c->cache, tree, block);
 		if (b) {
 			/* already in cache */
-			cache_put_and_wake(c, b);
+			cache_put_and_wake(c, tree, b);
 			continue;
 		}
 
 		dm_bufio_lock(c);
-		b = __bufio_new(c, block, NF_PREFETCH, &need_submit,
+		b = __bufio_new(c, tree, block, NF_PREFETCH, &need_submit,
 				&write_list);
 		if (unlikely(!list_empty(&write_list))) {
 			dm_bufio_unlock(c);
@@ -2025,6 +2055,7 @@ EXPORT_SYMBOL_GPL(dm_bufio_prefetch_with_ioprio);
 void dm_bufio_release(struct dm_buffer *b)
 {
 	struct dm_bufio_client *c = b->c;
+	struct buffer_tree *tree = cache_get_tree(&c->cache, b->block);
 
 	/*
 	 * If there were errors on the buffer, and the buffer is not
@@ -2038,7 +2069,7 @@ void dm_bufio_release(struct dm_buffer *b)
 		dm_bufio_lock(c);
 
 		/* cache remove can fail if there are other holders */
-		if (cache_remove(&c->cache, b)) {
+		if (cache_remove(&c->cache, tree, b)) {
 			__free_buffer_wake(b);
 			dm_bufio_unlock(c);
 			return;
@@ -2047,7 +2078,7 @@ void dm_bufio_release(struct dm_buffer *b)
 		dm_bufio_unlock(c);
 	}
 
-	cache_put_and_wake(c, b);
+	cache_put_and_wake(c, tree, b);
 }
 EXPORT_SYMBOL_GPL(dm_bufio_release);
 
@@ -2066,7 +2097,8 @@ void dm_bufio_mark_partial_buffer_dirty(struct dm_buffer *b,
 	if (!test_and_set_bit(B_DIRTY, &b->state)) {
 		b->dirty_start = start;
 		b->dirty_end = end;
-		cache_mark(&c->cache, b, LIST_DIRTY);
+		cache_mark(&c->cache, cache_get_tree(&c->cache, b->block), b,
+			   LIST_DIRTY);
 	} else {
 		if (start < b->dirty_start)
 			b->dirty_start = start;
@@ -2131,6 +2163,7 @@ int dm_bufio_write_dirty_buffers(struct dm_bufio_client *c)
 	lru_iter_begin(&c->cache.lru[LIST_DIRTY], &it);
 	while ((e = lru_iter_next(&it, is_writing, c))) {
 		struct dm_buffer *b = le_to_buffer(e);
+		struct buffer_tree *tree;
 		__cache_inc_buffer(b);
 
 		BUG_ON(test_bit(B_READING, &b->state));
@@ -2144,10 +2177,12 @@ int dm_bufio_write_dirty_buffers(struct dm_bufio_client *c)
 			wait_on_bit_io(&b->state, B_WRITING, TASK_UNINTERRUPTIBLE);
 		}
 
+		tree = cache_get_tree(&c->cache, b->block);
+
 		if (!test_bit(B_DIRTY, &b->state) && !test_bit(B_WRITING, &b->state))
-			cache_mark(&c->cache, b, LIST_CLEAN);
+			cache_mark(&c->cache, tree, b, LIST_CLEAN);
 
-		cache_put_and_wake(c, b);
+		cache_put_and_wake(c, tree, b);
 
 		cond_resched();
 	}
@@ -2215,17 +2250,18 @@ EXPORT_SYMBOL_GPL(dm_bufio_issue_discard);
 
 static void forget_buffer(struct dm_bufio_client *c, sector_t block)
 {
+	struct buffer_tree *tree = cache_get_tree(&c->cache, block);
 	struct dm_buffer *b;
 
-	b = cache_get(&c->cache, block);
+	b = cache_get(&c->cache, tree, block);
 	if (b) {
 		if (likely(!smp_load_acquire(&b->state))) {
-			if (cache_remove(&c->cache, b))
+			if (cache_remove(&c->cache, tree, b))
 				__free_buffer_wake(b);
 			else
-				cache_put_and_wake(c, b);
+				cache_put_and_wake(c, tree, b);
 		} else {
-			cache_put_and_wake(c, b);
+			cache_put_and_wake(c, tree, b);
 		}
 	}
 }