bcachefs: Rebalance now adds replicas if needed

Signed-off-by: Kent Overstreet <kent.overstreet@gmail.com>

2 files changed, 20 insertions, 27 deletions

diff --git a/fs/bcachefs/bcachefs_format.h b/fs/bcachefs/bcachefs_format.h
index 667170b54ce6..4577d77a9f38 100644
--- a/fs/bcachefs/bcachefs_format.h
+++ b/fs/bcachefs/bcachefs_format.h
@@ -657,7 +657,7 @@ struct bch_reservation {
 
 /* Maximum possible size of an entire extent value: */
 #define BKEY_EXTENT_VAL_U64s_MAX				\
-	(BKEY_EXTENT_PTR_U64s_MAX * (BCH_REPLICAS_MAX + 1))
+	(1 + BKEY_EXTENT_PTR_U64s_MAX * (BCH_REPLICAS_MAX + 1))
 
 #define BKEY_PADDED(key)	__BKEY_PADDED(key, BKEY_EXTENT_VAL_U64s_MAX)
 
diff --git a/fs/bcachefs/rebalance.c b/fs/bcachefs/rebalance.c
index 4797d620fe77..84b3fb6eb101 100644
--- a/fs/bcachefs/rebalance.c
+++ b/fs/bcachefs/rebalance.c
@@ -42,9 +42,6 @@ void bch2_rebalance_add_key(struct bch_fs *c,
 	const union bch_extent_entry *entry;
 	struct extent_ptr_decoded p;
 
-	if (!bkey_extent_is_data(k.k))
-		return;
-
 	if (!io_opts->background_target &&
 	    !io_opts->background_compression)
 		return;
@@ -72,30 +69,26 @@ static enum data_cmd rebalance_pred(struct bch_fs *c, void *arg,
 				    struct bch_io_opts *io_opts,
 				    struct data_opts *data_opts)
 {
-	switch (k.k->type) {
-	case KEY_TYPE_extent: {
-		struct bkey_s_c_extent e = bkey_s_c_to_extent(k);
-		const union bch_extent_entry *entry;
-		struct extent_ptr_decoded p;
-
-		/* Make sure we have room to add a new pointer: */
-		if (bkey_val_u64s(e.k) + BKEY_EXTENT_PTR_U64s_MAX >
-		    BKEY_EXTENT_VAL_U64s_MAX)
-			return DATA_SKIP;
-
-		extent_for_each_ptr_decode(e, p, entry)
-			if (rebalance_ptr_pred(c, p, io_opts))
-				goto found;
-
-		return DATA_SKIP;
-found:
-		data_opts->target		= io_opts->background_target;
-		data_opts->btree_insert_flags	= 0;
-		return DATA_ADD_REPLICAS;
-	}
-	default:
-		return DATA_SKIP;
+	struct bkey_ptrs_c ptrs = bch2_bkey_ptrs_c(k);
+	const union bch_extent_entry *entry;
+	struct extent_ptr_decoded p;
+	unsigned nr_replicas = 0;
+
+	bkey_for_each_ptr_decode(k.k, ptrs, p, entry) {
+		nr_replicas += !p.ptr.cached;
+
+		if (rebalance_ptr_pred(c, p, io_opts))
+			goto found;
 	}
+
+	if (nr_replicas < io_opts->data_replicas)
+		goto found;
+
+	return DATA_SKIP;
+found:
+	data_opts->target		= io_opts->background_target;
+	data_opts->btree_insert_flags	= 0;
+	return DATA_ADD_REPLICAS;
 }
 
 struct rebalance_work {