[RFC PATCH v3 6/6] blk-cgroup: make policy blkg creation nowait-safe

[Yu Kuai <yukuai@kernel.org>]

From: Yu Kuai <yukuai@fygo.io>

bio_blkg() is called by blkcg policy paths when they need a queue-local
blkg. Keep that allocation lazy instead of preparing every REQ_NOWAIT bio
from submit_bio_noacct().

If a policy first needs a blkg for a REQ_NOWAIT bio, use mutex_trylock()
and GFP_ATOMIC so the lookup never sleeps. If the mutex cannot be acquired,
look up and pin the closest existing blkg in the hierarchy under RCU. The
creation helper provides the same fallback if atomic allocation fails, so
valid policy I/O always gets a blkg without blocking.

Signed-off-by: Yu Kuai <yukuai@fygo.io>
---
 block/blk-cgroup.c | 35 ++++++++++++++++++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

diff --git a/block/blk-cgroup.c b/block/blk-cgroup.c
index 31afb433ab18..9895d6661070 100644
--- a/block/blk-cgroup.c
+++ b/block/blk-cgroup.c
@@ -28,10 +28,11 @@
 #include <linux/atomic.h>
 #include <linux/ctype.h>
 #include <linux/resume_user_mode.h>
 #include <linux/psi.h>
 #include <linux/part_stat.h>
+#include <linux/preempt.h>
 #include "blk.h"
 #include "blk-cgroup.h"
 #include "blk-ioprio.h"
 #include "blk-throttle.h"
 
@@ -469,10 +470,24 @@ static struct blkcg_gq *blkg_lookup_tryget(struct blkcg_gq *blkg)
 	while (!blkg_tryget(blkg))
 		blkg = blkg->parent;
 	return blkg;
 }
 
+static struct blkcg_gq *blkg_lookup_closest(struct blkcg *blkcg,
+					    struct request_queue *q)
+{
+	struct blkcg_gq *blkg;
+
+	rcu_read_lock();
+	while (!(blkg = blkg_lookup(blkcg, q)))
+		blkcg = blkcg_parent(blkcg);
+	blkg = blkg_lookup_tryget(blkg);
+	rcu_read_unlock();
+
+	return blkg;
+}
+
 /**
  * blkg_lookup_create - lookup blkg, try to create one if not there
  * @blkcg: blkcg of interest
  * @disk: gendisk of interest
  * @gfp_mask: allocation mask to use
@@ -2066,11 +2081,10 @@ 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;
-	int ret;
 
 	if (!blkcg || !bio->bi_bdev)
 		return NULL;
 
 	if (bio_flagged(bio, BIO_BLKG_REF))
@@ -2087,10 +2101,29 @@ struct blkcg_gq *bio_blkg(struct bio *bio)
 	if (blkg) {
 		bio_set_blkg_ref(bio, blkg);
 		return blkg;
 	}
 
+	if (bio->bi_opf & REQ_NOWAIT) {
+		/*
+		 * Nowait callers must not sleep on the mutex nor allocate with
+		 * sleeping GFPs.  Trylock the mutex and create the missing blkg
+		 * atomically.  If the mutex cannot be acquired, skip allocation
+		 * and pin the closest existing blkg instead.  blkg_lookup_create()
+		 * provides the same fallback if allocation fails.
+		 */
+		if (!preemptible() || !mutex_trylock(&q->blkcg_mutex)) {
+			blkg = blkg_lookup_closest(blkcg, q);
+		} else {
+			blkg_lookup_create(blkcg, disk, GFP_ATOMIC, &blkg);
+			blkg = blkg_lookup_tryget(blkg);
+			mutex_unlock(&q->blkcg_mutex);
+		}
+		bio_set_blkg_ref(bio, blkg);
+		return blkg;
+	}
+
 	mutex_lock(&q->blkcg_mutex);
 	blkg_lookup_create(blkcg, disk, GFP_NOIO, &blkg);
 	blkg = blkg_lookup_tryget(blkg);
 	mutex_unlock(&q->blkcg_mutex);
 
-- 
2.51.0