4 files changed, 22 insertions, 17 deletions

diff --git a/fs/bcachefs/btree_journal_iter.c b/fs/bcachefs/btree_journal_iter.c
index 2e9078b5f538..f63c349e09da 100644
--- a/fs/bcachefs/btree_journal_iter.c
+++ b/fs/bcachefs/btree_journal_iter.c
@@ -265,13 +265,8 @@ int bch2_journal_key_insert_take(struct bch_fs *c, enum btree_id id,
 	struct journal_key n = {
 		.btree_id	= id,
 		.level		= level,
-		.k		= k,
 		.allocated	= true,
-		/*
-		 * Ensure these keys are done last by journal replay, to unblock
-		 * journal reclaim:
-		 */
-		.journal_seq	= U64_MAX,
+		.allocated_k	= k,
 	};
 	struct journal_keys *keys = &c->journal_keys;
 	size_t idx = bch2_journal_key_search(keys, id, level, k->k.p);
@@ -292,7 +287,7 @@ int bch2_journal_key_insert_take(struct bch_fs *c, enum btree_id id,
 		}
 
 		if (keys->data[idx].allocated)
-			kfree(keys->data[idx].k);
+			kfree(keys->data[idx].allocated_k);
 		keys->data[idx] = n;
 		return 0;
 	}
@@ -702,7 +697,7 @@ void bch2_journal_keys_put(struct bch_fs *c)
 			kfree(i->overwritten_range);
 
 		if (i->allocated)
-			kfree(i->k);
+			kfree(i->allocated_k);
 	}
 
 	kvfree(keys->data);
@@ -781,7 +776,6 @@ int bch2_journal_keys_sort(struct bch_fs *c)
 					.btree_id	= entry->btree_id,
 					.level		= entry->level,
 					.rewind		= rewind,
-					.k		= k,
 					.journal_seq	= le64_to_cpu(i->j.seq),
 					.journal_offset	= k->_data - i->j._data,
 				};
@@ -828,6 +822,8 @@ void bch2_shoot_down_journal_keys(struct bch_fs *c, enum btree_id btree,
 		      bpos_ge(k->k.p, start) &&
 		      bpos_le(k->k.p, end)))
 			keys->data[dst++] = *i;
+		else if (i->allocated)
+			kfree(i->allocated_k);
 	}
 	keys->nr = keys->gap = dst;
 }
diff --git a/fs/bcachefs/btree_journal_iter.h b/fs/bcachefs/btree_journal_iter.h
index a1b26832869f..cfd2061bc966 100644
--- a/fs/bcachefs/btree_journal_iter.h
+++ b/fs/bcachefs/btree_journal_iter.h
@@ -29,10 +29,21 @@ struct btree_and_journal_iter {
 	bool			fail_if_too_many_whiteouts;
 };
 
+static inline u32 journal_entry_radix_idx(struct bch_fs *c, u64 seq)
+{
+	return (seq - c->journal_entries_base_seq) & (~0U >> 1);
+}
+
 static inline struct bkey_i *journal_key_k(struct bch_fs *c,
 					   const struct journal_key *k)
 {
-	return k->k;
+	if (k->allocated)
+		return k->allocated_k;
+
+	struct journal_replay *i =
+		*genradix_ptr(&c->journal_entries, journal_entry_radix_idx(c, k->journal_seq));
+
+	return (struct bkey_i *) (i->j._data + k->journal_offset);
 }
 
 static inline int __journal_key_btree_cmp(enum btree_id	l_btree_id,
diff --git a/fs/bcachefs/btree_journal_iter_types.h b/fs/bcachefs/btree_journal_iter_types.h
index 0b4f6c1cb183..e9d8628edec6 100644
--- a/fs/bcachefs/btree_journal_iter_types.h
+++ b/fs/bcachefs/btree_journal_iter_types.h
@@ -30,7 +30,10 @@ struct journal_key_range_overwritten {
 };
 
 struct journal_key {
-	u64			journal_seq;
+	union {
+		u64		journal_seq;
+		struct bkey_i	*allocated_k;
+	};
 	u32			journal_offset;
 	enum btree_id		btree_id:8;
 	unsigned		level:8;
@@ -39,7 +42,6 @@ struct journal_key {
 	bool			rewind:1;
 	struct journal_key_range_overwritten __rcu *
 				overwritten_range;
-	struct bkey_i		*k;
 };
 
 struct journal_keys {
diff --git a/fs/bcachefs/journal_io.c b/fs/bcachefs/journal_io.c
index 093e4acad085..6e8a89a0f244 100644
--- a/fs/bcachefs/journal_io.c
+++ b/fs/bcachefs/journal_io.c
@@ -3,6 +3,7 @@
 #include "alloc_background.h"
 #include "alloc_foreground.h"
 #include "btree_io.h"
+#include "btree_journal_iter.h"
 #include "btree_update_interior.h"
 #include "btree_write_buffer.h"
 #include "buckets.h"
@@ -106,11 +107,6 @@ static bool jset_csum_good(struct bch_fs *c, struct jset *j, struct bch_csum *cs
 	return !bch2_crc_cmp(j->csum, *csum);
 }
 
-static inline u32 journal_entry_radix_idx(struct bch_fs *c, u64 seq)
-{
-	return (seq - c->journal_entries_base_seq) & (~0U >> 1);
-}
-
 static void __journal_replay_free(struct bch_fs *c,
 				  struct journal_replay *i)
 {