diff --git a/fs/bcachefs/btree_cache.c b/fs/bcachefs/btree_cache.c
index 4ff57925fb2c..ef7ca552a0a3 100644
--- a/fs/bcachefs/btree_cache.c
+++ b/fs/bcachefs/btree_cache.c
@@ -554,7 +554,6 @@ struct btree *bch2_btree_node_mem_alloc(struct bch_fs *c)
 
 	list_del_init(&b->list);
 	mutex_unlock(&bc->lock);
-	memalloc_nofs_restore(flags);
 out:
 	b->flags		= 0;
 	b->written		= 0;
@@ -567,6 +566,7 @@ struct btree *bch2_btree_node_mem_alloc(struct bch_fs *c)
 	bch2_time_stats_update(&c->times[BCH_TIME_btree_node_mem_alloc],
 			       start_time);
 
+	memalloc_nofs_restore(flags);
 	return b;
 err:
 	/* Try to cannibalize another cached btree node: */
@@ -582,6 +582,7 @@ struct btree *bch2_btree_node_mem_alloc(struct bch_fs *c)
 	}
 
 	mutex_unlock(&bc->lock);
+	memalloc_nofs_restore(flags);
 	return ERR_PTR(-ENOMEM);
 }