btree: Implement faster binary search algorithm

This implements a binary search algorithm for B-Trees that reduces
branching to the absolute minimum necessary for a binary search
algorithm. It also enables the compiler to inlines the comparator to
ensure that the only slowdown when doing binary search is from waiting
for memory accesses.

Consumers must opt into using the faster algorithm.

Signed-off-by: Richard Yao <[email protected]>

Author: ryao

include/sys/btree.h

@@ -105,8 +105,12 @@ typedef struct zfs_btree_index {
 	boolean_t	bti_before;
 } zfs_btree_index_t;
 
-typedef struct btree {
+typedef struct btree zfs_btree_t;
+
+struct btree {
 	int (*bt_compar) (const void *, const void *);
+	void * (*bt_find_in_buf) (zfs_btree_t *, uint8_t *, uint32_t,
+	    const void *, zfs_btree_index_t *);
 	size_t			bt_elem_size;
 	size_t			bt_leaf_size;
 	uint32_t		bt_leaf_cap;
@@ -115,7 +119,33 @@ typedef struct btree {
 	uint64_t		bt_num_nodes;
 	zfs_btree_hdr_t		*bt_root;
 	zfs_btree_leaf_t	*bt_bulk; // non-null if bulk loading
-} zfs_btree_t;
+};
+
+#define	ZFS_BTREE_FIND_IN_BUF_FUNC(NAME, T, COMP)			\
+static void *								\
+NAME(zfs_btree_t *tree, uint8_t *buf, uint32_t nelems,			\
+    const void *value, zfs_btree_index_t *where)			\
+{									\
+	T *i = (T *)buf;						\
+	(void) tree;							\
+	while (nelems > 1) {						\
+		uint32_t half = nelems / 2;				\
+		nelems -= half;						\
+		i += (COMP(&i[half - 1], value) == -1) * half;		\
+	}								\
+									\
+	int comp = COMP(i, value);					\
+	where->bti_offset = (size_t)(i - (T *)buf) / sizeof (T) +	\
+	    (comp < 0);							\
+									\
+	if (comp == 0) {						\
+		where->bti_before = B_FALSE;				\
+		return (i);						\
+	}								\
+									\
+	where->bti_before = B_TRUE;					\
+	return (NULL);							\
+}
 
 /*
  * Allocate and deallocate caches for btree nodes.

module/zfs/btree.c

@@ -198,6 +198,10 @@ zfs_btree_create(zfs_btree_t *tree, int (*compar) (const void *, const void *),
 	zfs_btree_create_custom(tree, compar, size, BTREE_LEAF_SIZE);
 }
 
+static void *
+zfs_btree_find_in_buf(zfs_btree_t *tree, uint8_t *buf, uint32_t nelems,
+    const void *value, zfs_btree_index_t *where);
+
 void
 zfs_btree_create_custom(zfs_btree_t *tree,
     int (*compar) (const void *, const void *),
@@ -208,6 +212,7 @@ zfs_btree_create_custom(zfs_btree_t *tree,
 	ASSERT3U(size, <=, esize / 2);
 	memset(tree, 0, sizeof (*tree));
 	tree->bt_compar = compar;
+	tree->bt_find_in_buf = zfs_btree_find_in_buf;
 	tree->bt_elem_size = size;
 	tree->bt_leaf_size = lsize;
 	tree->bt_leaf_cap = P2ALIGN(esize / size, 2);
@@ -303,7 +308,7 @@ zfs_btree_find(zfs_btree_t *tree, const void *value, zfs_btree_index_t *where)
 			 * element in the last leaf, it's in the last leaf or
 			 * it's not in the tree.
 			 */
-			void *d = zfs_btree_find_in_buf(tree,
+			void *d = tree->bt_find_in_buf(tree,
 			    last_leaf->btl_elems +
 			    last_leaf->btl_hdr.bth_first * size,
 			    last_leaf->btl_hdr.bth_count, value, &idx);
@@ -327,7 +332,7 @@ zfs_btree_find(zfs_btree_t *tree, const void *value, zfs_btree_index_t *where)
 	for (node = (zfs_btree_core_t *)tree->bt_root; depth < tree->bt_height;
 	    node = (zfs_btree_core_t *)node->btc_children[child], depth++) {
 		ASSERT3P(node, !=, NULL);
-		void *d = zfs_btree_find_in_buf(tree, node->btc_elems,
+		void *d = tree->bt_find_in_buf(tree, node->btc_elems,
 		    node->btc_hdr.bth_count, value, &idx);
 		EQUIV(d != NULL, !idx.bti_before);
 		if (d != NULL) {
@@ -347,7 +352,7 @@ zfs_btree_find(zfs_btree_t *tree, const void *value, zfs_btree_index_t *where)
 	 */
 	zfs_btree_leaf_t *leaf = (depth == 0 ?
 	    (zfs_btree_leaf_t *)tree->bt_root : (zfs_btree_leaf_t *)node);
-	void *d = zfs_btree_find_in_buf(tree, leaf->btl_elems +
+	void *d = tree->bt_find_in_buf(tree, leaf->btl_elems +
 	    leaf->btl_hdr.bth_first * size,
 	    leaf->btl_hdr.bth_count, value, &idx);
 
@@ -671,7 +676,7 @@ zfs_btree_insert_into_parent(zfs_btree_t *tree, zfs_btree_hdr_t *old_node,
 	zfs_btree_hdr_t *par_hdr = &parent->btc_hdr;
 	zfs_btree_index_t idx;
 	ASSERT(zfs_btree_is_core(par_hdr));
-	VERIFY3P(zfs_btree_find_in_buf(tree, parent->btc_elems,
+	VERIFY3P(tree->bt_find_in_buf(tree, parent->btc_elems,
 	    par_hdr->bth_count, buf, &idx), ==, NULL);
 	ASSERT(idx.bti_before);
 	uint32_t offset = idx.bti_offset;
@@ -897,7 +902,7 @@ zfs_btree_find_parent_idx(zfs_btree_t *tree, zfs_btree_hdr_t *hdr)
 	}
 	zfs_btree_index_t idx;
 	zfs_btree_core_t *parent = hdr->bth_parent;
-	VERIFY3P(zfs_btree_find_in_buf(tree, parent->btc_elems,
+	VERIFY3P(tree->bt_find_in_buf(tree, parent->btc_elems,
 	    parent->btc_hdr.bth_count, buf, &idx), ==, NULL);
 	ASSERT(idx.bti_before);
 	ASSERT3U(idx.bti_offset, <=, parent->btc_hdr.bth_count);

module/zfs/dsl_scan.c

@@ -4203,13 +4203,17 @@ ext_size_compare(const void *x, const void *y)
 	return (TREE_CMP(*a, *b));
 }
 
+ZFS_BTREE_FIND_IN_BUF_FUNC(ext_size_find_in_buf, uint64_t,
+    ext_size_compare)
+
 static void
 ext_size_create(range_tree_t *rt, void *arg)
 {
 	(void) rt;
 	zfs_btree_t *size_tree = arg;
 
 	zfs_btree_create(size_tree, ext_size_compare, sizeof (uint64_t));
+	size_tree->bt_find_in_buf = ext_size_find_in_buf;
 }
 
 static void

module/zfs/metaslab.c

@@ -1412,6 +1412,13 @@ metaslab_size_tree_full_load(range_tree_t *rt)
 	range_tree_walk(rt, metaslab_size_sorted_add, &arg);
 }
 
+
+ZFS_BTREE_FIND_IN_BUF_FUNC(metaslab_rt_find_rangesize32_in_buf,
+    range_seg32_t, metaslab_rangesize32_compare)
+
+ZFS_BTREE_FIND_IN_BUF_FUNC(metaslab_rt_find_rangesize64_in_buf,
+    range_seg64_t, metaslab_rangesize64_compare)
+
 /*
  * Create any block allocator specific components. The current allocators
  * rely on using both a size-ordered range_tree_t and an array of uint64_t's.
@@ -1428,15 +1435,20 @@ metaslab_rt_create(range_tree_t *rt, void *arg)
 	case RANGE_SEG32:
 		size = sizeof (range_seg32_t);
 		compare = metaslab_rangesize32_compare;
+		zfs_btree_create(size_tree, compare, size);
+		size_tree->bt_find_in_buf =
+		    metaslab_rt_find_rangesize32_in_buf;
 		break;
 	case RANGE_SEG64:
 		size = sizeof (range_seg64_t);
 		compare = metaslab_rangesize64_compare;
+		zfs_btree_create(size_tree, compare, size);
+		size_tree->bt_find_in_buf =
+		    metaslab_rt_find_rangesize64_in_buf;
 		break;
 	default:
 		panic("Invalid range seg type %d", rt->rt_type);
 	}
-	zfs_btree_create(size_tree, compare, size);
 	mrap->mra_floor_shift = metaslab_by_size_min_shift;
 }
 

module/zfs/range_tree.c

@@ -187,6 +187,15 @@ range_tree_seg_gap_compare(const void *x1, const void *x2)
 	return ((r1->rs_start >= r2->rs_end) - (r1->rs_end <= r2->rs_start));
 }
 
+ZFS_BTREE_FIND_IN_BUF_FUNC(range_tree_seg32_find_in_buf, range_seg64_t,
+    range_tree_seg32_compare)
+
+ZFS_BTREE_FIND_IN_BUF_FUNC(range_tree_seg64_find_in_buf, range_seg64_t,
+    range_tree_seg64_compare)
+
+ZFS_BTREE_FIND_IN_BUF_FUNC(range_tree_seg_gap_find_in_buf, range_seg64_t,
+    range_tree_seg_gap_compare)
+
 range_tree_t *
 range_tree_create_gap(const range_tree_ops_t *ops, range_seg_type_t type,
     void *arg, uint64_t start, uint64_t shift, uint64_t gap)
@@ -201,19 +210,24 @@ range_tree_create_gap(const range_tree_ops_t *ops, range_seg_type_t type,
 	case RANGE_SEG32:
 		size = sizeof (range_seg32_t);
 		compare = range_tree_seg32_compare;
+		zfs_btree_create(&rt->rt_root, compare, size);
+		rt->rt_root.bt_find_in_buf = range_tree_seg32_find_in_buf;
 		break;
 	case RANGE_SEG64:
 		size = sizeof (range_seg64_t);
 		compare = range_tree_seg64_compare;
+		zfs_btree_create(&rt->rt_root, compare, size);
+		rt->rt_root.bt_find_in_buf = range_tree_seg64_find_in_buf;
 		break;
 	case RANGE_SEG_GAP:
 		size = sizeof (range_seg_gap_t);
 		compare = range_tree_seg_gap_compare;
+		zfs_btree_create(&rt->rt_root, compare, size);
+		rt->rt_root.bt_find_in_buf = range_tree_seg_gap_find_in_buf;
 		break;
 	default:
 		panic("Invalid range seg type %d", type);
 	}
-	zfs_btree_create(&rt->rt_root, compare, size);
 
 	rt->rt_ops = ops;
 	rt->rt_gap = gap;