btrfs: fix race between quota disable and quota assign ioctls
[platform/kernel/linux-rpi.git] / fs / btrfs / relocation.c
index b70be2a..becf339 100644 (file)
@@ -24,6 +24,7 @@
 #include "block-group.h"
 #include "backref.h"
 #include "misc.h"
+#include "subpage.h"
 
 /*
  * Relocation overview
@@ -1146,7 +1147,8 @@ int replace_file_extents(struct btrfs_trans_handle *trans,
                                       num_bytes, parent);
                ref.real_root = root->root_key.objectid;
                btrfs_init_data_ref(&ref, btrfs_header_owner(leaf),
-                                   key.objectid, key.offset);
+                                   key.objectid, key.offset,
+                                   root->root_key.objectid, false);
                ret = btrfs_inc_extent_ref(trans, &ref);
                if (ret) {
                        btrfs_abort_transaction(trans, ret);
@@ -1157,7 +1159,8 @@ int replace_file_extents(struct btrfs_trans_handle *trans,
                                       num_bytes, parent);
                ref.real_root = root->root_key.objectid;
                btrfs_init_data_ref(&ref, btrfs_header_owner(leaf),
-                                   key.objectid, key.offset);
+                                   key.objectid, key.offset,
+                                   root->root_key.objectid, false);
                ret = btrfs_free_extent(trans, &ref);
                if (ret) {
                        btrfs_abort_transaction(trans, ret);
@@ -1323,7 +1326,9 @@ again:
                btrfs_release_path(path);
 
                path->lowest_level = level;
+               set_bit(BTRFS_ROOT_RESET_LOCKDEP_CLASS, &src->state);
                ret = btrfs_search_slot(trans, src, &key, path, 0, 1);
+               clear_bit(BTRFS_ROOT_RESET_LOCKDEP_CLASS, &src->state);
                path->lowest_level = 0;
                if (ret) {
                        if (ret > 0)
@@ -1367,7 +1372,8 @@ again:
                btrfs_init_generic_ref(&ref, BTRFS_ADD_DELAYED_REF, old_bytenr,
                                       blocksize, path->nodes[level]->start);
                ref.skip_qgroup = true;
-               btrfs_init_tree_ref(&ref, level - 1, src->root_key.objectid);
+               btrfs_init_tree_ref(&ref, level - 1, src->root_key.objectid,
+                                   0, true);
                ret = btrfs_inc_extent_ref(trans, &ref);
                if (ret) {
                        btrfs_abort_transaction(trans, ret);
@@ -1376,7 +1382,8 @@ again:
                btrfs_init_generic_ref(&ref, BTRFS_ADD_DELAYED_REF, new_bytenr,
                                       blocksize, 0);
                ref.skip_qgroup = true;
-               btrfs_init_tree_ref(&ref, level - 1, dest->root_key.objectid);
+               btrfs_init_tree_ref(&ref, level - 1, dest->root_key.objectid, 0,
+                                   true);
                ret = btrfs_inc_extent_ref(trans, &ref);
                if (ret) {
                        btrfs_abort_transaction(trans, ret);
@@ -1385,7 +1392,8 @@ again:
 
                btrfs_init_generic_ref(&ref, BTRFS_DROP_DELAYED_REF, new_bytenr,
                                       blocksize, path->nodes[level]->start);
-               btrfs_init_tree_ref(&ref, level - 1, src->root_key.objectid);
+               btrfs_init_tree_ref(&ref, level - 1, src->root_key.objectid,
+                                   0, true);
                ref.skip_qgroup = true;
                ret = btrfs_free_extent(trans, &ref);
                if (ret) {
@@ -1395,7 +1403,8 @@ again:
 
                btrfs_init_generic_ref(&ref, BTRFS_DROP_DELAYED_REF, old_bytenr,
                                       blocksize, 0);
-               btrfs_init_tree_ref(&ref, level - 1, dest->root_key.objectid);
+               btrfs_init_tree_ref(&ref, level - 1, dest->root_key.objectid,
+                                   0, true);
                ref.skip_qgroup = true;
                ret = btrfs_free_extent(trans, &ref);
                if (ret) {
@@ -2474,7 +2483,8 @@ static int do_relocation(struct btrfs_trans_handle *trans,
                                               upper->eb->start);
                        ref.real_root = root->root_key.objectid;
                        btrfs_init_tree_ref(&ref, node->level,
-                                           btrfs_header_owner(upper->eb));
+                                           btrfs_header_owner(upper->eb),
+                                           root->root_key.objectid, false);
                        ret = btrfs_inc_extent_ref(trans, &ref);
                        if (!ret)
                                ret = btrfs_drop_subtree(trans, root, eb,
@@ -2690,8 +2700,12 @@ static int relocate_tree_block(struct btrfs_trans_handle *trans,
                        list_add_tail(&node->list, &rc->backref_cache.changed);
                } else {
                        path->lowest_level = node->level;
+                       if (root == root->fs_info->chunk_root)
+                               btrfs_reserve_chunk_metadata(trans, false);
                        ret = btrfs_search_slot(trans, root, key, path, 0, 1);
                        btrfs_release_path(path);
+                       if (root == root->fs_info->chunk_root)
+                               btrfs_trans_release_chunk_metadata(trans);
                        if (ret > 0)
                                ret = 0;
                }
@@ -2781,41 +2795,76 @@ static noinline_for_stack int prealloc_file_extent_cluster(
        u64 num_bytes;
        int nr;
        int ret = 0;
+       u64 i_size = i_size_read(&inode->vfs_inode);
        u64 prealloc_start = cluster->start - offset;
        u64 prealloc_end = cluster->end - offset;
        u64 cur_offset = prealloc_start;
 
-       BUG_ON(cluster->start != cluster->boundary[0]);
-       ret = btrfs_alloc_data_chunk_ondemand(inode,
-                                             prealloc_end + 1 - prealloc_start);
-       if (ret)
-               return ret;
-
        /*
-        * On a zoned filesystem, we cannot preallocate the file region.
-        * Instead, we dirty and fiemap_write the region.
+        * For subpage case, previous i_size may not be aligned to PAGE_SIZE.
+        * This means the range [i_size, PAGE_END + 1) is filled with zeros by
+        * btrfs_do_readpage() call of previously relocated file cluster.
+        *
+        * If the current cluster starts in the above range, btrfs_do_readpage()
+        * will skip the read, and relocate_one_page() will later writeback
+        * the padding zeros as new data, causing data corruption.
+        *
+        * Here we have to manually invalidate the range (i_size, PAGE_END + 1).
         */
