bcachefs: fix memalloc_nofs_restore() usage

Signed-off-by: Kent Overstreet <kent.overstreet@gmail.com>
Signed-off-by: Kent Overstreet <kent.overstreet@linux.dev>

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);
 }