3 files changed, 45 insertions, 14 deletions

diff --git a/fs/btrfs/extent_map.c b/fs/btrfs/extent_map.c
index 64d08f94485d..1874aee69c86 100644
--- a/fs/btrfs/extent_map.c
+++ b/fs/btrfs/extent_map.c
@@ -318,6 +318,20 @@ void clear_em_logging(struct extent_map_tree *tree, struct extent_map *em)
                 try_merge_map(tree, em);
 }
 
+static inline void setup_extent_mapping(struct extent_map_tree *tree,
+                                        struct extent_map *em,
+                                        int modified)
+{
+        atomic_inc(&em->refs);
+        em->mod_start = em->start;
+        em->mod_len = em->len;
+
+        if (modified)
+                list_move(&em->list, &tree->modified_extents);
+        else
+                try_merge_map(tree, em);
+}
+
 /**
  * add_extent_mapping - add new extent map to the extent tree
  * @tree:       tree to insert new map in
@@ -337,15 +351,7 @@ int add_extent_mapping(struct extent_map_tree *tree,
         if (ret)
                 goto out;
 
-        atomic_inc(&em->refs);
-
-        em->mod_start = em->start;
-        em->mod_len = em->len;
-
-        if (modified)
-                list_move(&em->list, &tree->modified_extents);
-        else
-                try_merge_map(tree, em);
+        setup_extent_mapping(tree, em, modified);
 out:
         return ret;
 }
@@ -432,3 +438,18 @@ int remove_extent_mapping(struct extent_map_tree *tree, struct extent_map *em)
         RB_CLEAR_NODE(&em->rb_node);
         return ret;
 }
+
+void replace_extent_mapping(struct extent_map_tree *tree,
+                            struct extent_map *cur,
+                            struct extent_map *new,
+                            int modified)
+{
+        WARN_ON(test_bit(EXTENT_FLAG_PINNED, &cur->flags));
+        ASSERT(extent_map_in_tree(cur));
+        if (!test_bit(EXTENT_FLAG_LOGGING, &cur->flags))
+                list_del_init(&cur->list);
+        rb_replace_node(&cur->rb_node, &new->rb_node, &tree->map);
+        RB_CLEAR_NODE(&cur->rb_node);
+
+        setup_extent_mapping(tree, new, modified);
+}
diff --git a/fs/btrfs/extent_map.h b/fs/btrfs/extent_map.h
index f0a645a14d6e..e7fd8a56a140 100644
--- a/fs/btrfs/extent_map.h
+++ b/fs/btrfs/extent_map.h
@@ -68,6 +68,10 @@ struct extent_map *lookup_extent_mapping(struct extent_map_tree *tree,
 int add_extent_mapping(struct extent_map_tree *tree,
                        struct extent_map *em, int modified);
 int remove_extent_mapping(struct extent_map_tree *tree, struct extent_map *em);
+void replace_extent_mapping(struct extent_map_tree *tree,
+                            struct extent_map *cur,
+                            struct extent_map *new,
+                            int modified);
 
 struct extent_map *alloc_extent_map(void);
 void free_extent_map(struct extent_map *em);
diff --git a/fs/btrfs/file.c b/fs/btrfs/file.c
index 762ca32bd988..31e48b947060 100644
--- a/fs/btrfs/file.c
+++ b/fs/btrfs/file.c
@@ -591,7 +591,6 @@ void btrfs_drop_extent_cache(struct inode *inode, u64 start, u64 end,
                 clear_bit(EXTENT_FLAG_PINNED, &em->flags);
                 clear_bit(EXTENT_FLAG_LOGGING, &flags);
                 modified = !list_empty(&em->list);
-                remove_extent_mapping(em_tree, em);
                 if (no_splits)
                         goto next;
 
@@ -622,8 +621,7 @@ void btrfs_drop_extent_cache(struct inode *inode, u64 start, u64 end,
                         split->bdev = em->bdev;
                         split->flags = flags;
                         split->compress_type = em->compress_type;
-                        ret = add_extent_mapping(em_tree, split, modified);
-                        BUG_ON(ret); /* Logic error */
+                        replace_extent_mapping(em_tree, em, split, modified);
                         free_extent_map(split);
                         split = split2;
                         split2 = NULL;
@@ -661,12 +659,20 @@ void btrfs_drop_extent_cache(struct inode *inode, u64 start, u64 end,
                                 split->orig_block_len = 0;
                         }
 
-                        ret = add_extent_mapping(em_tree, split, modified);
-                        BUG_ON(ret); /* Logic error */
+                        if (extent_map_in_tree(em)) {
+                                replace_extent_mapping(em_tree, em, split,
+                                                       modified);
+                        } else {
+                                ret = add_extent_mapping(em_tree, split,
+                                                         modified);
+                                ASSERT(ret == 0); /* Logic error */
+                        }
                         free_extent_map(split);
                         split = NULL;
                 }
 next:
+                if (extent_map_in_tree(em))
+                        remove_extent_mapping(em_tree, em);
                 write_unlock(&em_tree->lock);
 
                 /* once for us */