diff --git a/fs/btrfs/extent-tree.c b/fs/btrfs/extent-tree.c
index 5484256169dd..4c76677a54a9 100644
--- a/fs/btrfs/extent-tree.c
+++ b/fs/btrfs/extent-tree.c
@@ -7328,6 +7328,8 @@ struct find_free_extent_ctl {
 
 	/* Found result */
 	u64 found_offset;
+
+	struct btrfs_priority_tree *priority_tree;
 };
 
 
@@ -7692,6 +7694,10 @@ static int find_free_extent_update_loop(struct btrfs_fs_info *fs_info,
 	return -ENOSPC;
 }
 
+static inline bool
+is_priority_alloc_enabled(struct btrfs_fs_info *fs_info);
+
+static long compute_block_group_priority(struct btrfs_block_group_cache *bg);
 /*
  * walks the btree of allocated extents and find a hole of a given size.
  * The key ins is changed to record the hole:
@@ -7729,6 +7735,7 @@ static noinline int find_free_extent(struct btrfs_fs_info *fs_info,
 	struct btrfs_space_info *space_info;
 	bool use_cluster = true;
 	bool full_search = false;
+	bool use_priority = is_priority_alloc_enabled(fs_info);
 
 	WARN_ON(num_bytes < fs_info->sectorsize);
 
@@ -7744,6 +7751,7 @@ static noinline int find_free_extent(struct btrfs_fs_info *fs_info,
 	ffe_ctl.have_caching_bg = false;
 	ffe_ctl.orig_have_caching_bg = false;
 	ffe_ctl.found_offset = 0;
+	ffe_ctl.priority_tree = NULL;
 
 	ins->type = BTRFS_EXTENT_ITEM_KEY;
 	ins->objectid = 0;
@@ -7813,40 +7821,82 @@ static noinline int find_free_extent(struct btrfs_fs_info *fs_info,
 		 */
 		if (block_group && block_group_bits(block_group, flags) &&
 		    block_group->cached != BTRFS_CACHE_NO) {
-			down_read(&space_info->groups_sem);
-			if (list_empty(&block_group->list) ||
-			    block_group->ro) {
-				/*
-				 * someone is removing this block group,
-				 * we can't jump into the have_block_group
-				 * target because our list pointers are not
-				 * valid
-				 */
-				btrfs_put_block_group(block_group);
-				up_read(&space_info->groups_sem);
+			if (use_priority) {
+				spin_lock(&block_group->lock);
+				if (block_group->priority ==
+						    PRIORITY_BG_DELETED ||
+				    block_group->priority ==
+						    PRIORITY_BG_BUSY ||
+				    block_group->ro) {
+					spin_unlock(&block_group->lock);
+					btrfs_put_block_group(block_group);
+					goto search;
+				}
+
+				block_group->priority = PRIORITY_BG_BUSY;
+				spin_unlock(&block_group->lock);
+				ffe_ctl.priority_tree =
+					block_group->priority_tree;
+				down_read(&ffe_ctl.priority_tree->groups_sem);
 			} else {
-				ffe_ctl.index = btrfs_bg_flags_to_raid_index(
-						block_group->flags);
-				btrfs_lock_block_group(block_group, delalloc);
-				goto have_block_group;
+				down_read(&space_info->groups_sem);
+				if (list_empty(&block_group->list) ||
+				    block_group->ro) {
+					/*
+					 * someone is removing this block
+					 * group, we can't jump into the
+					 * have_block_group target because our
+					 * list pointers are not valid
+					 */
+					btrfs_put_block_group(block_group);
+					up_read(&space_info->groups_sem);
+					goto search;
+				}
 			}
+			ffe_ctl.index = btrfs_bg_flags_to_raid_index(
+				block_group->flags);
+			btrfs_lock_block_group(block_group, delalloc);
+			goto have_block_group;
 		} else if (block_group) {
 			btrfs_put_block_group(block_group);
 		}
 	}
 search:
+	ffe_ctl.priority_tree = NULL;
+	block_group = NULL;
 	ffe_ctl.have_caching_bg = false;
 	if (ffe_ctl.index == btrfs_bg_flags_to_raid_index(flags) ||
 	    ffe_ctl.index == 0)
 		full_search = true;
-	down_read(&space_info->groups_sem);
-	list_for_each_entry(block_group,
-			    &space_info->block_groups[ffe_ctl.index], list) {
+	if (!use_priority)
+		down_read(&space_info->groups_sem);
+	while (1) {
+		block_group = find_free_extent_search(space_info,
+				      ffe_ctl.index, block_group, use_priority,
+				      &ffe_ctl.priority_tree);
+
+		if (block_group == NULL)
+			break;
 		/* If the block group is read-only, we can skip it entirely. */
 		if (unlikely(block_group->ro))
 			continue;
 
 		btrfs_grab_block_group(block_group, delalloc);
+
+		if (use_priority) {
+			spin_lock(&block_group->lock);
+			/* someone the block group is unavaiable now */
+			if (block_group->priority == PRIORITY_BG_DELETED ||
+			    block_group->priority == PRIORITY_BG_UPDATING ||
+			    block_group->priority == PRIORITY_BG_BUSY) {
+				spin_unlock(&block_group->lock);
+				goto loop;
+			}
+
+			block_group->priority = PRIORITY_BG_BUSY;
+			spin_unlock(&block_group->lock);
+		}
+
 		ffe_ctl.search_start = block_group->key.objectid;
 
 		/*
@@ -7945,9 +7995,20 @@ static noinline int find_free_extent(struct btrfs_fs_info *fs_info,
 
 		trace_btrfs_reserve_extent(block_group, ffe_ctl.search_start,
 					   num_bytes);
-		btrfs_release_block_group(block_group, delalloc);
-		break;
+		goto out;
 loop:
+		if (use_priority) {
+			long priority;
+
+			spin_lock(&block_group->lock);
+			if (block_group->priority == PRIORITY_BG_BUSY) {
+				priority = compute_block_group_priority(
+					block_group);
+				block_group->priority = priority;
+			}
+			spin_unlock(&block_group->lock);
+		}
+
 		ffe_ctl.retry_clustered = false;
 		ffe_ctl.retry_unclustered = false;
 		BUG_ON(btrfs_bg_flags_to_raid_index(block_group->flags) !=
@@ -7955,8 +8016,16 @@ static noinline int find_free_extent(struct btrfs_fs_info *fs_info,
 		btrfs_release_block_group(block_group, delalloc);
 		cond_resched();
 	}
-	up_read(&space_info->groups_sem);
-
+	if (!use_priority)
+		up_read(&space_info->groups_sem);
+out:
+	if (ins->objectid) {
+		if (use_priority)
+			up_read(&ffe_ctl.priority_tree->groups_sem);
+		else
+			up_read(&space_info->groups_sem);
+		btrfs_release_block_group(block_group, delalloc);
+	}
 	ret = find_free_extent_update_loop(fs_info, last_ptr, ins, &ffe_ctl,
 					   full_search, use_cluster);
 	if (ret > 0)
@@ -10194,7 +10263,6 @@ static int check_chunk_block_group_mappings(struct btrfs_fs_info *fs_info)
 	return ret;
 }
 
-static long compute_block_group_priority(struct btrfs_block_group_cache *bg);
 static void add_block_group_priority(struct btrfs_block_group_cache *cache);
 int btrfs_read_block_groups(struct btrfs_fs_info *info)
 {
