Re: [PATCH v2 1/2] ublk: honor rq_affinity on request completion

[Ming Lei <tom.leiming@gmail.com>]

On Tue, Oct 06, 2026 at 10:03:47PM +0300, Yoav Cohen wrote:
> 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(). This covers every completion
> outcome -- reads, writes, flush/discard, errors, and the zero-copy/
> user-copy/auto-buf-reg/shmem-zc paths that skip the copy-back step --
> not just the classic-copy-mode successful-read case, since
> ublk_end_rq() re-derives which of those __ublk_complete_rq() took from
> @req and @io state rather than assuming one.
> 
> Gate it behind a new UBLK_F_SUPPORT_RQ_AFFINITY flag that a server
> must request at ADD_DEV time; servers that don't ask for it keep
> today's always-local completion unchanged. Add ublk_support_rq_affinity()
> alongside the existing ublk_support_*() helpers, and check it before
> calling blk_mq_complete_request_remote() at both call sites.
> 
> Signed-off-by: Yoav Cohen <yoav@nvidia.com>
> ---
>  drivers/block/ublk_drv.c      | 98 +++++++++++++++++++++++++++--------
>  include/uapi/linux/ublk_cmd.h |  8 +++
>  2 files changed, 84 insertions(+), 22 deletions(-)
> 
> diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
> index 66eb55e7162e..b94293296b73 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 \
> @@ -452,6 +453,11 @@ static inline bool ublk_support_user_copy(const struct ublk_queue *ubq)
>  	return ubq->flags & UBLK_F_USER_COPY;
>  }
>  
> +static inline bool ublk_support_rq_affinity(const struct ublk_queue *ubq)
> +{
> +	return ubq->flags & UBLK_F_SUPPORT_RQ_AFFINITY;
> +}
> +
>  static inline bool ublk_dev_support_user_copy(const struct ublk_device *ub)
>  {
>  	return ub->dev_info.flags & UBLK_F_USER_COPY;
> @@ -1468,6 +1474,14 @@ static inline bool ublk_need_unmap_req(const struct request *req)
>  	       (req_op(req) == REQ_OP_READ || req_op(req) == REQ_OP_DRV_IN);
>  }
>  
> +/* Whether __ublk_complete_rq() takes the copy-back/partial-completion path */
> +static inline bool ublk_rq_need_unmap(const struct ublk_queue *ubq,
> +				      struct request *req)
> +{
> +	return ublk_need_map_io(ubq) && ublk_need_unmap_req(req) &&
> +	       !ublk_iod_is_shmem_zc(ubq, req->tag);
> +}
> +
>  static unsigned int ublk_map_io(const struct request *req,
>  				const struct ublk_io *io)
>  {
> @@ -1549,13 +1563,62 @@ 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];
> +
> +	/* matches the two outcomes __ublk_complete_rq() may have redirected */
> +	if (io->res >= 0 && ublk_rq_need_unmap(ubq, req)) {
> +		if (!ublk_update_and_requeue(req, io) &&
> +		    likely(!blk_should_fake_timeout(req->q)))
> +			__blk_mq_end_request(req, BLK_STS_OK);
> +		return;
> +	}

ublk_end_rq() may be run from remove cpu, then it becomes fragile to
touch `io` here, and not efficient to touch `ubq` as Caleb mentioned.

It could be simplified a lot to just deal with normal io->res == blk_rq_bytes(req)
here, meantime keep everything else(short write, failure, ...) in __ublk_complete_rq().


Thanks, 
Ming