bcachefs: use memalloc_nofs_save() for vmalloc allocation

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

fs/bcachefs/btree_cache.c

@@ -9,6 +9,7 @@
 #include "trace.h"
 
 #include <linux/prefetch.h>
+#include <linux/sched/mm.h>
 
 const char * const bch2_btree_ids[] = {
 #define x(kwd, val, name) name,
@@ -509,7 +510,9 @@ struct btree *bch2_btree_node_mem_alloc(struct bch_fs *c)
 	struct btree_cache *bc = &c->btree_cache;
 	struct btree *b;
 	u64 start_time = local_clock();
+	unsigned flags;
 
+	flags = memalloc_nofs_save();
 	mutex_lock(&bc->lock);
 
 	/*
@@ -547,6 +550,7 @@ 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;