1 files changed, 49 insertions, 16 deletions

diff --git a/libbcachefs/btree_iter.c b/libbcachefs/btree_iter.c
index 47c10886..6ee3366e 100644
--- a/libbcachefs/btree_iter.c
+++ b/libbcachefs/btree_iter.c
@@ -27,6 +27,21 @@ static inline void btree_path_list_add(struct btree_trans *, struct btree_path *
 
 static struct btree_path *btree_path_alloc(struct btree_trans *, struct btree_path *);
 
+/*
+ * Unlocks before scheduling
+ * Note: does not revalidate iterator
+ */
+static inline int bch2_trans_cond_resched(struct btree_trans *trans)
+{
+	if (need_resched() || race_fault()) {
+		bch2_trans_unlock(trans);
+		schedule();
+		return bch2_trans_relock(trans) ? 0 : -EINTR;
+	} else {
+		return 0;
+	}
+}
+
 static inline int __btree_path_cmp(const struct btree_path *l,
 				   enum btree_id	r_btree_id,
 				   bool			r_cached,
@@ -1444,6 +1459,11 @@ static int btree_path_traverse_one(struct btree_trans *trans,
 	unsigned depth_want = path->level;
 	int ret = 0;
 
+	if (unlikely(trans->restarted)) {
+		ret = -EINTR;
+		goto out;
+	}
+
 	/*
 	 * Ensure we obey path->should_be_locked: if it's set, we can't unlock
 	 * and re-traverse the path without a transaction restart:
@@ -1911,30 +1931,41 @@ struct btree *bch2_btree_iter_next_node(struct btree_iter *iter)
 	struct btree_trans *trans = iter->trans;
 	struct btree_path *path = iter->path;
 	struct btree *b = NULL;
+	unsigned l;
 	int ret;
 
+	BUG_ON(trans->restarted);
 	EBUG_ON(iter->path->cached);
 	bch2_btree_iter_verify(iter);
 
-	/* already got to end? */
+	/* already at end? */
 	if (!btree_path_node(path, path->level))
-		goto out;
+		return NULL;
 
-	btree_node_unlock(path, path->level);
-	path->l[path->level].b = BTREE_ITER_NO_NODE_UP;
-	path->level++;
+	/* got to end? */
+	if (!btree_path_node(path, path->level + 1)) {
+		btree_node_unlock(path, path->level);
+		path->l[path->level].b = BTREE_ITER_NO_NODE_UP;
+		path->level++;
+		return NULL;
+	}
 
-	btree_path_set_dirty(path, BTREE_ITER_NEED_TRAVERSE);
-	ret = bch2_btree_path_traverse(trans, path, iter->flags);
-	if (ret)
+	if (!bch2_btree_node_relock(trans, path, path->level + 1)) {
+		__bch2_btree_path_unlock(path);
+		path->l[path->level].b = BTREE_ITER_NO_NODE_GET_LOCKS;
+		path->l[path->level + 1].b = BTREE_ITER_NO_NODE_GET_LOCKS;
+		btree_trans_restart(trans);
+		ret = -EINTR;
 		goto err;
+	}
 
-	/* got to end? */
-	b = btree_path_node(path, path->level);
-	if (!b)
-		goto out;
+	b = btree_path_node(path, path->level + 1);
 
-	if (bpos_cmp(iter->pos, b->key.k.p) < 0) {
+	if (!bpos_cmp(iter->pos, b->key.k.p)) {
+		btree_node_unlock(path, path->level);
+		path->l[path->level].b = BTREE_ITER_NO_NODE_UP;
+		path->level++;
+	} else {
 		/*
 		 * Haven't gotten to the end of the parent node: go back down to
 		 * the next child node
@@ -1943,10 +1974,12 @@ struct btree *bch2_btree_iter_next_node(struct btree_iter *iter)
 			btree_path_set_pos(trans, path, bpos_successor(iter->pos),
 					   iter->flags & BTREE_ITER_INTENT);
 
-		/* Unlock to avoid screwing up our lock invariants: */
-		btree_node_unlock(path, path->level);
-
 		path->level = iter->min_depth;
+
+		for (l = path->level + 1; l < BTREE_MAX_DEPTH; l++)
+			if (btree_lock_want(path, l) == BTREE_NODE_UNLOCKED)
+				btree_node_unlock(path, l);
+
 		btree_path_set_dirty(path, BTREE_ITER_NEED_TRAVERSE);
 		bch2_btree_iter_verify(iter);