[PATCH v2 0/3] blk-cgroup: store blkcg in bio before blkcg_mutex conversion

@@ -64,10 +64,17 @@ bool blkcg_debug_stats = false;
 
 static DEFINE_RAW_SPINLOCK(blkg_stat_lock);
 
 #define BLKG_DESTROY_BATCH_SIZE  64
 
+const struct rhashtable_params blkg_hash_params = {
+	.key_len		= sizeof_field(struct blkcg_gq, blkcg_id),
+	.key_offset		= offsetof(struct blkcg_gq, blkcg_id),
+	.head_offset		= offsetof(struct blkcg_gq, q_hash_node),
+	.automatic_shrinking	= true,
+};
+
 /*
  * Lockless lists for tracking IO stats update
  *
  * New IO stats are stored in the percpu iostat_cpu within blkcg_gq (blkg).
  * There are multiple blkg's (one for each block device) attached to each

@@ -194,10 +201,20 @@ static void blkg_release(struct percpu_ref *ref)
 {
 	struct blkcg_gq *blkg = container_of(ref, struct blkcg_gq, refcnt);
 	struct blkcg *blkcg = blkg->blkcg;
 	int cpu;
 
+	/*
+	 * A blkg that was never inserted into q->blkg_list has no hash
+	 * entry.  This happens when allocation or creation fails before
+	 * rhashtable_insert_fast() succeeds.
+	 */
+	if (!list_empty(&blkg->q_node))
+		WARN_ON_ONCE(rhashtable_remove_fast(&blkg->q->blkg_hash,
+						    &blkg->q_hash_node,
+						    blkg_hash_params));
+
 	/*
 	 * Flush all the non-empty percpu lockless lists before releasing
 	 * us, given these stat belongs to us.
 	 *
 	 * blkg_stat_lock is for serializing blkg stat update

@@ -327,10 +344,11 @@ static struct blkcg_gq *blkg_alloc(struct blkcg *blkcg, struct gendisk *disk,
 		goto out_put_queue;
 
 	blkg->q = disk->queue;
 	INIT_LIST_HEAD(&blkg->q_node);
 	blkg->blkcg = blkcg;
+	blkg->blkcg_id = blkcg->css.id;
 	blkg->iostat.blkg = blkg;
 #ifdef CONFIG_BLK_CGROUP_PUNT_BIO
 	spin_lock_init(&blkg->async_bio_lock);
 	bio_list_init(&blkg->async_bios);
 	INIT_WORK(&blkg->async_bio_work, blkg_async_bio_workfn);

@@ -423,11 +441,12 @@ static struct blkcg_gq *blkg_create(struct blkcg *blkcg, struct gendisk *disk,
 			pol->pd_init_fn(blkg->pd[i]);
 	}
 
 	/* insert */
 	spin_lock(&blkcg->lock);
