diff --git a/fs/bcachefs/compress.c b/fs/bcachefs/compress.c
index 920460a182b4..595d76aa3956 100644
--- a/fs/bcachefs/compress.c
+++ b/fs/bcachefs/compress.c
@@ -7,7 +7,6 @@
 #include "super-io.h"
 
 #include <linux/lz4.h>
-#include <linux/sched/mm.h>
 #include <linux/zlib.h>
 #include <linux/zstd.h>
 
@@ -64,7 +63,7 @@ static struct bbuf __bio_map_or_bounce(struct bch_fs *c, struct bio *bio,
 	struct bbuf ret;
 	struct bio_vec bv;
 	struct bvec_iter iter;
-	unsigned nr_pages = 0, flags;
+	unsigned nr_pages = 0;
 	struct page *stack_pages[16];
 	struct page **pages = NULL;
 	void *data;
@@ -104,10 +103,7 @@ static struct bbuf __bio_map_or_bounce(struct bch_fs *c, struct bio *bio,
 	__bio_for_each_segment(bv, bio, iter, start)
 		pages[nr_pages++] = bv.bv_page;
 
-	flags = memalloc_nofs_save();
 	data = vmap(pages, nr_pages, VM_MAP, PAGE_KERNEL);
-	memalloc_nofs_restore(flags);
-
 	if (pages != stack_pages)
 		kfree(pages);
 
diff --git a/fs/bcachefs/io.c b/fs/bcachefs/io.c
index be59b615b2db..ae2688d3aee6 100644
--- a/fs/bcachefs/io.c
+++ b/fs/bcachefs/io.c
@@ -32,6 +32,7 @@
 
 #include <linux/blkdev.h>
 #include <linux/random.h>
+#include <linux/sched/mm.h>
 
 #ifndef CONFIG_BCACHEFS_NO_LATENCY_ACCT
 
@@ -1052,7 +1053,10 @@ static void __bch2_write(struct closure *cl)
 	struct write_point *wp;
 	struct bio *bio;
 	bool skip_put = true;
+	unsigned nofs_flags;
 	int ret;
+
+	nofs_flags = memalloc_nofs_save();
 again:
 	memset(&op->failed, 0, sizeof(op->failed));
 
@@ -1134,13 +1138,15 @@ static void __bch2_write(struct closure *cl)
 
 	if (!skip_put)
 		continue_at(cl, bch2_write_index, index_update_wq(op));
+out:
+	memalloc_nofs_restore(nofs_flags);
 	return;
 err:
 	op->error = ret;
 	op->flags |= BCH_WRITE_DONE;
 
 	continue_at(cl, bch2_write_index, index_update_wq(op));
-	return;
+	goto out;
 flush_io:
 	/*
 	 * If the write can't all be submitted at once, we generally want to
@@ -1151,7 +1157,7 @@ static void __bch2_write(struct closure *cl)
 	 */
 	if (current->flags & PF_WQ_WORKER) {
 		continue_at(cl, bch2_write_index, index_update_wq(op));
-		return;
+		goto out;
 	}
 
 	closure_sync(cl);
@@ -1162,7 +1168,7 @@ static void __bch2_write(struct closure *cl)
 		if (op->error) {
 			op->flags |= BCH_WRITE_DONE;
 			continue_at_nobarrier(cl, bch2_write_done, NULL);
-			return;
+			goto out;
 		}
 	}