1 files changed, 141 insertions, 168 deletions

diff --git a/fs/btrfs/scrub.c b/fs/btrfs/scrub.c
index e427cb7ee12c..ec57687c9a4d 100644
--- a/fs/btrfs/scrub.c
+++ b/fs/btrfs/scrub.c
@@ -66,7 +66,6 @@ struct scrub_ctx;
 struct scrub_recover {
         atomic_t                refs;
         struct btrfs_bio        *bbio;
-        u64                     *raid_map;
         u64                     map_length;
 };
 
@@ -80,7 +79,7 @@ struct scrub_page {
         u64                     logical;
         u64                     physical;
         u64                     physical_for_dev_replace;
-        atomic_t                ref_count;
+        atomic_t                refs;
         struct {
                 unsigned int    mirror_num:8;
                 unsigned int    have_csum:1;
@@ -113,7 +112,7 @@ struct scrub_block {
         struct scrub_page       *pagev[SCRUB_MAX_PAGES_PER_BLOCK];
         int                     page_count;
         atomic_t                outstanding_pages;
-        atomic_t                ref_count; /* free mem on transition to zero */
+        atomic_t                refs; /* free mem on transition to zero */
         struct scrub_ctx        *sctx;
         struct scrub_parity     *sparity;
         struct {
@@ -142,7 +141,7 @@ struct scrub_parity {
 
         int                     stripe_len;
 
-        atomic_t                ref_count;
+        atomic_t                refs;
 
         struct list_head        spages;
 
@@ -194,6 +193,15 @@ struct scrub_ctx {
          */
         struct btrfs_scrub_progress stat;
         spinlock_t              stat_lock;
+
+        /*
+         * Use a ref counter to avoid use-after-free issues. Scrub workers
+         * decrement bios_in_flight and workers_pending and then do a wakeup
+         * on the list_wait wait queue. We must ensure the main scrub task
+         * doesn't free the scrub context before or while the workers are
+         * doing the wakeup() call.
+         */
+        atomic_t                refs;
 };
 
 struct scrub_fixup_nodatasum {
@@ -236,10 +244,7 @@ static void scrub_pending_bio_dec(struct scrub_ctx *sctx);
 static void scrub_pending_trans_workers_inc(struct scrub_ctx *sctx);
 static void scrub_pending_trans_workers_dec(struct scrub_ctx *sctx);
 static int scrub_handle_errored_block(struct scrub_block *sblock_to_check);
-static int scrub_setup_recheck_block(struct scrub_ctx *sctx,
-                                     struct btrfs_fs_info *fs_info,
-                                     struct scrub_block *original_sblock,
-                                     u64 length, u64 logical,
+static int scrub_setup_recheck_block(struct scrub_block *original_sblock,
                                      struct scrub_block *sblocks_for_recheck);
 static void scrub_recheck_block(struct btrfs_fs_info *fs_info,
                                 struct scrub_block *sblock, int is_metadata,
@@ -251,8 +256,7 @@ static void scrub_recheck_block_checksum(struct btrfs_fs_info *fs_info,
                                          const u8 *csum, u64 generation,
                                          u16 csum_size);
 static int scrub_repair_block_from_good_copy(struct scrub_block *sblock_bad,
-                                             struct scrub_block *sblock_good,
-                                             int force_write);
+                                             struct scrub_block *sblock_good);
 static int scrub_repair_page_from_good_copy(struct scrub_block *sblock_bad,
                                             struct scrub_block *sblock_good,
                                             int page_num, int force_write);
@@ -302,10 +306,12 @@ static int copy_nocow_pages(struct scrub_ctx *sctx, u64 logical, u64 len,
 static void copy_nocow_pages_worker(struct btrfs_work *work);
 static void __scrub_blocked_if_needed(struct btrfs_fs_info *fs_info);
 static void scrub_blocked_if_needed(struct btrfs_fs_info *fs_info);
+static void scrub_put_ctx(struct scrub_ctx *sctx);
 
 
 static void scrub_pending_bio_inc(struct scrub_ctx *sctx)
 {
+        atomic_inc(&sctx->refs);
         atomic_inc(&sctx->bios_in_flight);
 }
 
@@ -313,6 +319,7 @@ static void scrub_pending_bio_dec(struct scrub_ctx *sctx)
 {
         atomic_dec(&sctx->bios_in_flight);
         wake_up(&sctx->list_wait);
+        scrub_put_ctx(sctx);
 }
 
 static void __scrub_blocked_if_needed(struct btrfs_fs_info *fs_info)
@@ -346,6 +353,7 @@ static void scrub_pending_trans_workers_inc(struct scrub_ctx *sctx)
 {
         struct btrfs_fs_info *fs_info = sctx->dev_root->fs_info;
 
+        atomic_inc(&sctx->refs);
         /*
          * increment scrubs_running to prevent cancel requests from
          * completing as long as a worker is running. we must also
@@ -388,6 +396,7 @@ static void scrub_pending_trans_workers_dec(struct scrub_ctx *sctx)
         atomic_dec(&sctx->workers_pending);
         wake_up(&fs_info->scrub_pause_wait);
         wake_up(&sctx->list_wait);
+        scrub_put_ctx(sctx);
 }
 
 static void scrub_free_csums(struct scrub_ctx *sctx)
@@ -433,6 +442,12 @@ static noinline_for_stack void scrub_free_ctx(struct scrub_ctx *sctx)
         kfree(sctx);
 }
 
+static void scrub_put_ctx(struct scrub_ctx *sctx)
+{
+        if (atomic_dec_and_test(&sctx->refs))
+                scrub_free_ctx(sctx);
+}
+
 static noinline_for_stack
 struct scrub_ctx *scrub_setup_ctx(struct btrfs_device *dev, int is_dev_replace)
 {
@@ -457,6 +472,7 @@ struct scrub_ctx *scrub_setup_ctx(struct btrfs_device *dev, int is_dev_replace)
         sctx = kzalloc(sizeof(*sctx), GFP_NOFS);
         if (!sctx)
                 goto nomem;
+        atomic_set(&sctx->refs, 1);
         sctx->is_dev_replace = is_dev_replace;
         sctx->pages_per_rd_bio = pages_per_rd_bio;
         sctx->curr = -1;
@@ -520,6 +536,7 @@ static int scrub_print_warning_inode(u64 inum, u64 offset, u64 root,
         struct inode_fs_paths *ipath = NULL;
         struct btrfs_root *local_root;
         struct btrfs_key root_key;
+        struct btrfs_key key;
 
         root_key.objectid = root;
         root_key.type = BTRFS_ROOT_ITEM_KEY;
@@ -530,7 +547,14 @@ static int scrub_print_warning_inode(u64 inum, u64 offset, u64 root,
                 goto err;
         }
 
-        ret = inode_item_info(inum, 0, local_root, swarn->path);
+        /*
+         * this makes the path point to (inum INODE_ITEM ioff)
+         */
+        key.objectid = inum;
+        key.type = BTRFS_INODE_ITEM_KEY;
+        key.offset = 0;
+
+        ret = btrfs_search_slot(NULL, local_root, &key, swarn->path, 0, 0);
         if (ret) {
                 btrfs_release_path(swarn->path);
                 goto err;
@@ -848,8 +872,7 @@ static inline void scrub_get_recover(struct scrub_recover *recover)
 static inline void scrub_put_recover(struct scrub_recover *recover)
 {
         if (atomic_dec_and_test(&recover->refs)) {
-                kfree(recover->bbio);
-                kfree(recover->raid_map);
+                btrfs_put_bbio(recover->bbio);
                 kfree(recover);
         }
 }
@@ -955,8 +978,7 @@ static int scrub_handle_errored_block(struct scrub_block *sblock_to_check)
         }
 
         /* setup the context, map the logical blocks and alloc the pages */
-        ret = scrub_setup_recheck_block(sctx, fs_info, sblock_to_check, length,
-                                        logical, sblocks_for_recheck);
+        ret = scrub_setup_recheck_block(sblock_to_check, sblocks_for_recheck);
         if (ret) {
                 spin_lock(&sctx->stat_lock);
                 sctx->stat.read_errors++;
@@ -1030,9 +1052,10 @@ static int scrub_handle_errored_block(struct scrub_block *sblock_to_check)
         if (!is_metadata && !have_csum) {
                 struct scrub_fixup_nodatasum *fixup_nodatasum;
 
-nodatasum_case:
                 WARN_ON(sctx->is_dev_replace);
 
+nodatasum_case:
+
                 /*
                  * !is_metadata and !have_csum, this means that the data
                  * might not be COW'ed, that it might be modified
@@ -1091,76 +1114,20 @@ nodatasum_case:
                     sblock_other->no_io_error_seen) {
                         if (sctx->is_dev_replace) {
                                 scrub_write_block_to_dev_replace(sblock_other);
+                                goto corrected_error;
                         } else {
-                                int force_write = is_metadata || have_csum;
-
                                 ret = scrub_repair_block_from_good_copy(
-                                                sblock_bad, sblock_other,
-                                                force_write);
+                                                sblock_bad, sblock_other);
+                                if (!ret)
+                                        goto corrected_error;
                         }
-                        if (0 == ret)
-                                goto corrected_error;
                 }
         }
 
-        /*
-         * for dev_replace, pick good pages and write to the target device.
-         */
-        if (sctx->is_dev_replace) {
-                success = 1;
-                for (page_num = 0; page_num < sblock_bad->page_count;
-                     page_num++) {
-                        int sub_success;
-
-                        sub_success = 0;
-                        for (mirror_index = 0;
-                             mirror_index < BTRFS_MAX_MIRRORS &&
-                             sblocks_for_recheck[mirror_index].page_count > 0;
-                             mirror_index++) {
-                                struct scrub_block *sblock_other =
-                                        sblocks_for_recheck + mirror_index;
-                                struct scrub_page *page_other =
-                                        sblock_other->pagev[page_num];
-
-                                if (!page_other->io_error) {
-                                        ret = scrub_write_page_to_dev_replace(
-                                                        sblock_other, page_num);
-                                        if (ret == 0) {
-                                                /* succeeded for this page */
-                                                sub_success = 1;
-                                                break;
-                                        } else {
-                                                btrfs_dev_replace_stats_inc(
-                                                        &sctx->dev_root->
-                                                        fs_info->dev_replace.
-                                                        num_write_errors);
-                                        }
-                                }
-                        }
-
-                        if (!sub_success) {
-                                /*
-                                 * did not find a mirror to fetch the page
-                                 * from. scrub_write_page_to_dev_replace()
-                                 * handles this case (page->io_error), by
-                                 * filling the block with zeros before
-                                 * submitting the write request
-                                 */
-                                success = 0;
-                                ret = scrub_write_page_to_dev_replace(
-                                                sblock_bad, page_num);
-                                if (ret)
-                                        btrfs_dev_replace_stats_inc(
-                                                &sctx->dev_root->fs_info->
-                                                dev_replace.num_write_errors);
-                        }
-                }
-
-                goto out;
-        }
+        if (sblock_bad->no_io_error_seen && !sctx->is_dev_replace)
+                goto did_not_correct_error;
 
         /*
-         * for regular scrub, repair those pages that are errored.
          * In case of I/O errors in the area that is supposed to be
          * repaired, continue by picking good copies of those pages.
          * Select the good pages from mirrors to rewrite bad pages from
@@ -1184,44 +1151,64 @@ nodatasum_case:
          * mirror, even if other 512 byte sectors in the same PAGE_SIZE
          * area are unreadable.
          */
-
-        /* can only fix I/O errors from here on */
-        if (sblock_bad->no_io_error_seen)
-                goto did_not_correct_error;
-
         success = 1;
-        for (page_num = 0; page_num < sblock_bad->page_count; page_num++) {
+        for (page_num = 0; page_num < sblock_bad->page_count;
+             page_num++) {
                 struct scrub_page *page_bad = sblock_bad->pagev[page_num];
+                struct scrub_block *sblock_other = NULL;
 
-                if (!page_bad->io_error)
+                /* skip no-io-error page in scrub */
+                if (!page_bad->io_error && !sctx->is_dev_replace)
                         continue;
 
-                for (mirror_index = 0;
-                     mirror_index < BTRFS_MAX_MIRRORS &&
-                     sblocks_for_recheck[mirror_index].page_count > 0;
-                     mirror_index++) {
-                        struct scrub_block *sblock_other = sblocks_for_recheck +
-                                                           mirror_index;
-                        struct scrub_page *page_other = sblock_other->pagev[
-                                                        page_num];
-
-                        if (!page_other->io_error) {
-                                ret = scrub_repair_page_from_good_copy(
-                                        sblock_bad, sblock_other, page_num, 0);
-                                if (0 == ret) {
-                                        page_bad->io_error = 0;
-                                        break; /* succeeded for this page */
+                /* try to find no-io-error page in mirrors */
+                if (page_bad->io_error) {
+                        for (mirror_index = 0;
+                             mirror_index < BTRFS_MAX_MIRRORS &&
+                             sblocks_for_recheck[mirror_index].page_count > 0;
+                             mirror_index++) {
+                                if (!sblocks_for_recheck[mirror_index].
+                                    pagev[page_num]->io_error) {
+                                        sblock_other = sblocks_for_recheck +
+                                                       mirror_index;
+                                        break;
                                 }
                         }
+                        if (!sblock_other)
+                                success = 0;
                 }
 
-                if (page_bad->io_error) {
-                        /* did not find a mirror to copy the page from */
-                        success = 0;
+                if (sctx->is_dev_replace) {
+                        /*
+                         * did not find a mirror to fetch the page
+                         * from. scrub_write_page_to_dev_replace()
+                         * handles this case (page->io_error), by
+                         * filling the block with zeros before
+                         * submitting the write request
+                         */
+                        if (!sblock_other)
+                                sblock_other = sblock_bad;
+
+                        if (scrub_write_page_to_dev_replace(sblock_other,
+                                                            page_num) != 0) {
+                                btrfs_dev_replace_stats_inc(
+                                        &sctx->dev_root->
+                                        fs_info->dev_replace.
+                                        num_write_errors);
+                                success = 0;
+                        }
+                } else if (sblock_other) {
+                        ret = scrub_repair_page_from_good_copy(sblock_bad,
+                                                               sblock_other,
+                                                               page_num, 0);
+                        if (0 == ret)
+                                page_bad->io_error = 0;
+                        else
+                                success = 0;
                 }
         }
 
-        if (success) {
+        if (success && !sctx->is_dev_replace) {
                 if (is_metadata || have_csum) {
                         /*
                          * need to verify the checksum now that all
@@ -1288,19 +1275,18 @@ out:
         return 0;
 }
 
-static inline int scrub_nr_raid_mirrors(struct btrfs_bio *bbio, u64 *raid_map)
+static inline int scrub_nr_raid_mirrors(struct btrfs_bio *bbio)
 {
-        if (raid_map) {
-                if (raid_map[bbio->num_stripes - 1] == RAID6_Q_STRIPE)
-                        return 3;
-                else
-                        return 2;
-        } else {
+        if (bbio->map_type & BTRFS_BLOCK_GROUP_RAID5)
+                return 2;
+        else if (bbio->map_type & BTRFS_BLOCK_GROUP_RAID6)
+                return 3;
+        else
                 return (int)bbio->num_stripes;
-        }
 }
 
-static inline void scrub_stripe_index_and_offset(u64 logical, u64 *raid_map,
+static inline void scrub_stripe_index_and_offset(u64 logical, u64 map_type,
+                                                 u64 *raid_map,
                                                  u64 mapped_length,
                                                  int nstripes, int mirror,
                                                  int *stripe_index,
@@ -1308,7 +1294,7 @@ static inline void scrub_stripe_index_and_offset(u64 logical, u64 *raid_map,
 {
         int i;
 
-        if (raid_map) {
+        if (map_type & BTRFS_BLOCK_GROUP_RAID56_MASK) {
                 /* RAID5/6 */
                 for (i = 0; i < nstripes; i++) {
                         if (raid_map[i] == RAID6_Q_STRIPE ||
@@ -1329,72 +1315,65 @@ static inline void scrub_stripe_index_and_offset(u64 logical, u64 *raid_map,
         }
 }
 
-static int scrub_setup_recheck_block(struct scrub_ctx *sctx,
-                                     struct btrfs_fs_info *fs_info,
-                                     struct scrub_block *original_sblock,
-                                     u64 length, u64 logical,
+static int scrub_setup_recheck_block(struct scrub_block *original_sblock,
                                      struct scrub_block *sblocks_for_recheck)
 {
+        struct scrub_ctx *sctx = original_sblock->sctx;
+        struct btrfs_fs_info *fs_info = sctx->dev_root->fs_info;
+        u64 length = original_sblock->page_count * PAGE_SIZE;
+        u64 logical = original_sblock->pagev[0]->logical;
         struct scrub_recover *recover;
         struct btrfs_bio *bbio;
-        u64 *raid_map;
         u64 sublen;
         u64 mapped_length;
         u64 stripe_offset;
         int stripe_index;
-        int page_index;
+        int page_index = 0;
         int mirror_index;
         int nmirrors;
         int ret;
 
         /*
-         * note: the two members ref_count and outstanding_pages
+         * note: the two members refs and outstanding_pages
          * are not used (and not set) in the blocks that are used for
          * the recheck procedure
          */
 
-        page_index = 0;
         while (length > 0) {
                 sublen = min_t(u64, length, PAGE_SIZE);
                 mapped_length = sublen;
                 bbio = NULL;
-                raid_map = NULL;
 
                 /*
                  * with a length of PAGE_SIZE, each returned stripe
                  * represents one mirror
                  */
                 ret = btrfs_map_sblock(fs_info, REQ_GET_READ_MIRRORS, logical,
-                                       &mapped_length, &bbio, 0, &raid_map);
+                                       &mapped_length, &bbio, 0, 1);
                 if (ret || !bbio || mapped_length < sublen) {
-                        kfree(bbio);
-                        kfree(raid_map);
+                        btrfs_put_bbio(bbio);
                         return -EIO;
                 }
 
                 recover = kzalloc(sizeof(struct scrub_recover), GFP_NOFS);
                 if (!recover) {
-                        kfree(bbio);
-                        kfree(raid_map);
+                        btrfs_put_bbio(bbio);
                         return -ENOMEM;
                 }
 
                 atomic_set(&recover->refs, 1);
                 recover->bbio = bbio;
-                recover->raid_map = raid_map;
                 recover->map_length = mapped_length;
 
                 BUG_ON(page_index >= SCRUB_PAGES_PER_RD_BIO);
 
-                nmirrors = scrub_nr_raid_mirrors(bbio, raid_map);
+                nmirrors = min(scrub_nr_raid_mirrors(bbio), BTRFS_MAX_MIRRORS);
+
                 for (mirror_index = 0; mirror_index < nmirrors;
                      mirror_index++) {
                         struct scrub_block *sblock;
                         struct scrub_page *page;
 
-                        if (mirror_index >= BTRFS_MAX_MIRRORS)
-                                continue;
-
                         sblock = sblocks_for_recheck + mirror_index;
                         sblock->sctx = sctx;
                         page = kzalloc(sizeof(*page), GFP_NOFS);
@@ -1410,9 +1389,12 @@ leave_nomem:
                         sblock->pagev[page_index] = page;
                         page->logical = logical;
 
-                        scrub_stripe_index_and_offset(logical, raid_map,
+                        scrub_stripe_index_and_offset(logical,
+                                                      bbio->map_type,
+                                                      bbio->raid_map,
                                                       mapped_length,
-                                                      bbio->num_stripes,
+                                                      bbio->num_stripes -
+                                                      bbio->num_tgtdevs,
                                                       mirror_index,
                                                       &stripe_index,
                                                       &stripe_offset);
@@ -1458,7 +1440,8 @@ static void scrub_bio_wait_endio(struct bio *bio, int error)
 
 static inline int scrub_is_page_on_raid56(struct scrub_page *page)
 {
-        return page->recover && page->recover->raid_map;
+        return page->recover &&
+               (page->recover->bbio->map_type & BTRFS_BLOCK_GROUP_RAID56_MASK);
 }
 
 static int scrub_submit_raid56_bio_wait(struct btrfs_fs_info *fs_info,
@@ -1475,7 +1458,6 @@ static int scrub_submit_raid56_bio_wait(struct btrfs_fs_info *fs_info,
         bio->bi_end_io = scrub_bio_wait_endio;
 
         ret = raid56_parity_recover(fs_info->fs_root, bio, page->recover->bbio,
-                                    page->recover->raid_map,
                                     page->recover->map_length,
                                     page->mirror_num, 0);
         if (ret)
@@ -1615,8 +1597,7 @@ static void scrub_recheck_block_checksum(struct btrfs_fs_info *fs_info,
 }
 
 static int scrub_repair_block_from_good_copy(struct scrub_block *sblock_bad,
-                                             struct scrub_block *sblock_good,
-                                             int force_write)
+                                             struct scrub_block *sblock_good)
 {
         int page_num;
         int ret = 0;
@@ -1626,8 +1607,7 @@ static int scrub_repair_block_from_good_copy(struct scrub_block *sblock_bad,
 
                 ret_sub = scrub_repair_page_from_good_copy(sblock_bad,
                                                            sblock_good,
-                                                           page_num,
-                                                           force_write);
+                                                           page_num, 1);
                 if (ret_sub)
                         ret = ret_sub;
         }
@@ -2067,12 +2047,12 @@ static int scrub_checksum_super(struct scrub_block *sblock)
 
 static void scrub_block_get(struct scrub_block *sblock)
 {
-        atomic_inc(&sblock->ref_count);
+        atomic_inc(&sblock->refs);
 }
 
 static void scrub_block_put(struct scrub_block *sblock)
 {
-        if (atomic_dec_and_test(&sblock->ref_count)) {
+        if (atomic_dec_and_test(&sblock->refs)) {
                 int i;
 
                 if (sblock->sparity)
@@ -2086,12 +2066,12 @@ static void scrub_block_put(struct scrub_block *sblock)
 
 static void scrub_page_get(struct scrub_page *spage)
 {
-        atomic_inc(&spage->ref_count);
+        atomic_inc(&spage->refs);
 }
 
 static void scrub_page_put(struct scrub_page *spage)
 {
-        if (atomic_dec_and_test(&spage->ref_count)) {
+        if (atomic_dec_and_test(&spage->refs)) {
                 if (spage->page)
                         __free_page(spage->page);
                 kfree(spage);
@@ -2217,7 +2197,7 @@ static int scrub_pages(struct scrub_ctx *sctx, u64 logical, u64 len,
 
         /* one ref inside this function, plus one for each page added to
          * a bio later on */
-        atomic_set(&sblock->ref_count, 1);
+        atomic_set(&sblock->refs, 1);
         sblock->sctx = sctx;
         sblock->no_io_error_seen = 1;
 
@@ -2510,7 +2490,7 @@ static int scrub_pages_for_parity(struct scrub_parity *sparity,
 
         /* one ref inside this function, plus one for each page added to
          * a bio later on */
-        atomic_set(&sblock->ref_count, 1);
+        atomic_set(&sblock->refs, 1);
         sblock->sctx = sctx;
         sblock->no_io_error_seen = 1;
         sblock->sparity = sparity;
@@ -2705,7 +2685,6 @@ static void scrub_parity_check_and_repair(struct scrub_parity *sparity)
         struct btrfs_raid_bio *rbio;
         struct scrub_page *spage;
         struct btrfs_bio *bbio = NULL;
-        u64 *raid_map = NULL;
         u64 length;
         int ret;
 
@@ -2716,8 +2695,8 @@ static void scrub_parity_check_and_repair(struct scrub_parity *sparity)
         length = sparity->logic_end - sparity->logic_start + 1;
         ret = btrfs_map_sblock(sctx->dev_root->fs_info, WRITE,
                                sparity->logic_start,
-                               &length, &bbio, 0, &raid_map);
-        if (ret || !bbio || !raid_map)
+                               &length, &bbio, 0, 1);
+        if (ret || !bbio || !bbio->raid_map)
                 goto bbio_out;
 
         bio = btrfs_io_bio_alloc(GFP_NOFS, 0);
@@ -2729,8 +2708,7 @@ static void scrub_parity_check_and_repair(struct scrub_parity *sparity)
         bio->bi_end_io = scrub_parity_bio_endio;
 
         rbio = raid56_parity_alloc_scrub_rbio(sctx->dev_root, bio, bbio,
-                                              raid_map, length,
-                                              sparity->scrub_dev,
+                                              length, sparity->scrub_dev,
                                               sparity->dbitmap,
                                               sparity->nsectors);
         if (!rbio)
@@ -2747,8 +2725,7 @@ static void scrub_parity_check_and_repair(struct scrub_parity *sparity)
 rbio_out:
         bio_put(bio);
 bbio_out:
-        kfree(bbio);
-        kfree(raid_map);
+        btrfs_put_bbio(bbio);
         bitmap_or(sparity->ebitmap, sparity->ebitmap, sparity->dbitmap,
                   sparity->nsectors);
         spin_lock(&sctx->stat_lock);
@@ -2765,12 +2742,12 @@ static inline int scrub_calc_parity_bitmap_len(int nsectors)
 
 static void scrub_parity_get(struct scrub_parity *sparity)
 {
-        atomic_inc(&sparity->ref_count);
+        atomic_inc(&sparity->refs);
 }
 
 static void scrub_parity_put(struct scrub_parity *sparity)
 {
-        if (!atomic_dec_and_test(&sparity->ref_count))
+        if (!atomic_dec_and_test(&sparity->refs))
                 return;
 
         scrub_parity_check_and_repair(sparity);
@@ -2820,7 +2797,7 @@ static noinline_for_stack int scrub_raid56_parity(struct scrub_ctx *sctx,
         sparity->scrub_dev = sdev;
         sparity->logic_start = logic_start;
         sparity->logic_end = logic_end;
-        atomic_set(&sparity->ref_count, 1);
+        atomic_set(&sparity->refs, 1);
         INIT_LIST_HEAD(&sparity->spages);
         sparity->dbitmap = sparity->bitmap;
         sparity->ebitmap = (void *)sparity->bitmap + bitmap_len;
@@ -3037,8 +3014,7 @@ static noinline_for_stack int scrub_stripe(struct scrub_ctx *sctx,
         } else if (map->type & BTRFS_BLOCK_GROUP_DUP) {
                 increment = map->stripe_len;
                 mirror_num = num % map->num_stripes + 1;
-        } else if (map->type & (BTRFS_BLOCK_GROUP_RAID5 |
-                                BTRFS_BLOCK_GROUP_RAID6)) {
+        } else if (map->type & BTRFS_BLOCK_GROUP_RAID56_MASK) {
                 get_raid56_logic_offset(physical, num, map, &offset, NULL);
                 increment = map->stripe_len * nr_data_stripes(map);
                 mirror_num = 1;
@@ -3074,8 +3050,7 @@ static noinline_for_stack int scrub_stripe(struct scrub_ctx *sctx,
          */
         logical = base + offset;
         physical_end = physical + nstripes * map->stripe_len;
-        if (map->type & (BTRFS_BLOCK_GROUP_RAID5 |
-                         BTRFS_BLOCK_GROUP_RAID6)) {
+        if (map->type & BTRFS_BLOCK_GROUP_RAID56_MASK) {
                 get_raid56_logic_offset(physical_end, num,
                                         map, &logic_end, NULL);
                 logic_end += base;
@@ -3121,8 +3096,7 @@ static noinline_for_stack int scrub_stripe(struct scrub_ctx *sctx,
         ret = 0;
         while (physical < physical_end) {
                 /* for raid56, we skip parity stripe */
-                if (map->type & (BTRFS_BLOCK_GROUP_RAID5 |
-                                BTRFS_BLOCK_GROUP_RAID6)) {
+                if (map->type & BTRFS_BLOCK_GROUP_RAID56_MASK) {
                         ret = get_raid56_logic_offset(physical, num,
                                         map, &logical, &stripe_logical);
                         logical += base;
@@ -3280,8 +3254,7 @@ again:
                         scrub_free_csums(sctx);
                         if (extent_logical + extent_len <
                             key.objectid + bytes) {
-                                if (map->type & (BTRFS_BLOCK_GROUP_RAID5 |
-                                        BTRFS_BLOCK_GROUP_RAID6)) {
+                                if (map->type & BTRFS_BLOCK_GROUP_RAID56_MASK) {
                                         /*
                                          * loop until we find next data stripe
                                          * or we have finished all stripes.
@@ -3775,7 +3748,7 @@ int btrfs_scrub_dev(struct btrfs_fs_info *fs_info, u64 devid, u64 start,
         scrub_workers_put(fs_info);
         mutex_unlock(&fs_info->scrub_lock);
 
-        scrub_free_ctx(sctx);
+        scrub_put_ctx(sctx);
 
         return ret;
 }
@@ -3881,14 +3854,14 @@ static void scrub_remap_extent(struct btrfs_fs_info *fs_info,
                               &mapped_length, &bbio, 0);
         if (ret || !bbio || mapped_length < extent_len ||
             !bbio->stripes[0].dev->bdev) {
-                kfree(bbio);
+                btrfs_put_bbio(bbio);
                 return;
         }
 
         *extent_physical = bbio->stripes[0].physical;
         *extent_mirror_num = bbio->mirror_num;
         *extent_dev = bbio->stripes[0].dev;
-        kfree(bbio);
+        btrfs_put_bbio(bbio);
 }
 
 static int scrub_setup_wr_ctx(struct scrub_ctx *sctx,