[PATCH 1/2] ublk: complete requests via blk_mq_complete_request() so rq_affinity applies

[Yoav Cohen <yoav@nvidia.com>]

ublk ends successfully completed requests inline on the server thread
that issued UBLK_IO_COMMIT_AND_FETCH_REQ, so rq_affinity never
redirects completion back to the submitting CPU, unlike NVMe, SCSI,
virtio-blk, loop, nbd and rnbd.

Route completion through blk_mq_complete_request_remote(), with a new
->complete() callback, ublk_end_rq(). Add UBLK_F_SUPPORT_RQ_AFFINITY
so servers can detect support via UBLK_CMD_GET_FEATURES.

Signed-off-by: Yoav Cohen <yoav@nvidia.com>
---
 drivers/block/ublk_drv.c      | 72 +++++++++++++++++++++++++----------
 include/uapi/linux/ublk_cmd.h |  8 ++++
 2 files changed, 59 insertions(+), 21 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index 66eb55e7162e..48bdb5d782ee 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -90,7 +90,8 @@
 		| UBLK_F_BATCH_IO \
 		| UBLK_F_NO_AUTO_PART_SCAN \
 		| UBLK_F_SHMEM_ZC \
-		| UBLK_F_IO_DESC_SIZE)
+		| UBLK_F_IO_DESC_SIZE \
+		| UBLK_F_SUPPORT_RQ_AFFINITY)
 
 #define UBLK_F_ALL_RECOVERY_FLAGS (UBLK_F_USER_RECOVERY \
 		| UBLK_F_USER_RECOVERY_REISSUE \
@@ -1549,13 +1550,54 @@ static void ublk_end_request(struct request *req, blk_status_t error)
 	local_bh_enable();
 }
 
+/*
+ * Update @req with its result and requeue it if it was only partially
+ * completed. Returns true if @req was requeued, in which case the caller
+ * must not touch it any further.
+ *
+ * Run bio->bi_end_io() with softirqs disabled. If the final fput happens
+ * off this path, then that will prevent ublk's blkdev_release() from
+ * being called on current's task work, see fput() implementation.
+ *
+ * This matters for the caller completing @req locally on the ublk
+ * server's own thread: it may already be holding disk->open_mutex, e.g.
+ * reading the partition table from bdev_open(), and an fput() running
+ * inline there could deadlock on it. Preferably we would not be doing
+ * IO with a mutex held that is also used for release, but this
+ * work-around will suffice for now. A caller reached instead via
+ * blk_mq_complete_request_remote()'s softirq/IPI redirect never runs on
+ * that thread, so it isn't exposed to this hazard, but disabling
+ * softirqs here is harmless for it too.
+ */
+static inline bool ublk_update_and_requeue(struct request *req,
+					   struct ublk_io *io)
+{
+	bool requeue;
+
+	local_bh_disable();
+	requeue = blk_update_request(req, BLK_STS_OK, io->res);
+	local_bh_enable();
+	if (requeue)
+		blk_mq_requeue_request(req, true);
+	return requeue;
+}
+
+static void ublk_end_rq(struct request *req)
+{
+	struct ublk_queue *ubq = req->mq_hctx->driver_data;
+	struct ublk_io *io = &ubq->ios[req->tag];
+
+	if (!ublk_update_and_requeue(req, io) &&
+	    likely(!blk_should_fake_timeout(req->q)))
+		__blk_mq_end_request(req, BLK_STS_OK);
+}
+
 /* todo: handle partial completion */
 static inline void __ublk_complete_rq(struct request *req, struct ublk_io *io,
 				      bool need_map, struct io_comp_batch *iob)
 {
 	unsigned int unmapped_bytes;
 	blk_status_t res = BLK_STS_OK;
-	bool requeue;
 
 	/* failed read IO if nothing is read */
 	if (!io->res && req_op(req) == REQ_OP_READ)
@@ -1588,25 +1630,11 @@ static inline void __ublk_complete_rq(struct request *req, struct ublk_io *io,
 		io->res = unmapped_bytes;
 	}
 
-	/*
-	 * Run bio->bi_end_io() with softirqs disabled. If the final fput
-	 * happens off this path, then that will prevent ublk's blkdev_release()
-	 * from being called on current's task work, see fput() implementation.
-	 *
-	 * Otherwise, ublk server may not provide forward progress in case of
-	 * reading the partition table from bdev_open() with disk->open_mutex
-	 * held, and causes dead lock as we could already be holding
-	 * disk->open_mutex here.
-	 *
-	 * Preferably we would not be doing IO with a mutex held that is also
-	 * used for release, but this work-around will suffice for now.
-	 */
-	local_bh_disable();
-	requeue = blk_update_request(req, BLK_STS_OK, io->res);
-	local_bh_enable();
-	if (requeue)
-		blk_mq_requeue_request(req, true);
-	else if (likely(!blk_should_fake_timeout(req->q))) {
+	if (blk_mq_complete_request_remote(req))
+		return;
+
+	if (!ublk_update_and_requeue(req, io) &&
+	    likely(!blk_should_fake_timeout(req->q))) {
 		if (blk_mq_add_to_batch(req, iob, false, blk_mq_end_request_batch))
 			return;
 		__blk_mq_end_request(req, BLK_STS_OK);
@@ -2347,6 +2375,7 @@ static const struct blk_mq_ops ublk_mq_ops = {
 	.queue_rqs      = ublk_queue_rqs,
 	.init_hctx	= ublk_init_hctx,
 	.timeout	= ublk_timeout,
+	.complete	= ublk_end_rq,
 };
 
 static const struct blk_mq_ops ublk_batch_mq_ops = {
@@ -2355,6 +2384,7 @@ static const struct blk_mq_ops ublk_batch_mq_ops = {
 	.queue_rqs      = ublk_batch_queue_rqs,
 	.init_hctx	= ublk_init_hctx,
 	.timeout	= ublk_timeout,
+	.complete	= ublk_end_rq,
 };
 
 static void ublk_queue_reinit(struct ublk_device *ub, struct ublk_queue *ubq)
diff --git a/include/uapi/linux/ublk_cmd.h b/include/uapi/linux/ublk_cmd.h
index 33b25dd13965..47b5fba9519c 100644
--- a/include/uapi/linux/ublk_cmd.h
+++ b/include/uapi/linux/ublk_cmd.h
@@ -420,6 +420,14 @@ struct ublk_shmem_buf_reg {
 /* ublksrv_io_desc size is specified by ublksrv_ctrl_dev_info's io_desc_size */
 #define UBLK_F_IO_DESC_SIZE (1ULL << 20)
 
+/*
+ * Request completion honors the block device's rq_affinity setting
+ * (/sys/block/ublkbN/queue/rq_affinity): the submitting CPU's completion
+ * work can run there instead of always on the ublk server's CPU. Without
+ * this feature, rq_affinity has no effect on ublk devices.
+ */
+#define UBLK_F_SUPPORT_RQ_AFFINITY (1ULL << 21)
+
 /* device state */
 #define UBLK_S_DEV_DEAD	0
 #define UBLK_S_DEV_LIVE	1
-- 
2.50.1 (Apple Git-155)