3 files changed, 100 insertions, 36 deletions

diff --git a/fs/btrfs/extent_io.c b/fs/btrfs/extent_io.c
index 7c70613eb72c..c7a5e860fe21 100644
--- a/fs/btrfs/extent_io.c
+++ b/fs/btrfs/extent_io.c
@@ -471,10 +471,14 @@ static int clear_state_bit(struct extent_io_tree *tree,
  * bits were already set, or zero if none of the bits were already set.
  */
 int clear_extent_bit(struct extent_io_tree *tree, u64 start, u64 end,
-                     int bits, int wake, int delete, gfp_t mask)
+                     int bits, int wake, int delete,
+                     struct extent_state **cached_state,
+                     gfp_t mask)
 {
         struct extent_state *state;
+        struct extent_state *cached;
         struct extent_state *prealloc = NULL;
+        struct rb_node *next_node;
         struct rb_node *node;
         u64 last_end;
         int err;
@@ -488,6 +492,17 @@ again:
         }
 
         spin_lock(&tree->lock);
+        if (cached_state) {
+                cached = *cached_state;
+                *cached_state = NULL;
+                if (cached->tree && cached->start == start) {
+                        atomic_dec(&cached->refs);
+                        state = cached;
+                        last_end = state->end;
+                        goto found;
+                }
+                free_extent_state(cached);
+        }
         /*
          * this search will find the extents that end after
          * our range starts
@@ -496,6 +511,7 @@ again:
         if (!node)
                 goto out;
         state = rb_entry(node, struct extent_state, rb_node);
+hit_next:
         if (state->start > end)
                 goto out;
         WARN_ON(state->end < start);
@@ -555,11 +571,21 @@ again:
                 prealloc = NULL;
                 goto out;
         }
-
+found:
+        if (state->end < end && prealloc && !need_resched())
+                next_node = rb_next(&state->rb_node);
+        else
+                next_node = NULL;
         set |= clear_state_bit(tree, state, bits, wake, delete);
         if (last_end == (u64)-1)
                 goto out;
         start = last_end + 1;
+        if (start <= end && next_node) {
+                state = rb_entry(next_node, struct extent_state,
+                                 rb_node);
+                if (state->start == start)
+                        goto hit_next;
+        }
         goto search_again;
 
 out:
@@ -653,6 +679,17 @@ static void set_state_bits(struct extent_io_tree *tree,
         state->state |= bits;
 }
 
+static void cache_state(struct extent_state *state,
+                        struct extent_state **cached_ptr)
+{
+        if (cached_ptr && !(*cached_ptr)) {
+                if (state->state & (EXTENT_IOBITS | EXTENT_BOUNDARY)) {
+                        *cached_ptr = state;
+                        atomic_inc(&state->refs);
+                }
+        }
+}
+
 /*
  * 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.
@@ -666,6 +703,7 @@ static void set_state_bits(struct extent_io_tree *tree,
 
 static int set_extent_bit(struct extent_io_tree *tree, u64 start, u64 end,
                           int bits, int exclusive_bits, u64 *failed_start,
+                          struct extent_state **cached_state,
                           gfp_t mask)
 {
         struct extent_state *state;
@@ -712,6 +750,7 @@ hit_next:
                         goto out;
                 }
                 set_state_bits(tree, state, bits);
+                cache_state(state, cached_state);
                 merge_state(tree, state);
                 if (last_end == (u64)-1)
                         goto out;
@@ -758,6 +797,7 @@ hit_next:
                         goto out;
                 if (state->end <= end) {
                         set_state_bits(tree, state, bits);
+                        cache_state(state, cached_state);
                         merge_state(tree, state);
                         if (last_end == (u64)-1)
                                 goto out;
@@ -782,6 +822,7 @@ hit_next:
                         this_end = last_start - 1;
                 err = insert_state(tree, prealloc, start, this_end,
                                    bits);
+                cache_state(prealloc, cached_state);
                 prealloc = NULL;
                 BUG_ON(err == -EEXIST);
                 if (err)
@@ -805,6 +846,7 @@ hit_next:
                 BUG_ON(err == -EEXIST);
 
                 set_state_bits(tree, prealloc, bits);
+                cache_state(prealloc, cached_state);
                 merge_state(tree, prealloc);
                 prealloc = NULL;
                 goto out;
@@ -833,26 +875,27 @@ int set_extent_dirty(struct extent_io_tree *tree, u64 start, u64 end,
                      gfp_t mask)
 {
         return set_extent_bit(tree, start, end, EXTENT_DIRTY, 0, NULL,
-                              mask);
+                              NULL, mask);
 }
 
 int set_extent_ordered(struct extent_io_tree *tree, u64 start, u64 end,
                        gfp_t mask)
 {
-        return set_extent_bit(tree, start, end, EXTENT_ORDERED, 0, NULL, mask);
+        return set_extent_bit(tree, start, end, EXTENT_ORDERED, 0, NULL, NULL,
+                              mask);
 }
 
 int set_extent_bits(struct extent_io_tree *tree, u64 start, u64 end,
                     int bits, gfp_t mask)
 {
         return set_extent_bit(tree, start, end, bits, 0, NULL,
-                              mask);
+                              NULL, mask);
 }
 
 int clear_extent_bits(struct extent_io_tree *tree, u64 start, u64 end,
                       int bits, gfp_t mask)
 {
-        return clear_extent_bit(tree, start, end, bits, 0, 0, mask);
+        return clear_extent_bit(tree, start, end, bits, 0, 0, NULL, mask);
 }
 
 int set_extent_delalloc(struct extent_io_tree *tree, u64 start, u64 end,
@@ -860,46 +903,50 @@ int set_extent_delalloc(struct extent_io_tree *tree, u64 start, u64 end,
 {
         return set_extent_bit(tree, start, end,
                               EXTENT_DELALLOC | EXTENT_DIRTY | EXTENT_UPTODATE,
-                              0, NULL, mask);
+                              0, NULL, NULL, mask);