-	ret = radix_tree_insert(&blkcg->blkg_tree, disk->queue->id, blkg);
+	ret = rhashtable_insert_fast(&disk->queue->blkg_hash,
+				     &blkg->q_hash_node, blkg_hash_params);
 	if (likely(!ret)) {
 		hlist_add_head_rcu(&blkg->blkcg_node, &blkcg->blkg_list);
 		list_add(&blkg->q_node, &disk->queue->blkg_list);
 
 		for (i = 0; i < BLKCG_MAX_POLS; i++) {

@@ -476,13 +495,10 @@ static struct blkcg_gq *blkg_lookup_create(struct blkcg *blkcg,
 	struct blkcg_gq *blkg;
 
 	rcu_read_lock();
 	blkg = blkg_lookup(blkcg, q);
 	if (blkg) {
-		if (blkcg != &blkcg_root &&
-		    blkg != rcu_dereference(blkcg->blkg_hint))
-			rcu_assign_pointer(blkcg->blkg_hint, blkg);
 		rcu_read_unlock();
 		return blkg;
 	}
 	rcu_read_unlock();
 

@@ -546,21 +562,12 @@ static void blkg_destroy(struct blkcg_gq *blkg)
 		}
 	}
 
 	blkg->online = false;
 
-	radix_tree_delete(&blkcg->blkg_tree, blkg->q->id);
 	hlist_del_init_rcu(&blkg->blkcg_node);
 
-	/*
-	 * Both setting lookup hint to and clearing it from @blkg are done
-	 * under queue_lock.  If it's not pointing to @blkg now, it never
-	 * will.  Hint assignment itself can race safely.
-	 */
-	if (rcu_access_pointer(blkcg->blkg_hint) == blkg)
-		rcu_assign_pointer(blkcg->blkg_hint, NULL);
-
 	/*
 	 * Put the reference taken at the time of creation so that when all
 	 * queues are gone, group can be destroyed.
 	 */
 	percpu_ref_kill(&blkg->refcnt);

@@ -880,47 +887,37 @@ int blkg_conf_prep(struct blkcg *blkcg, const struct blkcg_policy *pol,
 		if (unlikely(!new_blkg)) {
 			ret = -ENOMEM;
 			goto fail_exit;
 		}
 
-		if (radix_tree_preload(GFP_KERNEL)) {
-			blkg_free(new_blkg);
-			ret = -ENOMEM;
-			goto fail_exit;
-		}
-
 		spin_lock_irq(&q->queue_lock);
 
 		if (!blkcg_policy_enabled(q, pol)) {
 			blkg_free(new_blkg);
 			ret = -EOPNOTSUPP;
-			goto fail_preloaded;
+			goto fail_unlock;
 		}
 
 		blkg = blkg_lookup(pos, q);
 		if (blkg) {
 			blkg_free(new_blkg);
 		} else {
 			blkg = blkg_create(pos, disk, new_blkg);
 			if (IS_ERR(blkg)) {
 				ret = PTR_ERR(blkg);
-				goto fail_preloaded;
+				goto fail_unlock;
 			}
 		}
 
-		radix_tree_preload_end();
-
 		if (pos == blkcg)
 			goto success;
 	}
 success:
 	mutex_unlock(&q->blkcg_mutex);
 	ctx->blkg = blkg;
 	return 0;
 
-fail_preloaded:
-	radix_tree_preload_end();
 fail_unlock:
 	spin_unlock_irq(&q->queue_lock);
 fail_exit:
 	mutex_unlock(&q->blkcg_mutex);
 	/*

@@ -1418,11 +1415,10 @@ blkcg_css_alloc(struct cgroup_subsys_state *parent_css)
 		cpd->plid = i;
 	}
 
 	spin_lock_init(&blkcg->lock);
 	refcount_set(&blkcg->online_pin, 1);
-	INIT_RADIX_TREE(&blkcg->blkg_tree, GFP_NOWAIT);
 	INIT_HLIST_HEAD(&blkcg->blkg_list);
 #ifdef CONFIG_CGROUP_WRITEBACK
 	INIT_LIST_HEAD(&blkcg->cgwb_list);
 #endif
 	list_add_tail(&blkcg->all_blkcgs_node, &all_blkcgs);

@@ -1455,21 +1451,26 @@ static int blkcg_css_online(struct cgroup_subsys_state *css)
 	if (parent)
 		blkcg_pin_online(&parent->css);
 	return 0;
 }
 
-void blkg_init_queue(struct request_queue *q)
+int blkg_init_queue(struct request_queue *q)
 {
 	INIT_LIST_HEAD(&q->blkg_list);
 	mutex_init(&q->blkcg_mutex);
+	return rhashtable_init(&q->blkg_hash, &blkg_hash_params);
+}
+
+void blkg_exit_queue(struct request_queue *q)
+{
+	rhashtable_destroy(&q->blkg_hash);
 }
 
 int blkcg_init_disk(struct gendisk *disk)
 {
 	struct request_queue *q = disk->queue;
 	struct blkcg_gq *new_blkg, *blkg;
-	bool preloaded;
 
 	/*
 	 * If the queue is shared across disk rebind (e.g., SCSI), the
 	 * previous disk's blkcg state is cleaned up asynchronously via
 	 * disk_release() -> blkcg_exit_disk(). Wait for all old blkgs to be

@@ -1479,30 +1480,23 @@ int blkcg_init_disk(struct gendisk *disk)
 
 	new_blkg = blkg_alloc(&blkcg_root, disk, GFP_KERNEL);
 	if (!new_blkg)
 		return -ENOMEM;
 
-	preloaded = !radix_tree_preload(GFP_KERNEL);
-
 	/* Make sure the root blkg exists. */
 	/* spin_lock_irq can serve as RCU read-side critical section. */
 	spin_lock_irq(&q->queue_lock);
 	blkg = blkg_create(&blkcg_root, disk, new_blkg);
 	if (IS_ERR(blkg))
 		goto err_unlock;
 	q->root_blkg = blkg;
 	spin_unlock_irq(&q->queue_lock);
 
-	if (preloaded)
-		radix_tree_preload_end();
-
 	return 0;
 
 err_unlock:
 	spin_unlock_irq(&q->queue_lock);
-	if (preloaded)
-		radix_tree_preload_end();
 	return PTR_ERR(blkg);
 }
 
 void blkcg_exit_disk(struct gendisk *disk)
 {

@@ -17,10 +17,12 @@
 #include <linux/blk-cgroup.h>
 #include <linux/cgroup.h>
 #include <linux/kthread.h>
 #include <linux/blk-mq.h>
 #include <linux/llist.h>
+#include <linux/rhashtable.h>
+#include <linux/rcupdate.h>
 #include "blk.h"
 
 struct blkcg_gq;
 struct blkg_policy_data;
 

@@ -54,13 +56,15 @@ struct blkg_iostat_set {
 
 /* association between a blk cgroup and a request queue */
 struct blkcg_gq {
 	/* Pointer to the associated request_queue */
 	struct request_queue		*q;
+	struct rhash_head		q_hash_node;
 	struct list_head		q_node;
 	struct hlist_node		blkcg_node;
 	struct blkcg			*blkcg;
+	int				blkcg_id;
 
 	/* all non-root blkcg_gq's are guaranteed to have access to parent */
 	struct blkcg_gq			*parent;
 
 	/* reference count */

@@ -96,12 +100,10 @@ struct blkcg {
 	spinlock_t			lock;
 	refcount_t			online_pin;
 	/* If there is block congestion on this cgroup. */
 	atomic_t			congestion_count;
 
-	struct radix_tree_root		blkg_tree;
-	struct blkcg_gq	__rcu		*blkg_hint;
 	struct hlist_head		blkg_list;
 
 	struct blkcg_policy_data	*cpd[BLKCG_MAX_POLS];
 
 	struct list_head		all_blkcgs_node;

@@ -190,12 +192,14 @@ struct blkcg_policy {
 	blkcg_pol_stat_pd_fn		*pd_stat_fn;
 };
 
 extern struct blkcg blkcg_root;
 extern bool blkcg_debug_stats;
+extern const struct rhashtable_params blkg_hash_params;
 
-void blkg_init_queue(struct request_queue *q);
+int blkg_init_queue(struct request_queue *q);
+void blkg_exit_queue(struct request_queue *q);
 int blkcg_init_disk(struct gendisk *disk);
 void blkcg_exit_disk(struct gendisk *disk);
 
 /* Blkio controller policy registration */
 int blkcg_policy_register(struct blkcg_policy *pol);

@@ -247,15 +251,37 @@ static inline bool bio_issue_as_root_blkg(struct bio *bio)
 {
 	return (bio->bi_opf & (REQ_META | REQ_SWAP)) != 0;
 }
 
 /**
- * blkg_lookup - lookup blkg for the specified blkcg - q pair
+ * blkg_lookup_any - lookup any blkg for the specified blkcg - q pair
  * @blkcg: blkcg of interest
  * @q: request_queue of interest
  *
- * Lookup blkg for the @blkcg - @q pair.
+ * Lookup a blkg for the @blkcg - @q pair, whether it is online or dying.
+ *
+ * Must be called in a RCU critical section.
+ *
+ * This does not acquire a reference.  The caller must already hold one, or
+ * have the blkg pinned by I/O.
+ */
+static inline struct blkcg_gq *blkg_lookup_any(struct blkcg *blkcg,
+					       struct request_queue *q)
+{
+	RCU_LOCKDEP_WARN(!rcu_read_lock_held(),
+			 "blkg_lookup_any() requires an RCU read lock");
+
+	return rhashtable_lookup(&q->blkg_hash, &blkcg->css.id,
+				 blkg_hash_params);
+}
+
+/**
+ * blkg_lookup - lookup an online blkg for the specified blkcg - q pair
+ * @blkcg: blkcg of interest
+ * @q: request_queue of interest
+ *
+ * Lookup an online blkg for the @blkcg - @q pair.
  *
  * Must be called in a RCU critical section.
  */
 static inline struct blkcg_gq *blkg_lookup(struct blkcg *blkcg,
 					   struct request_queue *q)

@@ -263,17 +289,12 @@ static inline struct blkcg_gq *blkg_lookup(struct blkcg *blkcg,
 	struct blkcg_gq *blkg;
 
 	if (blkcg == &blkcg_root)
 		return q->root_blkg;
 
-	blkg = rcu_dereference_check(blkcg->blkg_hint,
-			lockdep_is_held(&q->queue_lock));
-	if (blkg && blkg->q == q)
-		return blkg;
-
-	blkg = radix_tree_lookup(&blkcg->blkg_tree, q->id);
-	if (blkg && blkg->q != q)
+	blkg = blkg_lookup_any(blkcg, q);
+	if (blkg && !READ_ONCE(blkg->online))
 		blkg = NULL;
 	return blkg;
 }
 
 /**

@@ -496,12 +517,14 @@ struct blkcg_policy {
 };
 
 struct blkcg {
 };
 
+static inline struct blkcg_gq *blkg_lookup_any(struct blkcg *blkcg, void *key) { return NULL; }
 static inline struct blkcg_gq *blkg_lookup(struct blkcg *blkcg, void *key) { return NULL; }
-static inline void blkg_init_queue(struct request_queue *q) { }
+static inline int blkg_init_queue(struct request_queue *q) { return 0; }
+static inline void blkg_exit_queue(struct request_queue *q) { }
 static inline int blkcg_init_disk(struct gendisk *disk) { return 0; }
 static inline void blkcg_exit_disk(struct gendisk *disk) { }
 static inline int blkcg_policy_register(struct blkcg_policy *pol) { return 0; }
 static inline void blkcg_policy_unregister(struct blkcg_policy *pol) { }
 static inline int blkcg_activate_policy(struct gendisk *disk,

@@ -301,10 +301,11 @@ static void blk_free_queue(struct request_queue *q)
 {
 	blk_free_queue_stats(q->stats);
 	if (queue_is_mq(q))
 		blk_mq_release(q);
 
+	blkg_exit_queue(q);
 	ida_free(&blk_queue_ida, q->id);
 	lockdep_unregister_key(&q->io_lock_cls_key);
 	lockdep_unregister_key(&q->q_lock_cls_key);
 	call_rcu(&q->rcu_head, blk_free_queue_rcu);
 }

@@ -479,21 +480,23 @@ struct request_queue *blk_alloc_queue(struct queue_limits *lim, int node_id)
 	spin_lock_init(&q->queue_lock);
 
 	init_waitqueue_head(&q->mq_freeze_wq);
 	mutex_init(&q->mq_freeze_lock);
 
-	blkg_init_queue(q);
+	error = blkg_init_queue(q);
+	if (error)
+		goto fail_stats;
 
 	/*
 	 * Init percpu_ref in atomic mode so that it's faster to shutdown.
 	 * See blk_register_queue() for details.
 	 */
 	error = percpu_ref_init(&q->q_usage_counter,
 				blk_queue_usage_counter_release,
 				PERCPU_REF_INIT_ATOMIC, GFP_KERNEL);
 	if (error)
-		goto fail_stats;
+		goto fail_blkg;
 	lockdep_register_key(&q->io_lock_cls_key);
 	lockdep_register_key(&q->q_lock_cls_key);
 	lockdep_init_map(&q->io_lockdep_map, "&q->q_usage_counter(io)",
 			 &q->io_lock_cls_key, 0);
 	lockdep_init_map(&q->q_lockdep_map, "&q->q_usage_counter(queue)",

@@ -508,10 +511,12 @@ struct request_queue *blk_alloc_queue(struct queue_limits *lim, int node_id)
 	q->nr_requests = BLKDEV_DEFAULT_RQ;
 	q->async_depth = BLKDEV_DEFAULT_RQ;
 
 	return q;
 
+fail_blkg:
+	blkg_exit_queue(q);
 fail_stats:
 	blk_free_queue_stats(q->stats);
 fail_id:
 	ida_free(&blk_queue_ida, q->id);
 fail_q:

@@ -25,10 +25,11 @@
 #include <linux/sbitmap.h>
 #include <linux/uuid.h>
 #include <linux/xarray.h>
 #include <linux/file.h>
 #include <linux/lockdep.h>
+#include <linux/rhashtable-types.h>
 
 struct module;
 struct request_queue;
 struct elevator_queue;
 struct blk_trace;

@@ -577,10 +578,11 @@ struct request_queue {
 
 	struct list_head	icq_list;
 #ifdef CONFIG_BLK_CGROUP
 	DECLARE_BITMAP		(blkcg_pols, BLKCG_MAX_POLS);
 	struct blkcg_gq		*root_blkg;
+	struct rhashtable	blkg_hash;
 	struct list_head	blkg_list;
 	struct mutex		blkcg_mutex;
 #endif
 
 	int			node;

@@ -3306,11 +3306,11 @@ incompatible.
 wbc_init_bio() binds the specified bio to its cgroup.  Depending on
 the configuration, the bio may be executed at a lower priority and if
 the writeback session is holding shared resources, e.g. a journal
 entry, may lead to priority inversion.  There is no one easy solution
 for the problem.  Filesystems can try to work around specific problem
-cases by skipping wbc_init_bio() and using bio_associate_blkg()
+cases by skipping wbc_init_bio() and using bio_associate_blkcg()
 directly.
 
 
 Deprecated v1 Core Features
 ===========================

@@ -361,15 +361,17 @@ void bfqg_and_blkg_put(struct bfq_group *bfqg)
 	bfqg_put(bfqg);
 }
 
 void bfqg_stats_update_legacy_io(struct request_queue *q, struct request *rq)
 {
-	struct bfq_group *bfqg = blkg_to_bfqg(rq->bio->bi_blkg);
+	struct blkcg_gq *blkg = bio_blkg_lookup(rq->bio);
+	struct bfq_group *bfqg;
 
-	if (!bfqg)
+	if (!blkg)
 		return;
 
+	bfqg = blkg_to_bfqg(blkg);
 	blkg_rwstat_add(&bfqg->stats.bytes, rq->cmd_flags, blk_rq_bytes(rq));
 	blkg_rwstat_add(&bfqg->stats.ios, rq->cmd_flags, 1);
 }
 
 /* @stats = 0 */

@@ -604,27 +606,29 @@ static void bfq_link_bfqg(struct bfq_data *bfqd, struct bfq_group *bfqg)
 	}
 }
 
 struct bfq_group *bfq_bio_bfqg(struct bfq_data *bfqd, struct bio *bio)
 {
-	struct blkcg_gq *blkg = bio->bi_blkg;
+	struct blkcg_gq *blkg = bio_blkg_lookup(bio);
 	struct bfq_group *bfqg;
 
 	while (blkg) {
 		if (!data_race(blkg->online)) {
 			blkg = blkg->parent;
 			continue;
 		}
+
 		bfqg = blkg_to_bfqg(blkg);
 		if (bfqg->pd.online) {
-			bio_associate_blkg_from_css(bio, &blkg->blkcg->css);
+			bio_associate_blkcg_from_css(bio, &blkg->blkcg->css);
 			return bfqg;
 		}
 		blkg = blkg->parent;
 	}
-	bio_associate_blkg_from_css(bio,
-				&bfqg_to_blkg(bfqd->root_group)->blkcg->css);
+
+	blkg = bfqg_to_blkg(bfqd->root_group);
+	bio_associate_blkcg_from_css(bio, &blkg->blkcg->css);
 	return bfqd->root_group;
 }
 
 /**
  * bfq_bfqq_move - migrate @bfqq to @bfqg.

@@ -1130,10 +1134,11 @@ static int bfqg_print_rwstat(struct seq_file *sf, void *v)
 	return 0;
 }
 
 static u64 bfqg_prfill_rwstat_recursive(struct seq_file *sf,
 					struct blkg_policy_data *pd, int off)
+	__must_hold_shared(RCU)
 {
 	struct blkg_rwstat_sample sum;
 
 	blkg_rwstat_recursive_sum(pd_to_blkg(pd), &blkcg_policy_bfq, off, &sum);
 	return __blkg_prfill_rwstat(sf, pd, &sum);

@@ -1206,10 +1211,11 @@ static int bfqg_print_stat_sectors(struct seq_file *sf, void *v)
 	return 0;
 }
 
 static u64 bfqg_prfill_sectors_recursive(struct seq_file *sf,
 					 struct blkg_policy_data *pd, int off)
+	__must_hold_shared(RCU)
 {
 	struct blkg_rwstat_sample tmp;
 
 	blkg_rwstat_recursive_sum(pd->blkg, &blkcg_policy_bfq,
 			offsetof(struct bfq_group, stats.bytes), &tmp);

@@ -126,10 +126,11 @@
 
 #include <trace/events/block.h>
 
 #include "elevator.h"
 #include "blk.h"
+#include "blk-cgroup.h"
 #include "blk-mq.h"
 #include "blk-mq-sched.h"
 #include "bfq-iosched.h"
 #include "blk-wbt.h"
 

@@ -2450,19 +2451,20 @@ static bool bfq_bio_merge(struct request_queue *q, struct bio *bio,
 	struct bfq_data *bfqd = q->elevator->elevator_data;
 	struct bfq_io_cq *bic = bfq_bic_lookup(q);
 	struct request *free = NULL;
 	bool ret;
 
+#ifdef CONFIG_BFQ_GROUP_IOSCHED
+	/* blkg creation takes q->queue_lock, so do it before bfqd->lock. */
+	if (bic)
+		bio_blkg(bio);
+#endif
+
 	spin_lock_irq(&bfqd->lock);
 
 	if (bic) {
-		/*
-		 * Make sure cgroup info is uptodate for current process before
-		 * considering the merge.
-		 */
 		bfq_bic_update_cgroup(bic, bio);
-
 		bfqd->bio_bfqq = bic_to_bfqq(bic, op_is_sync(bio->bi_opf),
 					     bfq_actuator_index(bfqd, bio));
 	} else {
 		bfqd->bio_bfqq = NULL;
 	}

@@ -6243,10 +6245,17 @@ static void bfq_insert_request(struct blk_mq_hw_ctx *hctx, struct request *rq,
 	bool idle_timer_disabled = false;
 	blk_opf_t cmd_flags;
 	LIST_HEAD(free);
 
 #ifdef CONFIG_BFQ_GROUP_IOSCHED
+	/*
+	 * Pin the blkg used to look up bfqg.  If this is the first IO for
+	 * the blkcg on this queue, create the bfqg before holding bfqd->lock.
+	 */
+	if (rq->bio && !bio_flagged(rq->bio, BIO_BLKG_REF))
+		bio_blkg(rq->bio);
+
 	if (!cgroup_subsys_on_dfl(io_cgrp_subsys) && rq->bio)
 		bfqg_stats_update_legacy_io(q, rq);
 #endif
 	spin_lock_irq(&bfqd->lock);
 	bfqq = bfq_init_rq(rq);

@@ -177,16 +177,11 @@ static inline gfp_t try_alloc_gfp(gfp_t gfp)
 		__GFP_NOMEMALLOC | __GFP_NORETRY | __GFP_NOWARN;
 }
 
 void bio_uninit(struct bio *bio)
 {
-#ifdef CONFIG_BLK_CGROUP
-	if (bio->bi_blkg) {
-		blkg_put(bio->bi_blkg);
-		bio->bi_blkg = NULL;
-	}
-#endif
+	bio_clear_blkcg(bio);
 	if (bio_integrity(bio))
 		bio_integrity_free(bio);
 
 	bio_crypt_free_ctx(bio);
 }

@@ -231,14 +226,14 @@ void bio_init(struct bio *bio, struct block_device *bdev, struct bio_vec *table,
 	bio->bi_iter.bi_idx = 0;
 	bio->bi_iter.bi_offset = 0;
 	bio->bi_end_io = NULL;
 	bio->bi_private = NULL;
 #ifdef CONFIG_BLK_CGROUP
-	bio->bi_blkg = NULL;
+	bio->bi_blkcg = NULL;
 	bio->issue_time_ns = 0;
 	if (bdev)
-		bio_associate_blkg(bio);
+		bio_associate_blkcg(bio);
 #ifdef CONFIG_BLK_CGROUP_IOCOST
 	bio->bi_iocost_cost = 0;
 #endif
 #endif
 #ifdef CONFIG_BLK_INLINE_ENCRYPTION

@@ -279,11 +274,11 @@ void bio_reset(struct bio *bio, struct block_device *bdev, blk_opf_t opf)
 	memset(bio, 0, BIO_RESET_BYTES);
 	atomic_set(&bio->__bi_remaining, 1);
 	bio->bi_io_vec = bv;
 	bio->bi_bdev = bdev;
 	if (bio->bi_bdev)
-		bio_associate_blkg(bio);
+		bio_associate_blkcg(bio);
 	bio->bi_opf = opf;
 }
 EXPORT_SYMBOL(bio_reset);
 
 /**

@@ -865,11 +860,11 @@ static int __bio_clone(struct bio *bio, struct bio *bio_src, gfp_t gfp)
 
 	if (bio->bi_bdev) {
 		if (bio->bi_bdev == bio_src->bi_bdev &&
 		    bio_flagged(bio_src, BIO_REMAPPED))
 			bio_set_flag(bio, BIO_REMAPPED);
-		bio_clone_blkg_association(bio, bio_src);
+		bio_clone_blkcg_association(bio, bio_src);
 	}
 
 	if (bio_crypt_clone(bio, bio_src, gfp) < 0)
 		return -ENOMEM;
 	if (bio_integrity(bio_src) &&

@@ -1909,21 +1904,16 @@ void bio_endio(struct bio *bio)
 	if (bio->bi_end_io == bio_chain_endio) {
 		bio = __bio_chain_endio(bio);
 		goto again;
 	}
 
-#ifdef CONFIG_BLK_CGROUP
 	/*
 	 * Release cgroup info.  We shouldn't have to do this here, but quite
 	 * a few callers of bio_init fail to call bio_uninit, so we cover up
 	 * for that here at least for now.
 	 */
-	if (bio->bi_blkg) {
-		blkg_put(bio->bi_blkg);
-		bio->bi_blkg = NULL;
-	}
-#endif
+	bio_clear_blkcg(bio);
 
 	if (bio_flagged(bio, BIO_COMPLETE_IN_TASK) && bio_in_atomic())
 		__bio_complete_in_task(bio);
 	else if (bio->bi_end_io)
 		bio->bi_end_io(bio);

@@ -48,10 +48,16 @@ EXPORT_SYMBOL_GPL(blkcg_set_fc_appid);
  *
  * On success return the fc_app_id, on failure return NULL
  */
 char *blkcg_get_fc_appid(struct bio *bio)
 {
-	if (!bio->bi_blkg || bio->bi_blkg->blkcg->fc_app_id[0] == '\0')
+	struct blkcg *blkcg = bio_blkcg(bio);
+
+	if (!blkcg)
+		return NULL;
+
+	if (blkcg->fc_app_id[0] == '\0')
 		return NULL;
-	return bio->bi_blkg->blkcg->fc_app_id;
+
+	return blkcg->fc_app_id;
 }
 EXPORT_SYMBOL_GPL(blkcg_get_fc_appid);

@@ -44,11 +44,11 @@ void blkg_rwstat_exit(struct blkg_rwstat *rwstat);
 u64 __blkg_prfill_rwstat(struct seq_file *sf, struct blkg_policy_data *pd,
 			 const struct blkg_rwstat_sample *rwstat);
 u64 blkg_prfill_rwstat(struct seq_file *sf, struct blkg_policy_data *pd,
 		       int off);
 void blkg_rwstat_recursive_sum(struct blkcg_gq *blkg, struct blkcg_policy *pol,
-		int off, struct blkg_rwstat_sample *sum);
+		int off, struct blkg_rwstat_sample *sum) __must_hold_shared(RCU);
 
 
 /**
  * blkg_rwstat_add - add a value to a blkg_rwstat
  * @rwstat: target blkg_rwstat

@@ -259,19 +259,19 @@ static void blkg_async_bio_workfn(struct work_struct *work)
  * cgroup.  Use this helper instead of submit_bio to punt the actual issuing to
  * a dedicated per-blkcg work item to avoid such priority inversions.
  */
 void blkcg_punt_bio_submit(struct bio *bio)
 {
-	struct blkcg_gq *blkg = bio->bi_blkg;
+	struct blkcg_gq *blkg = bio_blkg(bio);
 
-	if (blkg->parent) {
+	if (blkg && blkg->parent) {
 		spin_lock(&blkg->async_bio_lock);
 		bio_list_add(&blkg->async_bios, bio);
 		spin_unlock(&blkg->async_bio_lock);
 		queue_work(blkcg_punt_bio_wq, &blkg->async_bio_work);
 	} else {
-		/* never bounce for the root cgroup */
+		/* Never bounce if there is no non-root blkg to queue on. */
 		submit_bio(bio);
 	}
 }
 EXPORT_SYMBOL_GPL(blkcg_punt_bio_submit);
 

@@ -295,13 +295,17 @@ subsys_initcall(blkcg_punt_bio_init);
  * associated. Callers are expected to either handle %NULL or know association
  * has been done prior to calling this.
  */
 struct cgroup_subsys_state *bio_blkcg_css(struct bio *bio)
 {
-	if (!bio || !bio->bi_blkg)
+	struct blkcg *blkcg;
+
+	if (!bio)
 		return NULL;
-	return &bio->bi_blkg->blkcg->css;
+
+	blkcg = bio_blkcg(bio);
+	return blkcg ? &blkcg->css : NULL;
 }
 EXPORT_SYMBOL_GPL(bio_blkcg_css);
 
 /**
  * blkcg_parent - get the parent of a blkcg

@@ -423,16 +427,19 @@ static struct blkcg_gq *blkg_create(struct blkcg *blkcg, struct gendisk *disk,
 	}
 	blkg = new_blkg;
 
 	/* link parent */
 	if (blkcg_parent(blkcg)) {
+		rcu_read_lock();
 		blkg->parent = blkg_lookup(blkcg_parent(blkcg), disk->queue);
 		if (WARN_ON_ONCE(!blkg->parent)) {
+			rcu_read_unlock();
 			ret = -ENODEV;
 			goto err_free_blkg;
 		}
 		blkg_get(blkg->parent);
+		rcu_read_unlock();
 	}
 
 	/* invoke per-policy init */
 	for (i = 0; i < BLKCG_MAX_POLS; i++) {
 		struct blkcg_policy *pol = blkcg_policy[i];

@@ -473,32 +480,48 @@ static struct blkcg_gq *blkg_create(struct blkcg *blkcg, struct gendisk *disk,
 	if (new_blkg)
 		blkg_free(new_blkg);
 	return ERR_PTR(ret);
 }
 
+/*
+ * The root blkg holds a live reference while the disk is active, so walking
+ * the parent chain always finds a blkg which can be pinned.
+ */
+static struct blkcg_gq *blkg_lookup_tryget(struct blkcg_gq *blkg)
+{
+	while (!blkg_tryget(blkg))
+		blkg = blkg->parent;
+	return blkg;
+}
+
 /**
  * blkg_lookup_create - lookup blkg, try to create one if not there
  * @blkcg: blkcg of interest
  * @disk: gendisk of interest
  *
  * Lookup blkg for the @blkcg - @disk pair.  If it doesn't exist, try to
  * create one.  blkg creation is performed recursively from blkcg_root such
- * that all non-root blkg's have access to the parent blkg.  This function
- * should be called under RCU read lock and takes @disk->queue->queue_lock.
+ * that all non-root blkg's have access to the parent blkg.
  *
- * Returns the blkg or the closest blkg if blkg_create() fails as it walks
- * down from root.
+ * Must be called with @disk->queue->queue_lock held.
+ *
+ * Returns the closest blkg with an extra reference acquired.  If
+ * blkg_create() fails while walking down from root, the returned blkg may
+ * belong to an ancestor of @blkcg.  This function never returns %NULL.
  */
 static struct blkcg_gq *blkg_lookup_create(struct blkcg *blkcg,
 		struct gendisk *disk)
 {
 	struct request_queue *q = disk->queue;
 	struct blkcg_gq *blkg;
 
+	lockdep_assert_held(&q->queue_lock);
+
 	rcu_read_lock();
 	blkg = blkg_lookup(blkcg, q);
 	if (blkg) {
+		blkg = blkg_lookup_tryget(blkg);
 		rcu_read_unlock();
 		return blkg;
 	}
 	rcu_read_unlock();
 

@@ -510,31 +533,33 @@ static struct blkcg_gq *blkg_lookup_create(struct blkcg *blkcg,
 	while (true) {
 		struct blkcg *pos = blkcg;
 		struct blkcg *parent = blkcg_parent(blkcg);
 		struct blkcg_gq *ret_blkg = q->root_blkg;
 
+		rcu_read_lock();
 		while (parent) {
 			blkg = blkg_lookup(parent, q);
 			if (blkg) {
 				/* remember closest blkg */
 				ret_blkg = blkg;
 				break;
 			}
 			pos = parent;
 			parent = blkcg_parent(parent);
 		}
+		rcu_read_unlock();
 
 		blkg = blkg_create(pos, disk, NULL);
 		if (IS_ERR(blkg)) {
 			blkg = ret_blkg;
 			break;
 		}
 		if (pos == blkcg)
 			break;
 	}
 
-	return blkg;
+	return blkg_lookup_tryget(blkg);
 }
 
 static void blkg_destroy(struct blkcg_gq *blkg)
 {
 	struct blkcg *blkcg = blkg->blkcg;

@@ -859,11 +884,13 @@ int blkg_conf_prep(struct blkcg *blkcg, const struct blkcg_policy *pol,
 	if (!blkcg_policy_enabled(q, pol)) {
 		ret = -EOPNOTSUPP;
 		goto fail_unlock;
 	}
 
+	rcu_read_lock();
 	blkg = blkg_lookup(blkcg, q);
+	rcu_read_unlock();
 	if (blkg)
 		goto success;
 
 	/*
 	 * Create blkgs walking down from blkcg_root to @blkcg, so that all

@@ -873,14 +900,16 @@ int blkg_conf_prep(struct blkcg *blkcg, const struct blkcg_policy *pol,
 		struct blkcg *pos = blkcg;
 		struct blkcg *parent;
 		struct blkcg_gq *new_blkg;
 
 		parent = blkcg_parent(blkcg);
+		rcu_read_lock();
 		while (parent && !blkg_lookup(parent, q)) {
 			pos = parent;
 			parent = blkcg_parent(parent);
 		}
+		rcu_read_unlock();
 
 		/* Drop locks to do new blkg allocation with GFP_KERNEL. */
 		spin_unlock_irq(&q->queue_lock);
 
 		new_blkg = blkg_alloc(pos, disk, GFP_NOIO);

@@ -895,11 +924,13 @@ int blkg_conf_prep(struct blkcg *blkcg, const struct blkcg_policy *pol,
 			blkg_free(new_blkg);
 			ret = -EOPNOTSUPP;
 			goto fail_unlock;
 		}
 
+		rcu_read_lock();
 		blkg = blkg_lookup(pos, q);
+		rcu_read_unlock();
 		if (blkg) {
 			blkg_free(new_blkg);
 		} else {
 			blkg = blkg_create(pos, disk, new_blkg);
 			if (IS_ERR(blkg)) {

@@ -2052,134 +2083,170 @@ void blkcg_add_delay(struct blkcg_gq *blkg, u64 now, u64 delta)
 		return;
 	blkcg_scale_delay(blkg, now);
 	atomic64_add(delta, &blkg->delay_nsec);
 }
 
-static inline struct blkcg_gq *blkg_lookup_tryget(struct blkcg_gq *blkg)
+/*
+ * Return the blkg pinned by @bio through BIO_BLKG_REF.  The returned blkg is
+ * already owned by @bio and no extra reference is acquired.
+ */
+static struct blkcg_gq *bio_pinned_blkg(struct bio *bio)
 {
-retry:
-	if (blkg_tryget(blkg))
-		return blkg;
+	struct request_queue *q = bdev_get_queue(bio->bi_bdev);
+	struct blkcg_gq *blkg;
 
-	blkg = blkg->parent;
-	if (blkg)
-		goto retry;
+	rcu_read_lock();
+	blkg = blkg_lookup_any(bio_blkcg(bio), q);
+	rcu_read_unlock();
 
-	return NULL;
+	WARN_ON_ONCE(!blkg);
+	return blkg;
 }
+
+static void bio_set_blkg_ref(struct bio *bio, struct blkcg_gq *blkg)
+{
+	if (bio_blkcg(bio) != blkg->blkcg) {
+		css_get(&blkg->blkcg->css);
+		bio_clear_blkcg(bio);
+		bio->bi_blkcg = blkg->blkcg;
+	}
+	bio_set_flag(bio, BIO_BLKG_REF);
+}
+
 /**
- * blkg_tryget_closest - try and get a blkg ref on the closet blkg
+ * bio_blkg_lookup - look up a blkg associated with a bio
  * @bio: target bio
- * @css: target css
  *
- * As the failure mode here is to walk up the blkg tree, this ensure that the
- * blkg->parent pointers are always valid.  This returns the blkg that it ended
- * up taking a reference on or %NULL if no reference was taken.
+ * Return the blkg already pinned by @bio, or %NULL if @bio doesn't own a blkg
+ * reference.  Call bio_blkg() instead when a missing blkg should be created.
  */
-static inline struct blkcg_gq *blkg_tryget_closest(struct bio *bio,
-		struct cgroup_subsys_state *css)
+struct blkcg_gq *bio_blkg_lookup(struct bio *bio)
 {
-	struct request_queue *q = bio->bi_bdev->bd_queue;
-	struct blkcg *blkcg = css_to_blkcg(css);
+	if (!bio_flagged(bio, BIO_BLKG_REF))
+		return NULL;
+	return bio_pinned_blkg(bio);
+}
+EXPORT_SYMBOL_GPL(bio_blkg_lookup);
+
+/**
+ * bio_put_blkg_ref - drop the blkg reference pinned by a bio
+ * @bio: target bio
+ *
+ * Drop the bio-owned blkg reference acquired by bio_blkg(), if any.
+ */
+void bio_put_blkg_ref(struct bio *bio)
+{
+	if (bio_flagged(bio, BIO_BLKG_REF)) {
+		struct blkcg_gq *blkg = bio_pinned_blkg(bio);
+
+		if (blkg)
+			blkg_put(blkg);
+		bio_clear_flag(bio, BIO_BLKG_REF);
+	}
+}
+EXPORT_SYMBOL_GPL(bio_put_blkg_ref);
+
+/**
+ * bio_blkg - look up the blkg associated with a bio
+ * @bio: target bio
+ *
+ * Look up the queue-local blkg for @bio's current device and blkcg.  If this
+ * is the first policy use of @bio, create the missing blkg hierarchy if
+ * necessary, pin the closest available blkg, and update @bio's blkcg
+ * association if allocation falls back to an ancestor.
+ */
+struct blkcg_gq *bio_blkg(struct bio *bio)
+{
+	struct blkcg *blkcg = bio_blkcg(bio);
+	struct gendisk *disk;
+	struct request_queue *q;
 	struct blkcg_gq *blkg;
 	unsigned long flags;
 
-	rcu_read_lock();
-	blkg = blkg_lookup(blkcg, q);
-	if (likely(blkg))
-		blkg = blkg_lookup_tryget(blkg);
-	rcu_read_unlock();
+	if (!blkcg || !bio->bi_bdev)
+		return NULL;
 
-	if (blkg)
-		return blkg;
+	if (bio_flagged(bio, BIO_BLKG_REF))
+		return bio_pinned_blkg(bio);
+
+	disk = bio->bi_bdev->bd_disk;
+	q = disk->queue;
 
-	/*
-	 * Fast path failed, we're probably issuing IO in this cgroup the first
-	 * time, hold lock to create new blkg.
-	 */
 	spin_lock_irqsave(&q->queue_lock, flags);
-	blkg = blkg_lookup_create(blkcg, bio->bi_bdev->bd_disk);
-	if (blkg)
-		blkg = blkg_lookup_tryget(blkg);
+	blkg = blkg_lookup_create(blkcg, disk);
 	spin_unlock_irqrestore(&q->queue_lock, flags);
 
+	bio_set_blkg_ref(bio, blkg);
 	return blkg;
 }
+EXPORT_SYMBOL_GPL(bio_blkg);
 
 /**
- * bio_associate_blkg_from_css - associate a bio with a specified css
+ * bio_associate_blkcg_from_css - associate a bio with a specified css
  * @bio: target bio
  * @css: target css
  *
- * Associate @bio with the blkg found by combining the css's blkg and the
- * request_queue of the @bio.  An association failure is handled by walking up
- * the blkg tree.  Therefore, the blkg associated can be anything between @blkg
- * and q->root_blkg.  This situation only happens when a cgroup is dying and
- * then the remaining bios will spill to the closest alive blkg.
+ * Associate @bio with the blkcg found from @css.  The queue-local blkg is
+ * created and pinned by bio_blkg() when blkcg policies need it.  If @css
+ * can't be referenced online, associate @bio with the root blkcg instead.
  *
- * A reference will be taken on the blkg and will be released when @bio is
+ * A reference will be taken on the blkcg and will be released when @bio is
  * freed.
  */
-void bio_associate_blkg_from_css(struct bio *bio,
-				 struct cgroup_subsys_state *css)
+void bio_associate_blkcg_from_css(struct bio *bio,
+				  struct cgroup_subsys_state *css)
 {
-	if (bio->bi_blkg)
-		blkg_put(bio->bi_blkg);
+	struct blkcg *blkcg;
 
-	if (css && css->parent) {
-		bio->bi_blkg = blkg_tryget_closest(bio, css);
-	} else {
-		blkg_get(bdev_get_queue(bio->bi_bdev)->root_blkg);
-		bio->bi_blkg = bdev_get_queue(bio->bi_bdev)->root_blkg;
+	if (!css || !css_tryget_online(css)) {
+		css = &blkcg_root.css;
+		css_get(css);
+	}
+
+	blkcg = css_to_blkcg(css);
+	if (bio_blkcg(bio) == blkcg) {
+		css_put(css);
+		return;
 	}
+
+	bio_clear_blkcg(bio);
+	bio->bi_blkcg = blkcg;
 }
-EXPORT_SYMBOL_GPL(bio_associate_blkg_from_css);
+EXPORT_SYMBOL_GPL(bio_associate_blkcg_from_css);
 
 /**
- * bio_associate_blkg - associate a bio with a blkg
+ * bio_associate_blkcg - associate a bio with a blkcg
  * @bio: target bio
  *
- * Associate @bio with the blkg found from the bio's css and request_queue.
- * If one is not found, bio_lookup_blkg() creates the blkg.  If a blkg is
- * already associated, the css is reused and association redone as the
- * request_queue may have changed.
+ * Associate @bio with the blkcg found from the bio's css.  If a blkcg is
+ * already associated, keep it as blkcg association is not queue-local.
  */
-void bio_associate_blkg(struct bio *bio)
+void bio_associate_blkcg(struct bio *bio)
 {
-	struct cgroup_subsys_state *css;
-
 	if (blk_op_is_passthrough(bio->bi_opf))
 		return;
 
-	if (bio->bi_blkg) {
-		css = bio_blkcg_css(bio);
-		bio_associate_blkg_from_css(bio, css);
-	} else {
-		rcu_read_lock();
-		css = blkcg_css();
-		if (!css_tryget_online(css))
-			css = NULL;
-		rcu_read_unlock();
+	if (bio_blkcg(bio))
+		return;
 
-		bio_associate_blkg_from_css(bio, css);
-		if (css)
-			css_put(css);
-	}
+	rcu_read_lock();
+	bio_associate_blkcg_from_css(bio, blkcg_css());
+	rcu_read_unlock();
 }
-EXPORT_SYMBOL_GPL(bio_associate_blkg);
+EXPORT_SYMBOL_GPL(bio_associate_blkcg);
 
 /**
- * bio_clone_blkg_association - clone blkg association from src to dst bio
+ * bio_clone_blkcg_association - clone blkcg association from src to dst bio
  * @dst: destination bio
  * @src: source bio
  */
-void bio_clone_blkg_association(struct bio *dst, struct bio *src)
+void bio_clone_blkcg_association(struct bio *dst, struct bio *src)
 {
-	if (src->bi_blkg)
-		bio_associate_blkg_from_css(dst, bio_blkcg_css(src));
+	if (bio_blkcg(src))
+		bio_associate_blkcg_from_css(dst, bio_blkcg_css(src));
 }
-EXPORT_SYMBOL_GPL(bio_clone_blkg_association);
+EXPORT_SYMBOL_GPL(bio_clone_blkcg_association);
 
 static int blk_cgroup_io_type(struct bio *bio)
 {
 	if (op_is_discard(bio->bi_opf))
 		return BLKG_IOSTAT_DISCARD;