1 files changed, 36 insertions, 9 deletions

diff --git a/fs/btrfs/extent_map.c b/fs/btrfs/extent_map.c
index 485cf0719b3c..010a287fbd71 100644
--- a/fs/btrfs/extent_map.c
+++ b/fs/btrfs/extent_map.c
@@ -204,10 +204,12 @@ static struct rb_node *tree_insert(struct rb_root *root, u64 offset,
 }
 
 static struct rb_node *__tree_search(struct rb_root *root, u64 offset,
-                                   struct rb_node **prev_ret)
+                                     struct rb_node **prev_ret,
+                                     struct rb_node **next_ret)
 {
         struct rb_node * n = root->rb_node;
         struct rb_node *prev = NULL;
+        struct rb_node *orig_prev = NULL;
         struct tree_entry *entry;
         struct tree_entry *prev_entry = NULL;
 
@@ -223,13 +225,25 @@ static struct rb_node *__tree_search(struct rb_root *root, u64 offset,
                 else
                         return n;
         }
-        if (!prev_ret)
-                return NULL;
-        while(prev && offset > prev_entry->end) {
-                prev = rb_next(prev);
+
+        if (prev_ret) {
+                orig_prev = prev;
+                while(prev && offset > prev_entry->end) {
+                        prev = rb_next(prev);
+                        prev_entry = rb_entry(prev, struct tree_entry, rb_node);
+                }
+                *prev_ret = prev;
+                prev = orig_prev;
+        }
+
+        if (next_ret) {
                 prev_entry = rb_entry(prev, struct tree_entry, rb_node);
+                while(prev && offset < prev_entry->start) {
+                        prev = rb_prev(prev);
+                        prev_entry = rb_entry(prev, struct tree_entry, rb_node);
+                }
+                *next_ret = prev;
         }
-        *prev_ret = prev;
         return NULL;
 }
 
@@ -237,7 +251,7 @@ static inline struct rb_node *tree_search(struct rb_root *root, u64 offset)
 {
         struct rb_node *prev;
         struct rb_node *ret;
-        ret = __tree_search(root, offset, &prev);
+        ret = __tree_search(root, offset, &prev, NULL);
         if (!ret)
                 return prev;
         return ret;
@@ -248,7 +262,7 @@ static int tree_delete(struct rb_root *root, u64 offset)
         struct rb_node *node;
         struct tree_entry *entry;
 
-        node = __tree_search(root, offset, NULL);
+        node = __tree_search(root, offset, NULL, NULL);
         if (!node)
                 return -ENOENT;
         entry = rb_entry(node, struct tree_entry, rb_node);
@@ -314,9 +328,21 @@ struct extent_map *lookup_extent_mapping(struct extent_map_tree *tree,
 {
         struct extent_map *em;
         struct rb_node *rb_node;
+        struct rb_node *prev = NULL;
+        struct rb_node *next = NULL;
 
         read_lock_irq(&tree->lock);
-        rb_node = tree_search(&tree->map, start);
+        rb_node = __tree_search(&tree->map, start, &prev, &next);
+        if (!rb_node && prev) {
+                em = rb_entry(prev, struct extent_map, rb_node);
+                if (em->start <= end && em->end >= start)
+                        goto found;
+        }
+        if (!rb_node && next) {
+                em = rb_entry(next, struct extent_map, rb_node);
+                if (em->start <= end && em->end >= start)
+                        goto found;
+        }
         if (!rb_node) {
                 em = NULL;
                 goto out;
@@ -330,6 +356,7 @@ struct extent_map *lookup_extent_mapping(struct extent_map_tree *tree,
                 em = NULL;
                 goto out;
         }
+found:
         atomic_inc(&em->refs);
 out:
         read_unlock_irq(&tree->lock);