2 files changed, 33 insertions, 10 deletions

diff --git a/fs/bcachefs/btree_cache.c b/fs/bcachefs/btree_cache.c
index 5991ebee228c..15f597ab03e7 100644
--- a/fs/bcachefs/btree_cache.c
+++ b/fs/bcachefs/btree_cache.c
@@ -663,13 +663,9 @@ static noinline struct btree *bch2_btree_node_fill(struct bch_fs *c,
 		return NULL;
 	}
 
-	/*
-	 * Unlock before doing IO:
-	 *
-	 * XXX: ideally should be dropping all btree node locks here
-	 */
-	if (iter && btree_node_read_locked(iter, level + 1))
-		btree_node_unlock(iter, level + 1);
+	/* Unlock before doing IO: */
+	if (iter && sync)
+		bch2_trans_unlock(iter->trans);
 
 	bch2_btree_node_read(c, b, sync);
 
@@ -680,6 +676,16 @@ static noinline struct btree *bch2_btree_node_fill(struct bch_fs *c,
 		return NULL;
 	}
 
+	/*
+	 * XXX: this will probably always fail because btree_iter_relock()
+	 * currently fails for iterators that aren't pointed at a valid btree
+	 * node
+	 */
+	if (iter && !bch2_trans_relock(iter->trans)) {
+		six_unlock_intent(&b->c.lock);
+		return ERR_PTR(-EINTR);
+	}
+
 	if (lock_type == SIX_LOCK_read)
 		six_lock_downgrade(&b->c.lock);
 
@@ -824,9 +830,22 @@ lock_node:
 		}
 	}
 
-	/* XXX: waiting on IO with btree locks held: */
-	wait_on_bit_io(&b->flags, BTREE_NODE_read_in_flight,
-		       TASK_UNINTERRUPTIBLE);
+	if (unlikely(btree_node_read_in_flight(b))) {
+		six_unlock_type(&b->c.lock, lock_type);
+		bch2_trans_unlock(iter->trans);
+
+		wait_on_bit_io(&b->flags, BTREE_NODE_read_in_flight,
+			       TASK_UNINTERRUPTIBLE);
+
+		/*
+		 * XXX: check if this always fails - btree_iter_relock()
+		 * currently fails for iterators that aren't pointed at a valid
+		 * btree node
+		 */
+		if (iter && !bch2_trans_relock(iter->trans))
+			return ERR_PTR(-EINTR);
+		goto retry;
+	}
 
 	prefetch(b->aux_data);
 
diff --git a/fs/bcachefs/btree_iter.c b/fs/bcachefs/btree_iter.c
index 17338410d1fe..5c38562ab206 100644
--- a/fs/bcachefs/btree_iter.c
+++ b/fs/bcachefs/btree_iter.c
@@ -1191,7 +1191,11 @@ static __always_inline int btree_iter_down(struct btree_iter *iter,
 	if (iter->flags & BTREE_ITER_PREFETCH)
 		btree_iter_prefetch(iter);
 
+	if (btree_node_read_locked(iter, level + 1))
+		btree_node_unlock(iter, level + 1);
 	iter->level = level;
+
+	bch2_btree_iter_verify_locks(iter);
 err:
 	bch2_bkey_buf_exit(&tmp, c);
 	return ret;