diff --git a/fs/btrfs/extent_map.c b/fs/btrfs/extent_map.c
index 370ed97e52f5b047132b57fe7e0168902d56133b..3021564ac62c41a9a6f7d1a07eecd641993fe657 100644
--- a/fs/btrfs/extent_map.c
+++ b/fs/btrfs/extent_map.c
@@ -456,6 +456,7 @@ int clear_extent_bit(struct extent_map_tree *tree, u64 start, u64 end,
 	struct extent_state *state;
 	struct extent_state *prealloc = NULL;
 	struct rb_node *node;
+	unsigned long flags;
 	int err;
 	int set = 0;
 
@@ -466,7 +467,7 @@ int clear_extent_bit(struct extent_map_tree *tree, u64 start, u64 end,
 			return -ENOMEM;
 	}
 
-	write_lock_irq(&tree->lock);
+	write_lock_irqsave(&tree->lock, flags);
 	/*
 	 * this search will find the extents that end after
 	 * our range starts
@@ -533,7 +534,7 @@ int clear_extent_bit(struct extent_map_tree *tree, u64 start, u64 end,
 	goto search_again;
 
 out:
-	write_unlock_irq(&tree->lock);
+	write_unlock_irqrestore(&tree->lock, flags);
 	if (prealloc)
 		free_extent_state(prealloc);
 
@@ -542,7 +543,7 @@ int clear_extent_bit(struct extent_map_tree *tree, u64 start, u64 end,
 search_again:
 	if (start >= end)
 		goto out;
-	write_unlock_irq(&tree->lock);
+	write_unlock_irqrestore(&tree->lock, flags);
 	if (mask & __GFP_WAIT)
 		cond_resched();
 	goto again;
@@ -628,6 +629,7 @@ int set_extent_bit(struct extent_map_tree *tree, u64 start, u64 end, int bits,
 	struct extent_state *state;
 	struct extent_state *prealloc = NULL;
 	struct rb_node *node;
+	unsigned long flags;
 	int err = 0;
 	int set;
 	u64 last_start;
@@ -639,7 +641,7 @@ int set_extent_bit(struct extent_map_tree *tree, u64 start, u64 end, int bits,
 			return -ENOMEM;
 	}
 
-	write_lock_irq(&tree->lock);
+	write_lock_irqsave(&tree->lock, flags);
 	/*
 	 * this search will find all the extents that end after
 	 * our range starts.
@@ -759,7 +761,7 @@ int set_extent_bit(struct extent_map_tree *tree, u64 start, u64 end, int bits,
 	goto search_again;
 
 out:
-	write_unlock_irq(&tree->lock);
+	write_unlock_irqrestore(&tree->lock, flags);
 	if (prealloc)
 		free_extent_state(prealloc);
 
@@ -768,7 +770,7 @@ int set_extent_bit(struct extent_map_tree *tree, u64 start, u64 end, int bits,
 search_again:
 	if (start > end)
 		goto out;
-	write_unlock_irq(&tree->lock);
+	write_unlock_irqrestore(&tree->lock, flags);
 	if (mask & __GFP_WAIT)
 		cond_resched();
 	goto again;