-       if (btrfs_is_zoned(inode->root->fs_info)) {
-               struct btrfs_root *root = inode->root;
-               struct btrfs_trans_handle *trans;
-
-               end = cluster->end - offset + 1;
-               trans = btrfs_start_transaction(root, 1);
-               if (IS_ERR(trans))
-                       return PTR_ERR(trans);
-
-               inode->vfs_inode.i_ctime = current_time(&inode->vfs_inode);
-               i_size_write(&inode->vfs_inode, end);
-               ret = btrfs_update_inode(trans, root, inode);
-               if (ret) {
-                       btrfs_abort_transaction(trans, ret);
-                       btrfs_end_transaction(trans);
+       if (!IS_ALIGNED(i_size, PAGE_SIZE)) {
+               struct address_space *mapping = inode->vfs_inode.i_mapping;
+               struct btrfs_fs_info *fs_info = inode->root->fs_info;
+               const u32 sectorsize = fs_info->sectorsize;
+               struct page *page;
+
+               ASSERT(sectorsize < PAGE_SIZE);
+               ASSERT(IS_ALIGNED(i_size, sectorsize));
+
+               /*
+                * Subpage can't handle page with DIRTY but without UPTODATE
+                * bit as it can lead to the following deadlock:
+                *
+                * btrfs_readpage()
+                * | Page already *locked*
+                * |- btrfs_lock_and_flush_ordered_range()
+                *    |- btrfs_start_ordered_extent()
+                *       |- extent_write_cache_pages()
+                *          |- lock_page()
+                *             We try to lock the page we already hold.
+                *
+                * Here we just writeback the whole data reloc inode, so that
+                * we will be ensured to have no dirty range in the page, and
+                * are safe to clear the uptodate bits.
+                *
+                * This shouldn't cause too much overhead, as we need to write
+                * the data back anyway.
+                */
+               ret = filemap_write_and_wait(mapping);
+               if (ret < 0)
                        return ret;
-               }
 
-               return btrfs_end_transaction(trans);
+               clear_extent_bits(&inode->io_tree, i_size,
+                                 round_up(i_size, PAGE_SIZE) - 1,
+                                 EXTENT_UPTODATE);
+               page = find_lock_page(mapping, i_size >> PAGE_SHIFT);
+               /*
+                * If page is freed we don't need to do anything then, as we
+                * will re-read the whole page anyway.
+                */
+               if (page) {
+                       btrfs_subpage_clear_uptodate(fs_info, page, i_size,
+                                       round_up(i_size, PAGE_SIZE) - i_size);
+                       unlock_page(page);
+                       put_page(page);
+               }
        }
 
+       BUG_ON(cluster->start != cluster->boundary[0]);
+       ret = btrfs_alloc_data_chunk_ondemand(inode,
+                                             prealloc_end + 1 - prealloc_start);
+       if (ret)
+               return ret;
+
        btrfs_inode_lock(&inode->vfs_inode, 0);
        for (nr = 0; nr < cluster->nr; nr++) {
                start = cluster->boundary[nr] - offset;
@@ -2876,28 +2925,158 @@ int setup_extent_mapping(struct inode *inode, u64 start, u64 end,
 }
 
 /*
- * Allow error injection to test balance cancellation
+ * Allow error injection to test balance/relocation cancellation
  */
 noinline int btrfs_should_cancel_balance(struct btrfs_fs_info *fs_info)
 {
        return atomic_read(&fs_info->balance_cancel_req) ||
+               atomic_read(&fs_info->reloc_cancel_req) ||
                fatal_signal_pending(current);
 }
 ALLOW_ERROR_INJECTION(btrfs_should_cancel_balance, TRUE);
 
-static int relocate_file_extent_cluster(struct inode *inode,
-                                       struct file_extent_cluster *cluster)
+static u64 get_cluster_boundary_end(struct file_extent_cluster *cluster,
+                                   int cluster_nr)
+{
+       /* Last extent, use cluster end directly */
+       if (cluster_nr >= cluster->nr - 1)
+               return cluster->end;
+
+       /* Use next boundary start*/
+       return cluster->boundary[cluster_nr + 1] - 1;
+}
+
+static int relocate_one_page(struct inode *inode, struct file_ra_state *ra,
+                            struct file_extent_cluster *cluster,
+                            int *cluster_nr, unsigned long page_index)
 {
        struct btrfs_fs_info *fs_info = btrfs_sb(inode->i_sb);
+       u64 offset = BTRFS_I(inode)->index_cnt;
+       const unsigned long last_index = (cluster->end - offset) >> PAGE_SHIFT;
+       gfp_t mask = btrfs_alloc_write_mask(inode->i_mapping);
+       struct page *page;
        u64 page_start;
        u64 page_end;
+       u64 cur;
+       int ret;
+
+       ASSERT(page_index <= last_index);
+       page = find_lock_page(inode->i_mapping, page_index);
+       if (!page) {
+               page_cache_sync_readahead(inode->i_mapping, ra, NULL,
+                               page_index, last_index + 1 - page_index);
+               page = find_or_create_page(inode->i_mapping, page_index, mask);
+               if (!page)
+                       return -ENOMEM;
+       }
+       ret = set_page_extent_mapped(page);
+       if (ret < 0)
+               goto release_page;
+
+       if (PageReadahead(page))
+               page_cache_async_readahead(inode->i_mapping, ra, NULL, page,
+                                  page_index, last_index + 1 - page_index);
+
+       if (!PageUptodate(page)) {
+               btrfs_readpage(NULL, page);
+               lock_page(page);
+               if (!PageUptodate(page)) {
+                       ret = -EIO;
+                       goto release_page;
+               }
+       }
+
+       page_start = page_offset(page);
+       page_end = page_start + PAGE_SIZE - 1;
+
+       /*
+        * Start from the cluster, as for subpage case, the cluster can start
+        * inside the page.
+        */
+       cur = max(page_start, cluster->boundary[*cluster_nr] - offset);
+       while (cur <= page_end) {
+               u64 extent_start = cluster->boundary[*cluster_nr] - offset;
+               u64 extent_end = get_cluster_boundary_end(cluster,
+                                               *cluster_nr) - offset;
+               u64 clamped_start = max(page_start, extent_start);
+               u64 clamped_end = min(page_end, extent_end);
+               u32 clamped_len = clamped_end + 1 - clamped_start;
+
+               /* Reserve metadata for this range */
+               ret = btrfs_delalloc_reserve_metadata(BTRFS_I(inode),
+                                                     clamped_len);
+               if (ret)
+                       goto release_page;
+
+               /* Mark the range delalloc and dirty for later writeback */
+               lock_extent(&BTRFS_I(inode)->io_tree, clamped_start, clamped_end);
+               ret = btrfs_set_extent_delalloc(BTRFS_I(inode), clamped_start,
+                                               clamped_end, 0, NULL);
+               if (ret) {
+                       clear_extent_bits(&BTRFS_I(inode)->io_tree,
+                                       clamped_start, clamped_end,
+                                       EXTENT_LOCKED | EXTENT_BOUNDARY);
+                       btrfs_delalloc_release_metadata(BTRFS_I(inode),
+                                                       clamped_len, true);
+                       btrfs_delalloc_release_extents(BTRFS_I(inode),
+                                                      clamped_len);
+                       goto release_page;
+               }
+               btrfs_page_set_dirty(fs_info, page, clamped_start, clamped_len);
+
+               /*
+                * Set the boundary if it's inside the page.
+                * Data relocation requires the destination extents to have the
+                * same size as the source.
+                * EXTENT_BOUNDARY bit prevents current extent from being merged
+                * with previous extent.
+                */
+               if (in_range(cluster->boundary[*cluster_nr] - offset,
+                            page_start, PAGE_SIZE)) {
+                       u64 boundary_start = cluster->boundary[*cluster_nr] -
+                                               offset;
+                       u64 boundary_end = boundary_start +
+                                          fs_info->sectorsize - 1;
+
+                       set_extent_bits(&BTRFS_I(inode)->io_tree,
+                                       boundary_start, boundary_end,
+                                       EXTENT_BOUNDARY);
+               }
+               unlock_extent(&BTRFS_I(inode)->io_tree, clamped_start, clamped_end);
+               btrfs_delalloc_release_extents(BTRFS_I(inode), clamped_len);
+               cur += clamped_len;
+
+               /* Crossed extent end, go to next extent */
+               if (cur >= extent_end) {
+                       (*cluster_nr)++;
+                       /* Just finished the last extent of the cluster, exit. */
+                       if (*cluster_nr >= cluster->nr)
+                               break;
+               }
+       }
+       unlock_page(page);
+       put_page(page);
+
+       balance_dirty_pages_ratelimited(inode->i_mapping);
+       btrfs_throttle(fs_info);
+       if (btrfs_should_cancel_balance(fs_info))
+               ret = -ECANCELED;
+       return ret;
+
+release_page:
+       unlock_page(page);
+       put_page(page);
+       return ret;
+}
+
+static int relocate_file_extent_cluster(struct inode *inode,
+                                       struct file_extent_cluster *cluster)
+{
        u64 offset = BTRFS_I(inode)->index_cnt;
        unsigned long index;
        unsigned long last_index;
-       struct page *page;
        struct file_ra_state *ra;
-       gfp_t mask = btrfs_alloc_write_mask(inode->i_mapping);
-       int nr = 0;
+       int cluster_nr = 0;
        int ret = 0;
 
        if (!cluster->nr)
@@ -2918,109 +3097,12 @@ static int relocate_file_extent_cluster(struct inode *inode,
        if (ret)
                goto out;
 
-       index = (cluster->start - offset) >> PAGE_SHIFT;
        last_index = (cluster->end - offset) >> PAGE_SHIFT;
-       while (index <= last_index) {
-               ret = btrfs_delalloc_reserve_metadata(BTRFS_I(inode),
-                               PAGE_SIZE);
-               if (ret)
-                       goto out;
-
-               page = find_lock_page(inode->i_mapping, index);
-               if (!page) {
-                       page_cache_sync_readahead(inode->i_mapping,
-                                                 ra, NULL, index,
-                                                 last_index + 1 - index);
-                       page = find_or_create_page(inode->i_mapping, index,
-                                                  mask);
-                       if (!page) {
-                               btrfs_delalloc_release_metadata(BTRFS_I(inode),
-                                                       PAGE_SIZE, true);
-                               btrfs_delalloc_release_extents(BTRFS_I(inode),
-                                                       PAGE_SIZE);
-                               ret = -ENOMEM;
-                               goto out;
-                       }
-               }
-               ret = set_page_extent_mapped(page);
-               if (ret < 0) {
-                       btrfs_delalloc_release_metadata(BTRFS_I(inode),
-                                                       PAGE_SIZE, true);
-                       btrfs_delalloc_release_extents(BTRFS_I(inode), PAGE_SIZE);
-                       unlock_page(page);
-                       put_page(page);
-                       goto out;
-               }
-
-               if (PageReadahead(page)) {
-                       page_cache_async_readahead(inode->i_mapping,
-                                                  ra, NULL, page, index,
-                                                  last_index + 1 - index);
-               }
-
-               if (!PageUptodate(page)) {
-                       btrfs_readpage(NULL, page);
-                       lock_page(page);
-                       if (!PageUptodate(page)) {
-                               unlock_page(page);
-                               put_page(page);
-                               btrfs_delalloc_release_metadata(BTRFS_I(inode),
-                                                       PAGE_SIZE, true);
-                               btrfs_delalloc_release_extents(BTRFS_I(inode),
-                                                              PAGE_SIZE);
-                               ret = -EIO;
-                               goto out;
-                       }
-               }
-
-               page_start = page_offset(page);
-               page_end = page_start + PAGE_SIZE - 1;
-
-               lock_extent(&BTRFS_I(inode)->io_tree, page_start, page_end);
-
-               if (nr < cluster->nr &&
-                   page_start + offset == cluster->boundary[nr]) {
-                       set_extent_bits(&BTRFS_I(inode)->io_tree,
-                                       page_start, page_end,
-                                       EXTENT_BOUNDARY);
-                       nr++;
-               }
-
-               ret = btrfs_set_extent_delalloc(BTRFS_I(inode), page_start,
-                                               page_end, 0, NULL);
-               if (ret) {
-                       unlock_page(page);
-                       put_page(page);
-                       btrfs_delalloc_release_metadata(BTRFS_I(inode),
-                                                        PAGE_SIZE, true);
-                       btrfs_delalloc_release_extents(BTRFS_I(inode),
-                                                      PAGE_SIZE);
-
-                       clear_extent_bits(&BTRFS_I(inode)->io_tree,
-                                         page_start, page_end,
-                                         EXTENT_LOCKED | EXTENT_BOUNDARY);
-                       goto out;
-
-               }
-               set_page_dirty(page);
-
-               unlock_extent(&BTRFS_I(inode)->io_tree,
-                             page_start, page_end);
-               unlock_page(page);
-               put_page(page);
-
-               index++;
-               btrfs_delalloc_release_extents(BTRFS_I(inode), PAGE_SIZE);
-               balance_dirty_pages_ratelimited(inode->i_mapping);
-               btrfs_throttle(fs_info);
-               if (btrfs_should_cancel_balance(fs_info)) {
-                       ret = -ECANCELED;
-                       goto out;
-               }
-       }
-       WARN_ON(nr != cluster->nr);
-       if (btrfs_is_zoned(fs_info) && !ret)
-               ret = btrfs_wait_ordered_range(inode, 0, (u64)-1);
+       for (index = (cluster->start - offset) >> PAGE_SHIFT;
+            index <= last_index && !ret; index++)
+               ret = relocate_one_page(inode, ra, cluster, &cluster_nr, index);
+       if (ret == 0)
+               WARN_ON(cluster_nr != cluster->nr);
 out:
        kfree(ra);
        return ret;
@@ -3496,7 +3578,12 @@ int prepare_to_relocate(struct reloc_control *rc)
                 */
                return PTR_ERR(trans);
        }
-       return btrfs_commit_transaction(trans);
+
+       ret = btrfs_commit_transaction(trans);
+       if (ret)
+               unset_reloc_control(rc);
+
+       return ret;
 }
 
 static noinline_for_stack int relocate_block_group(struct reloc_control *rc)
@@ -3673,12 +3760,8 @@ static int __insert_orphan_inode(struct btrfs_trans_handle *trans,
        struct btrfs_path *path;
        struct btrfs_inode_item *item;
        struct extent_buffer *leaf;
-       u64 flags = BTRFS_INODE_NOCOMPRESS | BTRFS_INODE_PREALLOC;
        int ret;
 
-       if (btrfs_is_zoned(trans->fs_info))
-               flags &= ~BTRFS_INODE_PREALLOC;
-
        path = btrfs_alloc_path();
        if (!path)
                return -ENOMEM;
@@ -3693,7 +3776,8 @@ static int __insert_orphan_inode(struct btrfs_trans_handle *trans,
        btrfs_set_inode_generation(leaf, item, 1);
        btrfs_set_inode_size(leaf, item, 0);
        btrfs_set_inode_mode(leaf, item, S_IFREG | 0600);
-       btrfs_set_inode_flags(leaf, item, flags);
+       btrfs_set_inode_flags(leaf, item, BTRFS_INODE_NOCOMPRESS |
+                                         BTRFS_INODE_PREALLOC);
        btrfs_mark_buffer_dirty(leaf);
 out:
        btrfs_free_path(path);
@@ -3780,6 +3864,47 @@ out:
        return inode;
 }
 
+/*
+ * Mark start of chunk relocation that is cancellable. Check if the cancellation
+ * has been requested meanwhile and don't start in that case.
+ *
+ * Return:
+ *   0             success
+ *   -EINPROGRESS  operation is already in progress, that's probably a bug
+ *   -ECANCELED    cancellation request was set before the operation started
+ */
+static int reloc_chunk_start(struct btrfs_fs_info *fs_info)
+{
+       if (test_and_set_bit(BTRFS_FS_RELOC_RUNNING, &fs_info->flags)) {
+               /* This should not happen */
+               btrfs_err(fs_info, "reloc already running, cannot start");
+               return -EINPROGRESS;
+       }
+
+       if (atomic_read(&fs_info->reloc_cancel_req) > 0) {
+               btrfs_info(fs_info, "chunk relocation canceled on start");
+               /*
+                * On cancel, clear all requests but let the caller mark
+                * the end after cleanup operations.
+                */
+               atomic_set(&fs_info->reloc_cancel_req, 0);
+               return -ECANCELED;
+       }
+       return 0;
+}
+
+/*
+ * Mark end of chunk relocation that is cancellable and wake any waiters.
+ */
+static void reloc_chunk_end(struct btrfs_fs_info *fs_info)
+{
+       /* Requested after start, clear bit first so any waiters can continue */
+       if (atomic_read(&fs_info->reloc_cancel_req) > 0)
+               btrfs_info(fs_info, "chunk relocation canceled during operation");
+       clear_and_wake_up_bit(BTRFS_FS_RELOC_RUNNING, &fs_info->flags);
+       atomic_set(&fs_info->reloc_cancel_req, 0);
+}
+
 static struct reloc_control *alloc_reloc_control(struct btrfs_fs_info *fs_info)
 {
        struct reloc_control *rc;
@@ -3847,6 +3972,19 @@ int btrfs_relocate_block_group(struct btrfs_fs_info *fs_info, u64 group_start)
        int rw = 0;
        int err = 0;
 
+       /*
+        * This only gets set if we had a half-deleted snapshot on mount.  We
+        * cannot allow relocation to start while we're still trying to clean up
+        * these pending deletions.
+        */
+       ret = wait_on_bit(&fs_info->flags, BTRFS_FS_UNFINISHED_DROPS, TASK_INTERRUPTIBLE);
+       if (ret)
+               return ret;
+
+       /* We may have been woken up by close_ctree, so bail if we're closing. */
+       if (btrfs_fs_closing(fs_info))
+               return -EINTR;
+
        bg = btrfs_lookup_block_group(fs_info, group_start);
        if (!bg)
                return -ENOENT;
@@ -3862,6 +4000,12 @@ int btrfs_relocate_block_group(struct btrfs_fs_info *fs_info, u64 group_start)
                return -ENOMEM;
        }
 
+       ret = reloc_chunk_start(fs_info);
+       if (ret < 0) {
+               err = ret;
+               goto out_put_bg;
+       }
+
        rc->extent_root = extent_root;
        rc->block_group = bg;
 
@@ -3952,7 +4096,9 @@ out:
        if (err && rw)
                btrfs_dec_block_group_ro(rc->block_group);
        iput(rc->data_inode);
-       btrfs_put_block_group(rc->block_group);
+out_put_bg:
+       btrfs_put_block_group(bg);
+       reloc_chunk_end(fs_info);
        free_reloc_control(rc);
        return err;
 }
@@ -4073,6 +4219,12 @@ int btrfs_recover_relocation(struct btrfs_root *root)
                goto out;
        }
 
+       ret = reloc_chunk_start(fs_info);
+       if (ret < 0) {
+               err = ret;
+               goto out_end;
+       }
+
        rc->extent_root = fs_info->extent_root;
 
        set_reloc_control(rc);
@@ -4137,6 +4289,8 @@ out_clean:
                err = ret;
 out_unset:
        unset_reloc_control(rc);
+out_end:
+       reloc_chunk_end(fs_info);
        free_reloc_control(rc);
 out:
        free_reloc_roots(&reloc_roots);
@@ -4219,8 +4373,7 @@ int btrfs_reloc_cow_block(struct btrfs_trans_handle *trans,
        if (!rc)
                return 0;
 
-       BUG_ON(rc->stage == UPDATE_DATA_PTRS &&
-              root->root_key.objectid == BTRFS_DATA_RELOC_TREE_OBJECTID);
+       BUG_ON(rc->stage == UPDATE_DATA_PTRS && btrfs_is_data_reloc_root(root));
 
        level = btrfs_header_level(buf);
        if (btrfs_header_generation(buf) <=