char *input, struct blkg_conf_ctx *ctx);
 void blkg_conf_finish(struct blkg_conf_ctx *ctx);
 
+/**
+ * blkcg_css - find the current css
+ *
+ * Find the css associated with either the kthread or the current task.
+ * This may return a dying css, so it is up to the caller to use tryget logic
+ * to confirm it is alive and well.
+ */
+static inline struct cgroup_subsys_state *blkcg_css(void)
+{
+       struct cgroup_subsys_state *css;
+
+       css = kthread_blkcg();
+       if (css)
+               return css;
+       return task_css(current, io_cgrp_id);
+}
+
+/**
+ * blkcg_get_css - find and get a reference to the css
+ *
+ * Find the css associated with either the kthread or the current task.
+ * This takes a reference on the blkcg which will need to be managed by the
+ * caller.
+ */
+static inline struct cgroup_subsys_state *blkcg_get_css(void)
+{
+       struct cgroup_subsys_state *css;
+
+       rcu_read_lock();
+
+       css = kthread_blkcg();
+       if (css) {
+               css_get(css);
+       } else {
+               /*
+                * This is a bit complicated.  It is possible task_css is seeing
+                * an old css pointer here.  This is caused by the current
+                * thread migrating away from this cgroup and this cgroup dying.
+                * css_tryget() will fail when trying to take a ref on a cgroup
+                * that's ref count has hit 0.
+                *
+                * Therefore, if it does fail, this means current must have
+                * been swapped away already and this is waiting for it to
+                * propagate on the polling cpu.  Hence the use of cpu_relax().
+                */
+               while (true) {
+                       css = task_css(current, io_cgrp_id);
+                       if (likely(css_tryget(css)))
+                               break;
+                       cpu_relax();
+               }
+       }
+
+       rcu_read_unlock();
+
+       return css;
+}
 
 static inline struct blkcg *css_to_blkcg(struct cgroup_subsys_state *css)
 {
        return css ? container_of(css, struct blkcg, css) : NULL;
 }
 
-static inline struct blkcg *bio_blkcg(struct bio *bio)
+/**
+ * __bio_blkcg - internal version of bio_blkcg for bfq and cfq
+ *
+ * DO NOT USE.
+ * There is a flaw using this version of the function.  In particular, this was
+ * used in a broken paradigm where association was called on the given css.  It
+ * is possible though that the returned css from task_css() is in the process
+ * of dying due to migration of the current task.  So it is improper to assume
+ * *_get() is going to succeed.  Both BFQ and CFQ rely on this logic and will
+ * take additional work to handle more gracefully.
+ */
+static inline struct blkcg *__bio_blkcg(struct bio *bio)
 {
-       struct cgroup_subsys_state *css;
+       if (bio && bio->bi_css)
+               return css_to_blkcg(bio->bi_css);
+       return css_to_blkcg(blkcg_css());
+}
 
+/**
+ * bio_blkcg - grab the blkcg associated with a bio
+ * @bio: target bio
+ *
+ * This returns the blkcg associated with a bio, NULL if not associated.
+ * Callers are expected to either handle NULL or know association has been
+ * done prior to calling this.
+ */
+static inline struct blkcg *bio_blkcg(struct bio *bio)
+{
        if (bio && bio->bi_css)
                return css_to_blkcg(bio->bi_css);
-       css = kthread_blkcg();
-       if (css)
-               return css_to_blkcg(css);
-       return css_to_blkcg(task_css(current, io_cgrp_id));
+       return NULL;
 }
 
 static inline bool blk_cgroup_congested(void)
        rcu_read_lock();
 
        blkcg = bio_blkcg(bio);
+       if (blkcg)
+               css_get(&blkcg->css);
+       else
+               blkcg = css_to_blkcg(blkcg_get_css());
 
        /* bypass blkg lookup and use @q->root_rl directly for root */
        if (blkcg == &blkcg_root)
  */
 static inline void blk_put_rl(struct request_list *rl)
 {
+       /* an additional ref is always taken for rl */
+       css_put(&rl->blkg->blkcg->css);
        if (rl->blkg->blkcg != &blkcg_root)
                blkg_put(rl->blkg);
 }
        bool throtl = false;
 
        rcu_read_lock();
-       blkcg = bio_blkcg(bio);
 
        /* associate blkcg if bio hasn't attached one */
-       bio_associate_blkcg(bio, &blkcg->css);
+       bio_associate_blkcg(bio, NULL);
+       blkcg = bio_blkcg(bio);
 
        blkg = blkg_lookup(blkcg, q);
        if (unlikely(!blkg)) {
 static inline void blkcg_deactivate_policy(struct request_queue *q,
                                           const struct blkcg_policy *pol) { }
 
+static inline struct blkcg *__bio_blkcg(struct bio *bio) { return NULL; }
 static inline struct blkcg *bio_blkcg(struct bio *bio) { return NULL; }
 
 static inline struct blkg_policy_data *blkg_to_pd(struct blkcg_gq *blkg,