4 files changed, 36 insertions, 20 deletions

diff --git a/fs/bcachefs/btree_iter.c b/fs/bcachefs/btree_iter.c
index 75815a1e90ec..200108c0c778 100644
--- a/fs/bcachefs/btree_iter.c
+++ b/fs/bcachefs/btree_iter.c
@@ -2199,30 +2199,14 @@ static inline struct bkey_i *btree_trans_peek_updates(struct btree_iter *iter)
 		: NULL;
 }
 
-static struct bkey_i *__btree_trans_peek_journal(struct btree_trans *trans,
-						 struct btree_path *path)
-{
-	struct journal_keys *keys = &trans->c->journal_keys;
-	size_t idx = bch2_journal_key_search(keys, path->btree_id,
-					     path->level, path->pos);
-
-	while (idx < keys->nr && keys->d[idx].overwritten)
-		idx++;
-
-	return (idx < keys->nr &&
-		keys->d[idx].btree_id	== path->btree_id &&
-		keys->d[idx].level	== path->level)
-		? keys->d[idx].k
-		: NULL;
-}
-
 static noinline
 struct bkey_s_c btree_trans_peek_slot_journal(struct btree_trans *trans,
 					      struct btree_iter *iter)
 {
-	struct bkey_i *k = __btree_trans_peek_journal(trans, iter->path);
+	struct bkey_i *k = bch2_journal_keys_peek(trans->c, iter->btree_id, 0,
+						  iter->path->pos);
 
-	if (k && !bpos_cmp(k->k.p, iter->pos)) {
+	if (k && !bpos_cmp(k->k.p, iter->path->pos)) {
 		iter->k = k->k;
 		return bkey_i_to_s_c(k);
 	} else {
@@ -2236,7 +2220,8 @@ struct bkey_s_c btree_trans_peek_journal(struct btree_trans *trans,
 					 struct bkey_s_c k)
 {
 	struct bkey_i *next_journal =
-		__btree_trans_peek_journal(trans, iter->path);
+		bch2_journal_keys_peek(trans->c, iter->btree_id, 0,
+				       iter->path->pos);
 
 	if (next_journal &&
 	    bpos_cmp(next_journal->k.p,
diff --git a/fs/bcachefs/btree_update_leaf.c b/fs/bcachefs/btree_update_leaf.c
index 5dd86c41c631..2af2d75a06c5 100644
--- a/fs/bcachefs/btree_update_leaf.c
+++ b/fs/bcachefs/btree_update_leaf.c
@@ -1372,6 +1372,7 @@ static int need_whiteout_for_snapshot(struct btree_trans *trans,
 int __must_check bch2_trans_update_by_path(struct btree_trans *trans, struct btree_path *path,
 				   struct bkey_i *k, enum btree_update_flags flags)
 {
+	struct bch_fs *c = trans->c;
 	struct btree_insert_entry *i, n;
 
 	BUG_ON(!path->should_be_locked);
@@ -1431,6 +1432,16 @@ int __must_check bch2_trans_update_by_path(struct btree_trans *trans, struct btr
 
 		i->old_v = bch2_btree_path_peek_slot(path, &i->old_k).v;
 		i->old_btree_u64s = !bkey_deleted(&i->old_k) ? i->old_k.u64s : 0;
+
+		if (unlikely(trans->journal_replay_not_finished)) {
+			struct bkey_i *j_k =
+				bch2_journal_keys_peek(c, n.btree_id, n.level, k->k.p);
+
+			if (j_k && !bpos_cmp(j_k->k.p, i->k->k.p)) {
+				i->old_k = j_k->k;
+				i->old_v = &j_k->v;
+			}
+		}
 	}
 
 	__btree_path_get(n.path, true);
diff --git a/fs/bcachefs/recovery.c b/fs/bcachefs/recovery.c
index b818093eab39..383838d66edf 100644
--- a/fs/bcachefs/recovery.c
+++ b/fs/bcachefs/recovery.c
@@ -94,6 +94,24 @@ size_t bch2_journal_key_search(struct journal_keys *journal_keys,
 	return l;
 }
 
+struct bkey_i *bch2_journal_keys_peek(struct bch_fs *c, enum btree_id btree_id,
+				      unsigned level, struct bpos pos)
+{
+	struct journal_keys *keys = &c->journal_keys;
+	struct journal_key *end = keys->d + keys->nr;
+	struct journal_key *k = keys->d +
+		bch2_journal_key_search(keys, btree_id, level, pos);
+
+	while (k < end && k->overwritten)
+		k++;
+
+	if (k < end &&
+	    k->btree_id	== btree_id &&
+	    k->level	== level)
+		return k->k;
+	return NULL;
+}
+
 static void journal_iter_fix(struct bch_fs *c, struct journal_iter *iter, unsigned idx)
 {
 	struct bkey_i *n = iter->keys->d[idx].k;
diff --git a/fs/bcachefs/recovery.h b/fs/bcachefs/recovery.h
index 21bdad9db249..e6927a918df3 100644
--- a/fs/bcachefs/recovery.h
+++ b/fs/bcachefs/recovery.h
@@ -33,6 +33,8 @@ struct btree_and_journal_iter {
 
 size_t bch2_journal_key_search(struct journal_keys *, enum btree_id,
 			       unsigned, struct bpos);
+struct bkey_i *bch2_journal_keys_peek(struct bch_fs *, enum btree_id,
+				      unsigned, struct bpos pos);
 
 int bch2_journal_key_insert_take(struct bch_fs *, enum btree_id,
 				 unsigned, struct bkey_i *);