3 files changed, 160 insertions, 102 deletions

diff --git a/fs/btrfs/disk-io.c b/fs/btrfs/disk-io.c
index 97f99f830795..0a7d80da9c94 100644
--- a/fs/btrfs/disk-io.c
+++ b/fs/btrfs/disk-io.c
@@ -2796,6 +2796,7 @@ void btrfs_init_fs_info(struct btrfs_fs_info *fs_info)
 	mutex_init(&fs_info->unused_bg_unpin_mutex);
 	mutex_init(&fs_info->reclaim_bgs_lock);
 	mutex_init(&fs_info->reloc_mutex);
+	spin_lock_init(&fs_info->reloc_ctl_lock);
 	mutex_init(&fs_info->delalloc_root_mutex);
 	mutex_init(&fs_info->zoned_meta_io_lock);
 	mutex_init(&fs_info->zoned_data_reloc_io_lock);
diff --git a/fs/btrfs/fs.h b/fs/btrfs/fs.h
index da87292420fa..5f0cfb0b5466 100644
--- a/fs/btrfs/fs.h
+++ b/fs/btrfs/fs.h
@@ -657,6 +657,8 @@ struct btrfs_fs_info {
 	 * to protect us from the relocation code.
 	 */
 	struct mutex reloc_mutex;
+	/* Protects setting, clearing and getting fs_info->reloc_ctl. */
+	spinlock_t reloc_ctl_lock;
 
 	struct list_head trans_list;
 	struct list_head dead_roots;
diff --git a/fs/btrfs/relocation.c b/fs/btrfs/relocation.c
index 5f1200e69692..fb85bc8b345c 100644
--- a/fs/btrfs/relocation.c
+++ b/fs/btrfs/relocation.c
@@ -178,8 +178,101 @@ struct reloc_control {
 	bool create_reloc_tree;
 	bool merge_reloc_tree;
 	bool found_file_extent;
+
+	refcount_t refs;
 };
 
+static struct reloc_control *get_reloc_control(struct btrfs_fs_info *fs_info)
+{
+	struct reloc_control *rc;
+
+	/* Quick path, avoid lock contention on fs_info->reloc_ctl_lock. */
+	if (!data_race(fs_info->reloc_ctl))
+		return NULL;
+
+	spin_lock(&fs_info->reloc_ctl_lock);
+	rc = fs_info->reloc_ctl;
+	if (rc)
+		refcount_inc(&rc->refs);
+	spin_unlock(&fs_info->reloc_ctl_lock);
+
+	return rc;
+}
+
+static void __del_reloc_root(struct btrfs_root *root);
+
+static noinline_for_stack void free_reloc_roots(struct list_head *list)
+{
+	struct btrfs_root *reloc_root, *tmp;
+
+	list_for_each_entry_safe(reloc_root, tmp, list, root_list)
+		__del_reloc_root(reloc_root);
+}
+
+static void put_reloc_control(struct reloc_control *rc)
+{
+	if (refcount_dec_and_test(&rc->refs)) {
+		struct mapping_node *node, *tmp;
+
+		if (rc->extent_root)
+			ASSERT(rc->extent_root->fs_info->reloc_ctl != rc);
+
+		free_reloc_roots(&rc->reloc_roots);
+		rbtree_postorder_for_each_entry_safe(node, tmp,
+						     &rc->reloc_root_tree.rb_root,
+						     rb_node)
+			kfree(node);
+
+		if (rc->block_group)
+			btrfs_put_block_group(rc->block_group);
+
+		kfree(rc);
+	}
+}
+
+/* Helper to delete the 'address of tree root -> reloc tree' mapping. */
+static void __del_reloc_root(struct btrfs_root *root)
+{
+	struct btrfs_fs_info *fs_info = root->fs_info;
+	struct rb_node *rb_node;
+	struct mapping_node AUTO_KFREE(node);
+	struct reloc_control *rc;
+	bool put_ref = false;
+
+	rc = get_reloc_control(fs_info);
+	if (rc && root->node) {
+		spin_lock(&rc->reloc_root_tree.lock);
+		rb_node = rb_simple_search(&rc->reloc_root_tree.rb_root,
+					   root->commit_root->start);
+		if (rb_node) {
+			node = rb_entry(rb_node, struct mapping_node, rb_node);
+			rb_erase(&node->rb_node, &rc->reloc_root_tree.rb_root);
+			RB_CLEAR_NODE(&node->rb_node);
+		}
+		spin_unlock(&rc->reloc_root_tree.lock);
+		ASSERT(!node || (struct btrfs_root *)node->data == root);
+	}
+
+	/*
+	 * We only put the reloc root here if it's on the list.  There's a lot
+	 * of places where the pattern is to splice the rc->reloc_roots, process
+	 * the reloc roots, and then add the reloc root back onto
+	 * rc->reloc_roots.  If we call __del_reloc_root while it's off of the
+	 * list we don't want the reference being dropped, because the guy
+	 * messing with the list is in charge of the reference.
+	 */
+	spin_lock(&fs_info->trans_lock);
+	if (!list_empty(&root->root_list)) {
+		put_ref = true;
+		list_del_init(&root->root_list);
+	}
+	spin_unlock(&fs_info->trans_lock);
+	if (put_ref)
+		btrfs_put_root(root);
+	if (rc)
+		put_reloc_control(rc);
+}
+
 static void mark_block_processed(struct reloc_control *rc,
 				 struct btrfs_backref_node *node)
 {
@@ -475,12 +568,11 @@ out:
 /*
  * helper to add 'address of tree root -> reloc tree' mapping
  */
-static int __add_reloc_root(struct btrfs_root *root)
+static int __add_reloc_root(struct btrfs_root *root, struct reloc_control *rc)
 {
 	struct btrfs_fs_info *fs_info = root->fs_info;
 	struct rb_node *rb_node;
 	struct mapping_node *node;
-	struct reloc_control *rc = fs_info->reloc_ctl;
 
 	node = kmalloc_obj(*node, GFP_NOFS);
 	if (!node)
@@ -504,49 +596,6 @@ static int __add_reloc_root(struct btrfs_root *root)
 }
 
 /*
- * helper to delete the 'address of tree root -> reloc tree'
- * mapping
- */
-static void __del_reloc_root(struct btrfs_root *root)
-{
-	struct btrfs_fs_info *fs_info = root->fs_info;
-	struct rb_node *rb_node;
-	struct mapping_node AUTO_KFREE(node);
-	struct reloc_control *rc = fs_info->reloc_ctl;
-	bool put_ref = false;
-
-	if (rc && root->node) {
-		spin_lock(&rc->reloc_root_tree.lock);
-		rb_node = rb_simple_search(&rc->reloc_root_tree.rb_root,
-					   root->commit_root->start);
-		if (rb_node) {
-			node = rb_entry(rb_node, struct mapping_node, rb_node);
-			rb_erase(&node->rb_node, &rc->reloc_root_tree.rb_root);
-			RB_CLEAR_NODE(&node->rb_node);
-		}
-		spin_unlock(&rc->reloc_root_tree.lock);
-		ASSERT(!node || (struct btrfs_root *)node->data == root);
-	}
-
-	/*
-	 * We only put the reloc root here if it's on the list.  There's a lot
-	 * of places where the pattern is to splice the rc->reloc_roots, process
-	 * the reloc roots, and then add the reloc root back onto
-	 * rc->reloc_roots.  If we call __del_reloc_root while it's off of the
-	 * list we don't want the reference being dropped, because the guy
-	 * messing with the list is in charge of the reference.
-	 */
-	spin_lock(&fs_info->trans_lock);
-	if (!list_empty(&root->root_list)) {
-		put_ref = true;
-		list_del_init(&root->root_list);
-	}
-	spin_unlock(&fs_info->trans_lock);
-	if (put_ref)
-		btrfs_put_root(root);
-}
-
-/*
  * helper to update the 'address of tree root -> reloc tree'
  * mapping
  */
@@ -699,11 +748,12 @@ int btrfs_init_reloc_root(struct btrfs_trans_handle *trans,
 {
 	struct btrfs_fs_info *fs_info = root->fs_info;
 	struct btrfs_root *reloc_root;
-	struct reloc_control *rc = fs_info->reloc_ctl;
+	struct reloc_control *rc;
 	struct btrfs_block_rsv *rsv;
 	bool clear_rsv = false;
-	int ret;
+	int ret = 0;
 
+	rc = get_reloc_control(fs_info);
 	if (!rc)
 		return 0;
 
@@ -712,7 +762,7 @@ int btrfs_init_reloc_root(struct btrfs_trans_handle *trans,
 	 * create/update the dead reloc tree
 	 */
 	if (reloc_root_is_dead(root))
-		return 0;
+		goto out;
 
 	/*
 	 * This is subtle but important.  We do not do
@@ -723,9 +773,8 @@ int btrfs_init_reloc_root(struct btrfs_trans_handle *trans,
 	 * in.
 	 */
 	if (root->reloc_root) {
-		reloc_root = root->reloc_root;
-		btrfs_set_root_last_trans(reloc_root, trans->transid);
-		return 0;
+		btrfs_set_root_last_trans(root->reloc_root, trans->transid);
+		goto out;
 	}
 
 	/*
@@ -733,7 +782,7 @@ int btrfs_init_reloc_root(struct btrfs_trans_handle *trans,
 	 * reloc trees never need their own reloc tree.
 	 */
 	if (!rc->create_reloc_tree || btrfs_root_id(root) == BTRFS_TREE_RELOC_OBJECTID)
-		return 0;
+		goto out;
 
 	if (!trans->reloc_reserved) {
 		rsv = trans->block_rsv;
@@ -743,18 +792,23 @@ int btrfs_init_reloc_root(struct btrfs_trans_handle *trans,
 	reloc_root = create_reloc_root(trans, root, btrfs_root_id(root));
 	if (clear_rsv)
 		trans->block_rsv = rsv;
-	if (IS_ERR(reloc_root))
-		return PTR_ERR(reloc_root);
+	if (IS_ERR(reloc_root)) {
+		ret = PTR_ERR(reloc_root);
+		goto out;
+	}
 
-	ret = __add_reloc_root(reloc_root);
+	ret = __add_reloc_root(reloc_root, rc);
 	ASSERT(ret != -EEXIST);
 	if (ret) {
 		/* Pairs with create_reloc_root */
 		btrfs_put_root(reloc_root);
-		return ret;
+		goto out;
 	}
 	root->reloc_root = btrfs_grab_root(reloc_root);
-	return 0;
+out:
+	put_reloc_control(rc);
+
+	return ret;
 }
 
 /*
@@ -766,6 +820,7 @@ int btrfs_update_reloc_root(struct btrfs_trans_handle *trans,
 	struct btrfs_fs_info *fs_info = root->fs_info;
 	struct btrfs_root *reloc_root;
 	struct btrfs_root_item *root_item;
+	struct reloc_control *rc;
 	int ret;
 
 	if (!have_reloc_root(root))
@@ -781,9 +836,9 @@ int btrfs_update_reloc_root(struct btrfs_trans_handle *trans,
 	 */
 	btrfs_grab_root(reloc_root);
 
+	rc = get_reloc_control(fs_info);
 	/* root->reloc_root will stay until current relocation finished */
-	if (fs_info->reloc_ctl && fs_info->reloc_ctl->merge_reloc_tree &&
-	    btrfs_root_refs(root_item) == 0) {
+	if (rc && rc->merge_reloc_tree && btrfs_root_refs(root_item) == 0) {
 		set_bit(BTRFS_ROOT_DEAD_RELOC_TREE, &root->state);
 		/*
 		 * Mark the tree as dead before we change reloc_root so
@@ -803,6 +858,9 @@ int btrfs_update_reloc_root(struct btrfs_trans_handle *trans,
 	ret = btrfs_update_root(trans, fs_info->tree_root,
 				&reloc_root->root_key, root_item);
 	btrfs_put_root(reloc_root);
+	if (rc)
+		put_reloc_control(rc);
+
 	return ret;
 }
 
@@ -1808,15 +1866,6 @@ again:
 }
 
 static noinline_for_stack
-void free_reloc_roots(struct list_head *list)
-{
-	struct btrfs_root *reloc_root, *tmp;
-
-	list_for_each_entry_safe(reloc_root, tmp, list, root_list)
-		__del_reloc_root(reloc_root);
-}
-
-static noinline_for_stack
 void merge_reloc_roots(struct reloc_control *rc)
 {
 	struct btrfs_fs_info *fs_info = rc->extent_root->fs_info;
@@ -1920,7 +1969,7 @@ out:
 	 * do the reloc_dirty_list afterwards.  Meanwhile the root->reloc_root
 	 * will be cleaned up on unmount.
 	 *
-	 * The remaining nodes will be cleaned up by free_reloc_control.
+	 * The remaining nodes will be cleaned up by put_reloc_control().
 	 */
 }
 
@@ -3433,7 +3482,9 @@ static void set_reloc_control(struct reloc_control *rc)
 	struct btrfs_fs_info *fs_info = rc->extent_root->fs_info;
 
 	mutex_lock(&fs_info->reloc_mutex);
+	spin_lock(&fs_info->reloc_ctl_lock);
 	fs_info->reloc_ctl = rc;
+	spin_unlock(&fs_info->reloc_ctl_lock);
 	mutex_unlock(&fs_info->reloc_mutex);
 }
 
@@ -3442,7 +3493,9 @@ static void unset_reloc_control(struct reloc_control *rc)
 	struct btrfs_fs_info *fs_info = rc->extent_root->fs_info;
 
 	mutex_lock(&fs_info->reloc_mutex);
+	spin_lock(&fs_info->reloc_ctl_lock);
 	fs_info->reloc_ctl = NULL;
+	spin_unlock(&fs_info->reloc_ctl_lock);
 	mutex_unlock(&fs_info->reloc_mutex);
 }
 
@@ -3827,19 +3880,9 @@ static struct reloc_control *alloc_reloc_control(struct btrfs_fs_info *fs_info)
 	rc->reloc_root_tree.rb_root = RB_ROOT;
 	spin_lock_init(&rc->reloc_root_tree.lock);
 	btrfs_extent_io_tree_init(fs_info, &rc->processed_blocks, IO_TREE_RELOC_BLOCKS);
-	return rc;
-}
-
-static void free_reloc_control(struct reloc_control *rc)
-{
-	struct mapping_node *node, *tmp;
+	refcount_set(&rc->refs, 1);
 
-	free_reloc_roots(&rc->reloc_roots);
-	rbtree_postorder_for_each_entry_safe(node, tmp,
-			&rc->reloc_root_tree.rb_root, rb_node)
-		kfree(node);
-
-	kfree(rc);
+	return rc;
 }
 
 /*
@@ -5379,13 +5422,14 @@ 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)
-		goto out_put_bg;
-
 	rc->extent_root = extent_root;
+	/* Block group ref now owned by rc, put_reloc_control() will drop it. */
 	rc->block_group = bg;
 
+	ret = reloc_chunk_start(fs_info);
+	if (ret < 0)
+		goto out_put_rc;
+
 	ret = btrfs_inc_block_group_ro(rc->block_group, true);
 	if (ret)
 		goto out;
@@ -5453,9 +5497,8 @@ out:
 		iput(rc->data_inode);
 	btrfs_free_path(path);
 	reloc_chunk_end(fs_info);
-out_put_bg:
-	btrfs_put_block_group(bg);
-	free_reloc_control(rc);
+out_put_rc:
+	put_reloc_control(rc);
 	return ret;
 }
 
@@ -5610,7 +5653,7 @@ int btrfs_recover_relocation(struct btrfs_fs_info *fs_info)
 			goto out_unset;
 		}
 
-		ret = __add_reloc_root(reloc_root);
+		ret = __add_reloc_root(reloc_root, rc);
 		ASSERT(ret != -EEXIST);
 		if (ret) {
 			list_add_tail(&reloc_root->root_list, &reloc_roots);
@@ -5644,7 +5687,7 @@ out_unset:
 	unset_reloc_control(rc);
 	reloc_chunk_end(fs_info);
 out_end:
-	free_reloc_control(rc);
+	put_reloc_control(rc);
 out:
 	free_reloc_roots(&reloc_roots);
 
@@ -5728,7 +5771,7 @@ int btrfs_reloc_cow_block(struct btrfs_trans_handle *trans,
 	int level;
 	int ret = 0;
 
-	rc = fs_info->reloc_ctl;
+	rc = get_reloc_control(fs_info);
 	if (!rc)
 		return 0;
 
@@ -5753,7 +5796,8 @@ int btrfs_reloc_cow_block(struct btrfs_trans_handle *trans,
 			btrfs_err(fs_info,
 "bytenr %llu was found but our backref cache was expecting %llu or %llu",
 				  buf->start, node->bytenr, node->new_bytenr);
-			return -EUCLEAN;
+			ret = -EUCLEAN;
+			goto out;
 		}
 
 		btrfs_backref_drop_node_buffer(node);
@@ -5776,6 +5820,9 @@ int btrfs_reloc_cow_block(struct btrfs_trans_handle *trans,
 
 	if (level == 0 && first_cow && rc->stage == UPDATE_DATA_PTRS)
 		ret = replace_file_extents(trans, rc, root, cow);
+out:
+	put_reloc_control(rc);
+
 	return ret;
 }
 
@@ -5824,13 +5871,16 @@ int btrfs_reloc_post_snapshot(struct btrfs_trans_handle *trans,
 	struct btrfs_root *root = pending->root;
 	struct btrfs_root *reloc_root;
 	struct btrfs_root *new_root;
-	struct reloc_control *rc = root->fs_info->reloc_ctl;
-	int ret;
+	struct reloc_control *rc;
+	int ret = 0;
 
-	if (!rc || !have_reloc_root(root))
+	rc = get_reloc_control(trans->fs_info);
+	if (!rc)
 		return 0;
 
-	rc = root->fs_info->reloc_ctl;
+	if (!have_reloc_root(root))
+		goto out;
+
 	rc->merging_rsv_size += rc->nodes_relocated;
 
 	if (rc->merge_reloc_tree) {
@@ -5838,23 +5888,28 @@ int btrfs_reloc_post_snapshot(struct btrfs_trans_handle *trans,
 					      rc->block_rsv,
 					      rc->nodes_relocated, true);
 		if (ret)
-			return ret;
+			goto out;
 	}
 
 	new_root = pending->snap;
 	reloc_root = create_reloc_root(trans, root->reloc_root, btrfs_root_id(new_root));
-	if (IS_ERR(reloc_root))
-		return PTR_ERR(reloc_root);
+	if (IS_ERR(reloc_root)) {
+		ret = PTR_ERR(reloc_root);
+		goto out;
+	}
 
-	ret = __add_reloc_root(reloc_root);
+	ret = __add_reloc_root(reloc_root, rc);
 	ASSERT(ret != -EEXIST);
 	if (ret) {
 		/* Pairs with create_reloc_root */
 		btrfs_put_root(reloc_root);
-		return ret;
+		goto out;
 	}
 	new_root->reloc_root = btrfs_grab_root(reloc_root);
-	return 0;
+out:
+	put_reloc_control(rc);
+
+	return ret;
 }
 
 /*