5 files changed, 28 insertions, 4 deletions

diff --git a/fs/btrfs/disk-io.c b/fs/btrfs/disk-io.c
index f7536bcf7cee..230546b45474 100644
--- a/fs/btrfs/disk-io.c
+++ b/fs/btrfs/disk-io.c
@@ -1724,6 +1724,7 @@ static int setup_bdi(struct btrfs_fs_info *info, struct backing_dev_info *bdi)
         bdi->ra_pages = VM_MAX_READAHEAD * 1024 / PAGE_CACHE_SIZE;
         bdi->congested_fn       = btrfs_congested_fn;
         bdi->congested_data     = info;
+        bdi->capabilities |= BDI_CAP_CGROUP_WRITEBACK;
         return 0;
 }
 
diff --git a/fs/btrfs/extent_io.c b/fs/btrfs/extent_io.c
index 02d05817cbdf..b9755ce98218 100644
--- a/fs/btrfs/extent_io.c
+++ b/fs/btrfs/extent_io.c
@@ -2730,6 +2730,9 @@ struct bio *btrfs_bio_clone(struct bio *bio, gfp_t gfp_mask)
                 btrfs_bio->csum = NULL;
                 btrfs_bio->csum_allocated = NULL;
                 btrfs_bio->end_io = NULL;
+                /* FIXME, put this into bio_clone_bioset */
+                if (bio->bi_css)
+                        bio_associate_blkcg(new, bio->bi_css);
         }
         return new;
 }
@@ -2790,6 +2793,7 @@ static int merge_bio(int rw, struct extent_io_tree *tree, struct page *page,
 }
 
 static int submit_extent_page(int rw, struct extent_io_tree *tree,
+                              struct writeback_control *wbc,
                               struct page *page, sector_t sector,
                               size_t size, unsigned long offset,
                               struct block_device *bdev,
@@ -2826,6 +2830,8 @@ static int submit_extent_page(int rw, struct extent_io_tree *tree,
                         }
                         bio = NULL;
                 } else {
+                        if (wbc)
+                                wbc_account_io(wbc, page, page_size);
                         return 0;
                 }
         }
@@ -2841,6 +2847,10 @@ static int submit_extent_page(int rw, struct extent_io_tree *tree,
         bio_add_page(bio, page, page_size, offset);
         bio->bi_end_io = end_io_func;
         bio->bi_private = tree;
+        if (wbc) {
+                wbc_init_bio(wbc, bio);
+                wbc_account_io(wbc, page, page_size);
+        }
 
         if (bio_ret)
                 *bio_ret = bio;
@@ -3051,7 +3061,7 @@ static int __do_readpage(struct extent_io_tree *tree,
                 }
 
                 pnr -= page->index;
-                ret = submit_extent_page(rw, tree, page,
+                ret = submit_extent_page(rw, tree, NULL, page,
                                          sector, disk_io_size, pg_offset,
                                          bdev, bio, pnr,
                                          end_bio_extent_readpage, mirror_num,
@@ -3446,7 +3456,7 @@ static noinline_for_stack int __extent_writepage_io(struct inode *inode,
                                        page->index, cur, end);
                         }
 
-                        ret = submit_extent_page(write_flags, tree, page,
+                        ret = submit_extent_page(write_flags, tree, wbc, page,
                                                  sector, iosize, pg_offset,
                                                  bdev, &epd->bio, max_nr,
                                                  end_bio_extent_writepage,
@@ -3749,7 +3759,7 @@ static noinline_for_stack int write_one_eb(struct extent_buffer *eb,
 
                 clear_page_dirty_for_io(p);
                 set_page_writeback(p);
-                ret = submit_extent_page(rw, tree, p, offset >> 9,
+                ret = submit_extent_page(rw, tree, wbc, p, offset >> 9,
                                          PAGE_CACHE_SIZE, 0, bdev, &epd->bio,
                                          -1, end_bio_extent_buffer_writepage,
                                          0, epd->bio_flags, bio_flags);
diff --git a/fs/btrfs/inode.c b/fs/btrfs/inode.c
index 79a73645346e..bda3c41dc9d5 100644
--- a/fs/btrfs/inode.c
+++ b/fs/btrfs/inode.c
@@ -7987,7 +7987,11 @@ static struct bio *btrfs_dio_bio_alloc(struct block_device *bdev,
                                        u64 first_sector, gfp_t gfp_flags)
 {
         int nr_vecs = bio_get_nr_vecs(bdev);
-        return btrfs_bio_alloc(bdev, first_sector, nr_vecs, gfp_flags);
+        struct bio *bio;
+        bio = btrfs_bio_alloc(bdev, first_sector, nr_vecs, gfp_flags);
+        if (bio)
+                bio_associate_current(bio);
+        return bio;
 }
 
 static inline int btrfs_lookup_and_bind_dio_csum(struct btrfs_root *root,
diff --git a/fs/btrfs/super.c b/fs/btrfs/super.c
index cd7ef34d2dce..d366dd4664d0 100644
--- a/fs/btrfs/super.c
+++ b/fs/btrfs/super.c
@@ -1033,6 +1033,7 @@ static int btrfs_fill_super(struct super_block *sb,
         sb->s_flags |= MS_POSIXACL;
 #endif
         sb->s_flags |= MS_I_VERSION;
+        sb->s_iflags |= SB_I_CGROUPWB;
         err = open_ctree(sb, fs_devices, (char *)data);
         if (err) {
                 printk(KERN_ERR "BTRFS: open_ctree failed\n");
diff --git a/fs/btrfs/volumes.c b/fs/btrfs/volumes.c
index fb9abf1678d0..88e2fe931bde 100644
--- a/fs/btrfs/volumes.c
+++ b/fs/btrfs/volumes.c
@@ -5942,6 +5942,14 @@ again:
         if (!bio)
                 return -ENOMEM;
 
+        if (first_bio->bi_ioc) {
+                get_io_context_active(first_bio->bi_ioc);
+                bio->bi_ioc = first_bio->bi_ioc;
+        }
+        if (first_bio->bi_css) {
+                css_get(first_bio->bi_css);
+                bio->bi_css = first_bio->bi_css;
+        }
         while (bvec <= (first_bio->bi_io_vec + first_bio->bi_vcnt - 1)) {
                 if (bio_add_page(bio, bvec->bv_page, bvec->bv_len,
                                  bvec->bv_offset) < bvec->bv_len) {