2 files changed, 15 insertions, 4 deletions

diff --git a/fs/btrfs/extent_io.c b/fs/btrfs/extent_io.c
index 420fe26d32d5..0d931b143c00 100644
--- a/fs/btrfs/extent_io.c
+++ b/fs/btrfs/extent_io.c
@@ -796,17 +796,25 @@ static void set_state_bits(struct extent_io_tree *tree,
         state->state |= bits_to_set;
 }
 
-static void cache_state(struct extent_state *state,
-                        struct extent_state **cached_ptr)
+static void cache_state_if_flags(struct extent_state *state,
+                                 struct extent_state **cached_ptr,
+                                 const u64 flags)
 {
         if (cached_ptr && !(*cached_ptr)) {
-                if (state->state & (EXTENT_IOBITS | EXTENT_BOUNDARY)) {
+                if (!flags || (state->state & flags)) {
                         *cached_ptr = state;
                         atomic_inc(&state->refs);
                 }
         }
 }
 
+static void cache_state(struct extent_state *state,
+                        struct extent_state **cached_ptr)
+{
+        return cache_state_if_flags(state, cached_ptr,
+                                    EXTENT_IOBITS | EXTENT_BOUNDARY);
+}
+
 /*
  * set some bits on a range in the tree.  This may require allocations or
  * sleeping, so the gfp mask is used to indicate what is allowed.
@@ -1482,7 +1490,7 @@ int find_first_extent_bit(struct extent_io_tree *tree, u64 start,
         state = find_first_extent_bit_state(tree, start, bits);
 got_it:
         if (state) {
-                cache_state(state, cached_state);
+                cache_state_if_flags(state, cached_state, 0);
                 *start_ret = state->start;
                 *end_ret = state->end;
                 ret = 0;
diff --git a/fs/btrfs/transaction.c b/fs/btrfs/transaction.c
index 04dbc800c209..16c704b68704 100644
--- a/fs/btrfs/transaction.c
+++ b/fs/btrfs/transaction.c
@@ -883,6 +883,7 @@ int btrfs_write_marked_extents(struct btrfs_root *root,
                         werr = err;
                 else if (wait_writeback)
                         werr = filemap_fdatawait_range(mapping, start, end);
+                free_extent_state(cached_state);
                 cached_state = NULL;
                 cond_resched();
                 start = end + 1;
@@ -927,6 +928,8 @@ int btrfs_wait_marked_extents(struct btrfs_root *root,
                         err = filemap_fdatawait_range(mapping, start, end);
                 if (err)
                         werr = err;
+                free_extent_state(cached_state);
+                cached_state = NULL;
                 cond_resched();
                 start = end + 1;
         }