return cache;
 }
 
-bool btrfs_inc_nocow_writers(struct btrfs_fs_info *fs_info, u64 bytenr)
+/**
+ * Check if we can do a NOCOW write for a given extent.
+ *
+ * @fs_info:       The filesystem information object.
+ * @bytenr:        Logical start address of the extent.
+ *
+ * Check if we can do a NOCOW write for the given extent, and increments the
+ * number of NOCOW writers in the block group that contains the extent, as long
+ * as the block group exists and it's currently not in read-only mode.
+ *
+ * Returns: A non-NULL block group pointer if we can do a NOCOW write, the caller
+ *          is responsible for calling btrfs_dec_nocow_writers() later.
+ *
+ *          Or NULL if we can not do a NOCOW write
+ */
+struct btrfs_block_group *btrfs_inc_nocow_writers(struct btrfs_fs_info *fs_info,
+                                                 u64 bytenr)
 {
        struct btrfs_block_group *bg;
-       bool ret = true;
+       bool can_nocow = true;
 
        bg = btrfs_lookup_block_group(fs_info, bytenr);
        if (!bg)
-               return false;
+               return NULL;
 
        spin_lock(&bg->lock);
        if (bg->ro)
-               ret = false;
+               can_nocow = false;
        else
                atomic_inc(&bg->nocow_writers);
        spin_unlock(&bg->lock);
 
-       /* No put on block group, done by btrfs_dec_nocow_writers */
-       if (!ret)
+       if (!can_nocow) {
                btrfs_put_block_group(bg);
+               return NULL;
+       }
 
-       return ret;
+       /* No put on block group, done by btrfs_dec_nocow_writers(). */
+       return bg;
 }
 
-void btrfs_dec_nocow_writers(struct btrfs_fs_info *fs_info, u64 bytenr)
+/**
+ * Decrement the number of NOCOW writers in a block group.
+ *
+ * @bg:       The block group.
+ *
+ * This is meant to be called after a previous call to btrfs_inc_nocow_writers(),
+ * and on the block group returned by that call. Typically this is called after
+ * creating an ordered extent for a NOCOW write, to prevent races with scrub and
+ * relocation.
+ *
+ * After this call, the caller should not use the block group anymore. It it wants
+ * to use it, then it should get a reference on it before calling this function.
+ */
+void btrfs_dec_nocow_writers(struct btrfs_block_group *bg)
 {
-       struct btrfs_block_group *bg;
-
-       bg = btrfs_lookup_block_group(fs_info, bytenr);
-       ASSERT(bg);
        if (atomic_dec_and_test(&bg->nocow_writers))
                wake_up_var(&bg->nocow_writers);
-       /*
-        * Once for our lookup and once for the lookup done by a previous call
-        * to btrfs_inc_nocow_writers()
-        */
-       btrfs_put_block_group(bg);
+
+       /* For the lookup done by a previous call to btrfs_inc_nocow_writers(). */
        btrfs_put_block_group(bg);
 }
 
 
 void btrfs_dec_block_group_reservations(struct btrfs_fs_info *fs_info,
                                        const u64 start);
 void btrfs_wait_block_group_reservations(struct btrfs_block_group *bg);
-bool btrfs_inc_nocow_writers(struct btrfs_fs_info *fs_info, u64 bytenr);
-void btrfs_dec_nocow_writers(struct btrfs_fs_info *fs_info, u64 bytenr);
+struct btrfs_block_group *btrfs_inc_nocow_writers(struct btrfs_fs_info *fs_info,
+                                                 u64 bytenr);
+void btrfs_dec_nocow_writers(struct btrfs_block_group *bg);
 void btrfs_wait_nocow_writers(struct btrfs_block_group *bg);
 void btrfs_wait_block_group_cache_progress(struct btrfs_block_group *cache,
                                           u64 num_bytes);
 
        int ret;
        bool check_prev = true;
        u64 ino = btrfs_ino(inode);
+       struct btrfs_block_group *bg;
        bool nocow = false;
        struct can_nocow_file_extent_args nocow_args = { 0 };
 
                }
 
                ret = 0;
-               if (btrfs_inc_nocow_writers(fs_info, nocow_args.disk_bytenr))
+               bg = btrfs_inc_nocow_writers(fs_info, nocow_args.disk_bytenr);
+               if (bg)
                        nocow = true;
 out_check:
                /*
                                goto error;
                }
 
-               if (nocow)
-                       btrfs_dec_nocow_writers(fs_info, nocow_args.disk_bytenr);
-               nocow = false;
+               if (nocow) {
+                       btrfs_dec_nocow_writers(bg);
+                       nocow = false;
+               }
 
                if (btrfs_is_data_reloc_root(root))
                        /*
 
 error:
        if (nocow)
-               btrfs_dec_nocow_writers(fs_info, nocow_args.disk_bytenr);
+               btrfs_dec_nocow_writers(bg);
 
        if (ret && cur_offset < end)
                extent_clear_unlock_delalloc(inode, cur_offset, end,
        struct extent_map *em = *map;
        int type;
        u64 block_start, orig_start, orig_block_len, ram_bytes;
+       struct btrfs_block_group *bg;
        bool can_nocow = false;
        bool space_reserved = false;
        u64 prev_len;
                block_start = em->block_start + (start - em->start);
 
                if (can_nocow_extent(inode, start, &len, &orig_start,
-                                    &orig_block_len, &ram_bytes, false) == 1 &&
-                   btrfs_inc_nocow_writers(fs_info, block_start))
-                       can_nocow = true;
+                                    &orig_block_len, &ram_bytes, false) == 1) {
+                       bg = btrfs_inc_nocow_writers(fs_info, block_start);
+                       if (bg)
+                               can_nocow = true;
+               }
        }
 
        prev_len = len;
                        /* Our caller expects us to free the input extent map. */
                        free_extent_map(em);
                        *map = NULL;
-                       btrfs_dec_nocow_writers(fs_info, block_start);
+                       btrfs_dec_nocow_writers(bg);
                        if (nowait && (ret == -ENOSPC || ret == -EDQUOT))
                                ret = -EAGAIN;
                        goto out;
                                              orig_start, block_start,
                                              len, orig_block_len,
                                              ram_bytes, type);
-               btrfs_dec_nocow_writers(fs_info, block_start);
+               btrfs_dec_nocow_writers(bg);
                if (type == BTRFS_ORDERED_PREALLOC) {
                        free_extent_map(em);
                        *map = em = em2;