2 files changed, 5 insertions, 3 deletions

diff --git a/block/bio.c b/block/bio.c
index b12966e415d3..8c680a776171 100644
--- a/block/bio.c
+++ b/block/bio.c
@@ -2015,7 +2015,8 @@ int bio_associate_blkg(struct bio *bio, struct blkcg_gq *blkg)
 {
         if (unlikely(bio->bi_blkg))
                 return -EBUSY;
-        blkg_get(blkg);
+        if (!blkg_try_get(blkg))
+                return -ENODEV;
         bio->bi_blkg = blkg;
         return 0;
 }
diff --git a/block/blk-throttle.c b/block/blk-throttle.c
index a3eede00d302..01d0620a4e4a 100644
--- a/block/blk-throttle.c
+++ b/block/blk-throttle.c
@@ -2129,8 +2129,9 @@ static inline void throtl_update_latency_buckets(struct throtl_data *td)
 static void blk_throtl_assoc_bio(struct throtl_grp *tg, struct bio *bio)
 {
 #ifdef CONFIG_BLK_DEV_THROTTLING_LOW
-        if (bio->bi_css)
-                bio_associate_blkg(bio, tg_to_blkg(tg));
+        /* fallback to root_blkg if we fail to get a blkg ref */
+        if (bio->bi_css && (bio_associate_blkg(bio, tg_to_blkg(tg)) == -ENODEV))
+                bio_associate_blkg(bio, bio->bi_disk->queue->root_blkg);
         bio_issue_init(&bio->bi_issue, bio_sectors(bio));
 #endif
 }