4 files changed, 139 insertions, 60 deletions

diff --git a/fs/bcachefs/alloc_background.c b/fs/bcachefs/alloc_background.c
index 5bc8e7531403..9ce53164d9ac 100644
--- a/fs/bcachefs/alloc_background.c
+++ b/fs/bcachefs/alloc_background.c
@@ -211,33 +211,31 @@ void bch2_alloc_to_text(struct printbuf *out, struct bch_fs *c,
 int bch2_alloc_read(struct bch_fs *c, struct journal_keys *journal_keys)
 {
 	struct btree_trans trans;
-	struct btree_iter *iter;
+	struct btree_and_journal_iter iter;
 	struct bkey_s_c k;
 	struct bch_dev *ca;
-	struct journal_key *j;
 	unsigned i;
-	int ret;
+	int ret = 0;
 
 	bch2_trans_init(&trans, c, 0, 0);
 
-	for_each_btree_key(&trans, iter, BTREE_ID_ALLOC, POS_MIN, 0, k, ret)
+	bch2_btree_and_journal_iter_init(&iter, &trans, journal_keys,
+					 BTREE_ID_ALLOC, POS_MIN);
+
+	while ((k = bch2_btree_and_journal_iter_peek(&iter)).k) {
 		bch2_mark_key(c, k, 0, 0, NULL, 0,
 			      BTREE_TRIGGER_ALLOC_READ|
 			      BTREE_TRIGGER_NOATOMIC);
 
+		bch2_btree_and_journal_iter_advance(&iter);
+	}
+
 	ret = bch2_trans_exit(&trans) ?: ret;
 	if (ret) {
 		bch_err(c, "error reading alloc info: %i", ret);
 		return ret;
 	}
 
-	for_each_journal_key(*journal_keys, j)
-		if (j->btree_id == BTREE_ID_ALLOC)
-			bch2_mark_key(c, bkey_i_to_s_c(j->k),
-				      0, 0, NULL, 0,
-				      BTREE_TRIGGER_ALLOC_READ|
-				      BTREE_TRIGGER_NOATOMIC);
-
 	percpu_down_write(&c->mark_lock);
 	bch2_dev_usage_from_buckets(c);
 	percpu_up_write(&c->mark_lock);
diff --git a/fs/bcachefs/ec.c b/fs/bcachefs/ec.c
index 0e2acd4f5712..59985227385b 100644
--- a/fs/bcachefs/ec.c
+++ b/fs/bcachefs/ec.c
@@ -1280,9 +1280,8 @@ int bch2_stripes_write(struct bch_fs *c, unsigned flags, bool *wrote)
 int bch2_stripes_read(struct bch_fs *c, struct journal_keys *journal_keys)
 {
 	struct btree_trans trans;
-	struct btree_iter *btree_iter;
-	struct journal_iter journal_iter;
-	struct bkey_s_c btree_k, journal_k;
+	struct btree_and_journal_iter iter;
+	struct bkey_s_c k;
 	int ret;
 
 	ret = bch2_fs_ec_start(c);
@@ -1291,38 +1290,16 @@ int bch2_stripes_read(struct bch_fs *c, struct journal_keys *journal_keys)
 
 	bch2_trans_init(&trans, c, 0, 0);
 
-	btree_iter	= bch2_trans_get_iter(&trans, BTREE_ID_EC, POS_MIN, 0);
-	journal_iter	= bch2_journal_iter_init(journal_keys, BTREE_ID_EC);
+	bch2_btree_and_journal_iter_init(&iter, &trans, journal_keys,
+					 BTREE_ID_EC, POS_MIN);
 
-	btree_k		= bch2_btree_iter_peek(btree_iter);
-	journal_k	= bch2_journal_iter_peek(&journal_iter);
 
-	while (1) {
-		bool btree;
-
-		if (btree_k.k && journal_k.k) {
-			int cmp = bkey_cmp(btree_k.k->p, journal_k.k->p);
-
-			if (!cmp)
-				btree_k = bch2_btree_iter_next(btree_iter);
-			btree = cmp < 0;
-		} else if (btree_k.k) {
-			btree = true;
-		} else if (journal_k.k) {
-			btree = false;
-		} else {
-			break;
-		}
-
-		bch2_mark_key(c, btree ? btree_k : journal_k,
-			      0, 0, NULL, 0,
+	while ((k = bch2_btree_and_journal_iter_peek(&iter)).k) {
+		bch2_mark_key(c, k, 0, 0, NULL, 0,
 			      BTREE_TRIGGER_ALLOC_READ|
 			      BTREE_TRIGGER_NOATOMIC);
 
-		if (btree)
-			btree_k = bch2_btree_iter_next(btree_iter);
-		else
-			journal_k = bch2_journal_iter_next(&journal_iter);
+		bch2_btree_and_journal_iter_advance(&iter);
 	}
 
 	ret = bch2_trans_exit(&trans) ?: ret;
diff --git a/fs/bcachefs/recovery.c b/fs/bcachefs/recovery.c
index 8ecd4abc8eeb..29e6f9f00bad 100644
--- a/fs/bcachefs/recovery.c
+++ b/fs/bcachefs/recovery.c
@@ -27,26 +27,15 @@
 
 /* iterate over keys read from the journal: */
 
-struct journal_iter bch2_journal_iter_init(struct journal_keys *keys,
-					   enum btree_id id)
-{
-	return (struct journal_iter) {
-		.keys		= keys,
-		.k		= keys->d,
-		.btree_id	= id,
-	};
-}
-
 struct bkey_s_c bch2_journal_iter_peek(struct journal_iter *iter)
 {
-	while (1) {
-		if (iter->k == iter->keys->d + iter->keys->nr)
-			return bkey_s_c_null;
-
+	while (iter->k) {
 		if (iter->k->btree_id == iter->btree_id)
 			return bkey_i_to_s_c(iter->k->k);
 
 		iter->k++;
+		if (iter->k == iter->keys->d + iter->keys->nr)
+			iter->k = NULL;
 	}
 
 	return bkey_s_c_null;
@@ -54,13 +43,110 @@ struct bkey_s_c bch2_journal_iter_peek(struct journal_iter *iter)
 
 struct bkey_s_c bch2_journal_iter_next(struct journal_iter *iter)
 {
-	if (iter->k == iter->keys->d + iter->keys->nr)
+	if (!iter->k)
 		return bkey_s_c_null;
 
 	iter->k++;
+	if (iter->k == iter->keys->d + iter->keys->nr)
+		iter->k = NULL;
+
 	return bch2_journal_iter_peek(iter);
 }
 
+void bch2_btree_and_journal_iter_advance(struct btree_and_journal_iter *iter)
+{
+	switch (iter->last) {
+	case none:
+		break;
+	case btree:
+		bch2_btree_iter_next(iter->btree);
+		break;
+	case journal:
+		bch2_journal_iter_next(&iter->journal);
+		break;
+	}
+
+	iter->last = none;
+}
+
+struct bkey_s_c bch2_btree_and_journal_iter_peek(struct btree_and_journal_iter *iter)
+{
+	struct bkey_s_c ret;
+
+	while (1) {
+		struct bkey_s_c btree_k		= bch2_btree_iter_peek(iter->btree);
+		struct bkey_s_c journal_k	= bch2_journal_iter_peek(&iter->journal);
+
+		if (btree_k.k && journal_k.k) {
+			int cmp = bkey_cmp(btree_k.k->p, journal_k.k->p);
+
+			if (!cmp)
+				bch2_btree_iter_next(iter->btree);
+
+			iter->last = cmp < 0 ? btree : journal;
+		} else if (btree_k.k) {
+			iter->last = btree;
+		} else if (journal_k.k) {
+			iter->last = journal;
+		} else {
+			iter->last = none;
+			return bkey_s_c_null;
+		}
+
+		ret = iter->last == journal ? journal_k : btree_k;
+		if (!bkey_deleted(ret.k))
+			break;
+
+		bch2_btree_and_journal_iter_advance(iter);
+	}
+
+	return ret;
+}
+
+struct bkey_s_c bch2_btree_and_journal_iter_next(struct btree_and_journal_iter *iter)
+{
+	bch2_btree_and_journal_iter_advance(iter);
+
+	return bch2_btree_and_journal_iter_peek(iter);
+}
+
+struct journal_key *journal_key_search(struct journal_keys *journal_keys,
+				       enum btree_id id, struct bpos pos)
+{
+	size_t l = 0, r = journal_keys->nr, m;
+
+	while (l < r) {
+		m = l + ((r - l) >> 1);
+		if ((cmp_int(id, journal_keys->d[m].btree_id) ?:
+		     bkey_cmp(pos, journal_keys->d[m].k->k.p)) > 0)
+			l = m + 1;
+		else
+			r = m;
+	}
+
+	BUG_ON(l < journal_keys->nr &&
+	       (cmp_int(id, journal_keys->d[l].btree_id) ?:
+		bkey_cmp(pos, journal_keys->d[l].k->k.p)) > 0);
+
+	BUG_ON(l &&
+	       (cmp_int(id, journal_keys->d[l - 1].btree_id) ?:
+		bkey_cmp(pos, journal_keys->d[l - 1].k->k.p)) <= 0);
+
+	return l < journal_keys->nr ? journal_keys->d + l : NULL;
+}
+
+void bch2_btree_and_journal_iter_init(struct btree_and_journal_iter *iter,
+				      struct btree_trans *trans,
+				      struct journal_keys *journal_keys,
+				      enum btree_id id, struct bpos pos)
+{
+	iter->journal.keys	= journal_keys;
+	iter->journal.k		= journal_key_search(journal_keys, id, pos);
+	iter->journal.btree_id	= id;
+
+	iter->btree = bch2_trans_get_iter(trans, id, pos, 0);
+}
+
 /* sort and dedup all keys in the journal: */
 
 static void journal_entries_free(struct list_head *list)
diff --git a/fs/bcachefs/recovery.h b/fs/bcachefs/recovery.h
index 479ea46f8dcb..ccd84a8fe60d 100644
--- a/fs/bcachefs/recovery.h
+++ b/fs/bcachefs/recovery.h
@@ -24,10 +24,28 @@ struct journal_iter {
 	enum btree_id		btree_id;
 };
 
-struct journal_iter bch2_journal_iter_init(struct journal_keys *,
-					   enum btree_id);
-struct bkey_s_c bch2_journal_iter_peek(struct journal_iter *);
-struct bkey_s_c bch2_journal_iter_next(struct journal_iter *);
+struct btree_and_journal_iter {
+	enum btree_id		btree_id;
+
+	struct btree_iter	*btree;
+	struct journal_iter	journal;
+
+	enum last_key_returned {
+		none,
+		btree,
+		journal,
+	}			last;
+};
+
+void bch2_btree_and_journal_iter_advance(struct btree_and_journal_iter *);
+struct bkey_s_c bch2_btree_and_journal_iter_peek(struct btree_and_journal_iter *);
+struct bkey_s_c bch2_btree_and_journal_iter_next(struct btree_and_journal_iter *);
+struct journal_key *journal_key_search(struct journal_keys *,
+				       enum btree_id, struct bpos);
+void bch2_btree_and_journal_iter_init(struct btree_and_journal_iter *,
+				      struct btree_trans *,
+				      struct journal_keys *,
+				      enum btree_id, struct bpos);
 
 int bch2_fs_recovery(struct bch_fs *);
 int bch2_fs_initialize(struct bch_fs *);