[PATCH v2 00/14] ublk: allow off-daemon zero-copy buffer registration

[PATCH v2 00/14] ublk: allow off-daemon zero-copy buffer registration

[Caleb Sander Mateos]

Currently ublk zero-copy requires ublk request buffers to be registered
and unregistered by the ublk I/O's daemon task. However, as currently
implemented, there is no reason for this restriction. Registration looks
up the request via the ublk device's tagset rather than the daemon-local
ublk_io structure and takes an atomic reference to prevent racing with
dispatch or completion of the request. Ming has expressed interest in
relaxing this restriction[1] so the ublk server can offload the I/O
operation that uses the zero-copy buffer to another thread.

Additionally, optimize the buffer registration for the common case where
the buffer is registered and unregistered by the daemon task. Skip the
expensive atomic reference count increment and decrement and the several
pointer dereferences involved in looking up the request on the tagset.
On our ublk server threads, this results in a 24% decrease in the CPU
time spent in the ublk functions on a 4K zero-copy read workload, from
8.75% to 6.68%.

Two preliminary fixes are included as well:
- Move the ublk request reference count from ublk_rq_data (a tail
  allocation of struct request) to ublk_io to prevent a use after free
  if the struct request is freed concurrently with taking a reference.
- Don't allocate physically contiguous memory for __queues, which can be
  very large. Doubling the size of struct ublk_io caused ENOMEM errors
  when adding a ublk device before this change.

[1]: https://lore.kernel.org/linux-block/aAmYJxaV1-yWEMRo@fedora/

v2:
- Add 2 fix patches
- Optimize buffer unregistration too
- Cache-align ublk_io
- Restore check for zero copy support in ublk_unregister_io_buf()
- Check for registered buffer count overflow

Caleb Sander Mateos (14):
  ublk: use vmalloc for ublk_device's __queues
  ublk: remove struct ublk_rq_data
  ublk: check cmd_op first
  ublk: handle UBLK_IO_FETCH_REQ earlier
  ublk: remove task variable from __ublk_ch_uring_cmd()
  ublk: consolidate UBLK_IO_FLAG_{ACTIVE,OWNED_BY_SRV} checks
  ublk: move ublk_prep_cancel() to case UBLK_IO_COMMIT_AND_FETCH_REQ
  ublk: don't take ublk_queue in ublk_unregister_io_buf()
  ublk: allow UBLK_IO_(UN)REGISTER_IO_BUF on any task
  ublk: return early if blk_should_fake_timeout()
  ublk: optimize UBLK_IO_REGISTER_IO_BUF on daemon task
  ublk: optimize UBLK_IO_UNREGISTER_IO_BUF on daemon task
  ublk: remove ubq checks from ublk_{get,put}_req_ref()
  ublk: cache-align struct ublk_io

 drivers/block/ublk_drv.c      | 283 +++++++++++++++++++++-------------
 include/uapi/linux/ublk_cmd.h |  10 ++
 2 files changed, 187 insertions(+), 106 deletions(-)

-- 
2.45.2

[PATCH v2 01/14] ublk: use vmalloc for ublk_device's __queues

[Caleb Sander Mateos]

struct ublk_device's __queues points to an allocation with up to
UBLK_MAX_NR_QUEUES (4096) queues, each of which have:
- struct ublk_queue (48 bytes)
- Tail array of up to UBLK_MAX_QUEUE_DEPTH (4096) struct ublk_io's,
  32 bytes each
This means the full allocation can exceed 512 MB, which may well be
impossible to service with contiguous physical pages. Switch to
kvcalloc() and kvfree(), since there is no need for physically
contiguous memory.

Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
Fixes: 71f28f3136af ("ublk_drv: add io_uring based userspace block driver")
---
 drivers/block/ublk_drv.c | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index d36f44f5ee80..467769fc388f 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -2486,22 +2486,22 @@ static void ublk_deinit_queues(struct ublk_device *ub)
 	if (!ub->__queues)
 		return;
 
 	for (i = 0; i < nr_queues; i++)
 		ublk_deinit_queue(ub, i);
-	kfree(ub->__queues);
+	kvfree(ub->__queues);
 }
 
 static int ublk_init_queues(struct ublk_device *ub)
 {
 	int nr_queues = ub->dev_info.nr_hw_queues;
 	int depth = ub->dev_info.queue_depth;
 	int ubq_size = sizeof(struct ublk_queue) + depth * sizeof(struct ublk_io);
 	int i, ret = -ENOMEM;
 
 	ub->queue_size = ubq_size;
-	ub->__queues = kcalloc(nr_queues, ubq_size, GFP_KERNEL);
+	ub->__queues = kvcalloc(nr_queues, ubq_size, GFP_KERNEL);
 	if (!ub->__queues)
 		return ret;
 
 	for (i = 0; i < nr_queues; i++) {
 		if (ublk_init_queue(ub, i))
-- 
2.45.2

[PATCH v2 02/14] ublk: remove struct ublk_rq_data

[Caleb Sander Mateos]

__ublk_check_and_get_req() attempts to atomically look up the struct
request for a ublk I/O and take a reference on it. However, the request
can be freed between the lookup on the tagset in blk_mq_tag_to_rq() and
the increment of its reference count in ublk_get_req_ref(), for example
if an elevator switch happens concurrently.

Fix the potential use after free by moving the reference count from
ublk_rq_data to ublk_io. Move the fields buf_index and buf_ctx_handle
too to reduce the number of cache lines touched when dispatching and
completing a ublk I/O, allowing ublk_rq_data to be removed entirely.

Suggested-by: Ming Lei <ming.lei@redhat.com>
Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
Fixes: 62fe99cef94a ("ublk: add read()/write() support for ublk char device")
---
 drivers/block/ublk_drv.c | 131 +++++++++++++++++++++------------------
 1 file changed, 71 insertions(+), 60 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index 467769fc388f..55d5435f1ee2 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -80,18 +80,10 @@
 #define UBLK_PARAM_TYPE_ALL                                \
 	(UBLK_PARAM_TYPE_BASIC | UBLK_PARAM_TYPE_DISCARD | \
 	 UBLK_PARAM_TYPE_DEVT | UBLK_PARAM_TYPE_ZONED |    \
 	 UBLK_PARAM_TYPE_DMA_ALIGN | UBLK_PARAM_TYPE_SEGMENT)
 
-struct ublk_rq_data {
-	refcount_t ref;
-
-	/* for auto-unregister buffer in case of UBLK_F_AUTO_BUF_REG */
-	u16 buf_index;
-	void *buf_ctx_handle;
-};
-
 struct ublk_uring_cmd_pdu {
 	/*
 	 * Store requests in same batch temporarily for queuing them to
 	 * daemon context.
 	 *
@@ -167,10 +159,26 @@ struct ublk_io {
 		/* valid if UBLK_IO_FLAG_OWNED_BY_SRV is set */
 		struct request *req;
 	};
 
 	struct task_struct *task;
+
+	/*
+	 * The number of uses of this I/O by the ublk server
+	 * if user copy or zero copy are enabled:
+	 * - 1 from dispatch to the server until UBLK_IO_COMMIT_AND_FETCH_REQ
+	 * - 1 for each inflight ublk_ch_{read,write}_iter() call
+	 * - 1 for each io_uring registered buffer
+	 * The I/O can only be completed once all references are dropped.
+	 * User copy and buffer registration operations are only permitted
+	 * if the reference count is nonzero.
+	 */
+	refcount_t ref;
+
+	/* auto-registered buffer, valid if UBLK_IO_FLAG_AUTO_BUF_REG is set */
+	u16 buf_index;
+	void *buf_ctx_handle;
 };
 
 struct ublk_queue {
 	int q_id;
 	int q_depth;
@@ -226,11 +234,12 @@ struct ublk_params_header {
 
 static void ublk_io_release(void *priv);
 static void ublk_stop_dev_unlocked(struct ublk_device *ub);
 static void ublk_abort_queue(struct ublk_device *ub, struct ublk_queue *ubq);
 static inline struct request *__ublk_check_and_get_req(struct ublk_device *ub,
-		const struct ublk_queue *ubq, int tag, size_t offset);
+		const struct ublk_queue *ubq, struct ublk_io *io,
+		size_t offset);
 static inline unsigned int ublk_req_build_flags(struct request *req);
 
 static inline struct ublksrv_io_desc *
 ublk_get_iod(const struct ublk_queue *ubq, unsigned tag)
 {
@@ -671,38 +680,30 @@ static inline bool ublk_need_req_ref(const struct ublk_queue *ubq)
 	return ublk_support_user_copy(ubq) || ublk_support_zero_copy(ubq) ||
 		ublk_support_auto_buf_reg(ubq);
 }
 
 static inline void ublk_init_req_ref(const struct ublk_queue *ubq,
-		struct request *req)
+		struct ublk_io *io)
 {
-	if (ublk_need_req_ref(ubq)) {
-		struct ublk_rq_data *data = blk_mq_rq_to_pdu(req);
-
-		refcount_set(&data->ref, 1);
-	}
+	if (ublk_need_req_ref(ubq))
+		refcount_set(&io->ref, 1);
 }
 
 static inline bool ublk_get_req_ref(const struct ublk_queue *ubq,
-		struct request *req)
+		struct ublk_io *io)
 {
-	if (ublk_need_req_ref(ubq)) {
-		struct ublk_rq_data *data = blk_mq_rq_to_pdu(req);
-
-		return refcount_inc_not_zero(&data->ref);
-	}
+	if (ublk_need_req_ref(ubq))
+		return refcount_inc_not_zero(&io->ref);
 
 	return true;
 }
 
 static inline void ublk_put_req_ref(const struct ublk_queue *ubq,
-		struct request *req)
+		struct ublk_io *io, struct request *req)
 {
 	if (ublk_need_req_ref(ubq)) {
-		struct ublk_rq_data *data = blk_mq_rq_to_pdu(req);
-
-		if (refcount_dec_and_test(&data->ref))
+		if (refcount_dec_and_test(&io->ref))
 			__ublk_complete_rq(req);
 	} else {
 		__ublk_complete_rq(req);
 	}
 }
@@ -1179,55 +1180,54 @@ static inline void __ublk_abort_rq(struct ublk_queue *ubq,
 		blk_mq_requeue_request(rq, false);
 	else
 		blk_mq_end_request(rq, BLK_STS_IOERR);
 }
 
-static void ublk_auto_buf_reg_fallback(struct request *req)
+static void
+ublk_auto_buf_reg_fallback(const struct ublk_queue *ubq, struct ublk_io *io)
 {
-	const struct ublk_queue *ubq = req->mq_hctx->driver_data;
-	struct ublksrv_io_desc *iod = ublk_get_iod(ubq, req->tag);
-	struct ublk_rq_data *data = blk_mq_rq_to_pdu(req);
+	unsigned tag = io - ubq->ios;
+	struct ublksrv_io_desc *iod = ublk_get_iod(ubq, tag);
 
 	iod->op_flags |= UBLK_IO_F_NEED_REG_BUF;
-	refcount_set(&data->ref, 1);
+	refcount_set(&io->ref, 1);
 }
 
-static bool ublk_auto_buf_reg(struct request *req, struct ublk_io *io,
-			      unsigned int issue_flags)
+static bool ublk_auto_buf_reg(const struct ublk_queue *ubq, struct request *req,
+			      struct ublk_io *io, unsigned int issue_flags)
 {
 	struct ublk_uring_cmd_pdu *pdu = ublk_get_uring_cmd_pdu(io->cmd);
-	struct ublk_rq_data *data = blk_mq_rq_to_pdu(req);
 	int ret;
 
 	ret = io_buffer_register_bvec(io->cmd, req, ublk_io_release,
 				      pdu->buf.index, issue_flags);
 	if (ret) {
 		if (pdu->buf.flags & UBLK_AUTO_BUF_REG_FALLBACK) {
-			ublk_auto_buf_reg_fallback(req);
+			ublk_auto_buf_reg_fallback(ubq, io);
 			return true;
 		}
 		blk_mq_end_request(req, BLK_STS_IOERR);
 		return false;
 	}
 	/* one extra reference is dropped by ublk_io_release */
-	refcount_set(&data->ref, 2);
+	refcount_set(&io->ref, 2);
 
-	data->buf_ctx_handle = io_uring_cmd_ctx_handle(io->cmd);
+	io->buf_ctx_handle = io_uring_cmd_ctx_handle(io->cmd);
 	/* store buffer index in request payload */
-	data->buf_index = pdu->buf.index;
+	io->buf_index = pdu->buf.index;
 	io->flags |= UBLK_IO_FLAG_AUTO_BUF_REG;
 	return true;
 }
 
 static bool ublk_prep_auto_buf_reg(struct ublk_queue *ubq,
 				   struct request *req, struct ublk_io *io,
 				   unsigned int issue_flags)
 {
 	if (ublk_support_auto_buf_reg(ubq) && ublk_rq_has_data(req))
-		return ublk_auto_buf_reg(req, io, issue_flags);
+		return ublk_auto_buf_reg(ubq, req, io, issue_flags);
 
-	ublk_init_req_ref(ubq, req);
+	ublk_init_req_ref(ubq, io);
 	return true;
 }
 
 static bool ublk_start_io(const struct ublk_queue *ubq, struct request *req,
 			  struct ublk_io *io)
@@ -1485,10 +1485,12 @@ static void ublk_queue_reinit(struct ublk_device *ub, struct ublk_queue *ubq)
 		 */
 		if (io->task) {
 			put_task_struct(io->task);
 			io->task = NULL;
 		}
+
+		WARN_ON_ONCE(refcount_read(&io->ref));
 	}
 }
 
 static int ublk_ch_open(struct inode *inode, struct file *filp)
 {
@@ -1989,33 +1991,35 @@ static inline int ublk_set_auto_buf_reg(struct io_uring_cmd *cmd)
 
 static void ublk_io_release(void *priv)
 {
 	struct request *rq = priv;
 	struct ublk_queue *ubq = rq->mq_hctx->driver_data;
+	struct ublk_io *io = &ubq->ios[rq->tag];
 
-	ublk_put_req_ref(ubq, rq);
+	ublk_put_req_ref(ubq, io, rq);
 }
 
 static int ublk_register_io_buf(struct io_uring_cmd *cmd,
-				const struct ublk_queue *ubq, unsigned int tag,
+				const struct ublk_queue *ubq,
+				struct ublk_io *io,
 				unsigned int index, unsigned int issue_flags)
 {
 	struct ublk_device *ub = cmd->file->private_data;
 	struct request *req;
 	int ret;
 
 	if (!ublk_support_zero_copy(ubq))
 		return -EINVAL;
 
-	req = __ublk_check_and_get_req(ub, ubq, tag, 0);
+	req = __ublk_check_and_get_req(ub, ubq, io, 0);
 	if (!req)
 		return -EINVAL;
 
 	ret = io_buffer_register_bvec(cmd, req, ublk_io_release, index,
 				      issue_flags);
 	if (ret) {
-		ublk_put_req_ref(ubq, req);
+		ublk_put_req_ref(ubq, io, req);
 		return ret;
 	}
 
 	return 0;
 }
@@ -2118,14 +2122,12 @@ static int ublk_commit_and_fetch(const struct ublk_queue *ubq,
 		 * one for registering the buffer, it is ublk server's
 		 * responsibility for unregistering the buffer, otherwise
 		 * this ublk request gets stuck.
 		 */
 		if (io->flags & UBLK_IO_FLAG_AUTO_BUF_REG) {
-			struct ublk_rq_data *data = blk_mq_rq_to_pdu(req);
-
-			if (data->buf_ctx_handle == io_uring_cmd_ctx_handle(cmd))
-				io_buffer_unregister_bvec(cmd, data->buf_index,
+			if (io->buf_ctx_handle == io_uring_cmd_ctx_handle(cmd))
+				io_buffer_unregister_bvec(cmd, io->buf_index,
 						issue_flags);
 			io->flags &= ~UBLK_IO_FLAG_AUTO_BUF_REG;
 		}
 
 		ret = ublk_set_auto_buf_reg(cmd);
@@ -2141,11 +2143,11 @@ static int ublk_commit_and_fetch(const struct ublk_queue *ubq,
 
 	if (req_op(req) == REQ_OP_ZONE_APPEND)
 		req->__sector = ub_cmd->zone_append_lba;
 
 	if (likely(!blk_should_fake_timeout(req->q)))
-		ublk_put_req_ref(ubq, req);
+		ublk_put_req_ref(ubq, io, req);
 
 	return 0;
 }
 
 static bool ublk_get_data(const struct ublk_queue *ubq, struct ublk_io *io)
@@ -2220,11 +2222,11 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 		goto out;
 
 	ret = -EINVAL;
 	switch (_IOC_NR(cmd_op)) {
 	case UBLK_IO_REGISTER_IO_BUF:
-		return ublk_register_io_buf(cmd, ubq, tag, ub_cmd->addr, issue_flags);
+		return ublk_register_io_buf(cmd, ubq, io, ub_cmd->addr, issue_flags);
 	case UBLK_IO_UNREGISTER_IO_BUF:
 		return ublk_unregister_io_buf(cmd, ubq, ub_cmd->addr, issue_flags);
 	case UBLK_IO_FETCH_REQ:
 		ret = ublk_fetch(cmd, ubq, io, ub_cmd->addr);
 		if (ret)
@@ -2252,19 +2254,24 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 			__func__, cmd_op, tag, ret, io->flags);
 	return ret;
 }
 
 static inline struct request *__ublk_check_and_get_req(struct ublk_device *ub,
-		const struct ublk_queue *ubq, int tag, size_t offset)
+		const struct ublk_queue *ubq, struct ublk_io *io, size_t offset)
 {
+	unsigned tag = io - ubq->ios;
 	struct request *req;
 
+	/*
+	 * can't use io->req in case of concurrent UBLK_IO_COMMIT_AND_FETCH_REQ,
+	 * which would overwrite it with io->cmd
+	 */
 	req = blk_mq_tag_to_rq(ub->tag_set.tags[ubq->q_id], tag);
 	if (!req)
 		return NULL;
 
-	if (!ublk_get_req_ref(ubq, req))
+	if (!ublk_get_req_ref(ubq, io))
 		return NULL;
 
 	if (unlikely(!blk_mq_request_started(req) || req->tag != tag))
 		goto fail_put;
 
@@ -2274,11 +2281,11 @@ static inline struct request *__ublk_check_and_get_req(struct ublk_device *ub,
 	if (offset > blk_rq_bytes(req))
 		goto fail_put;
 
 	return req;
 fail_put:
-	ublk_put_req_ref(ubq, req);
+	ublk_put_req_ref(ubq, io, req);
 	return NULL;
 }
 
 static inline int ublk_ch_uring_cmd_local(struct io_uring_cmd *cmd,
 		unsigned int issue_flags)
@@ -2341,11 +2348,12 @@ static inline bool ublk_check_ubuf_dir(const struct request *req,
 
 	return false;
 }
 
 static struct request *ublk_check_and_get_req(struct kiocb *iocb,
-		struct iov_iter *iter, size_t *off, int dir)
+		struct iov_iter *iter, size_t *off, int dir,
+		struct ublk_io **io)
 {
 	struct ublk_device *ub = iocb->ki_filp->private_data;
 	struct ublk_queue *ubq;
 	struct request *req;
 	size_t buf_off;
@@ -2375,11 +2383,12 @@ static struct request *ublk_check_and_get_req(struct kiocb *iocb,
 		return ERR_PTR(-EACCES);
 
 	if (tag >= ubq->q_depth)
 		return ERR_PTR(-EINVAL);
 
-	req = __ublk_check_and_get_req(ub, ubq, tag, buf_off);
+	*io = &ubq->ios[tag];
+	req = __ublk_check_and_get_req(ub, ubq, *io, buf_off);
 	if (!req)
 		return ERR_PTR(-EINVAL);
 
 	if (!req->mq_hctx || !req->mq_hctx->driver_data)
 		goto fail;
@@ -2388,46 +2397,48 @@ static struct request *ublk_check_and_get_req(struct kiocb *iocb,
 		goto fail;
 
 	*off = buf_off;
 	return req;
 fail:
-	ublk_put_req_ref(ubq, req);
+	ublk_put_req_ref(ubq, *io, req);
 	return ERR_PTR(-EACCES);
 }
 
 static ssize_t ublk_ch_read_iter(struct kiocb *iocb, struct iov_iter *to)
 {
 	struct ublk_queue *ubq;
 	struct request *req;
+	struct ublk_io *io;
 	size_t buf_off;
 	size_t ret;
 
-	req = ublk_check_and_get_req(iocb, to, &buf_off, ITER_DEST);
+	req = ublk_check_and_get_req(iocb, to, &buf_off, ITER_DEST, &io);
 	if (IS_ERR(req))
 		return PTR_ERR(req);
 
 	ret = ublk_copy_user_pages(req, buf_off, to, ITER_DEST);
 	ubq = req->mq_hctx->driver_data;
-	ublk_put_req_ref(ubq, req);
+	ublk_put_req_ref(ubq, io, req);
 
 	return ret;
 }
 
 static ssize_t ublk_ch_write_iter(struct kiocb *iocb, struct iov_iter *from)
 {
 	struct ublk_queue *ubq;
 	struct request *req;
+	struct ublk_io *io;
 	size_t buf_off;
 	size_t ret;
 
-	req = ublk_check_and_get_req(iocb, from, &buf_off, ITER_SOURCE);
+	req = ublk_check_and_get_req(iocb, from, &buf_off, ITER_SOURCE, &io);
 	if (IS_ERR(req))
 		return PTR_ERR(req);
 
 	ret = ublk_copy_user_pages(req, buf_off, from, ITER_SOURCE);
 	ubq = req->mq_hctx->driver_data;
-	ublk_put_req_ref(ubq, req);
+	ublk_put_req_ref(ubq, io, req);
 
 	return ret;
 }
 
 static const struct file_operations ublk_ch_fops = {
@@ -2448,10 +2459,11 @@ static void ublk_deinit_queue(struct ublk_device *ub, int q_id)
 
 	for (i = 0; i < ubq->q_depth; i++) {
 		struct ublk_io *io = &ubq->ios[i];
 		if (io->task)
 			put_task_struct(io->task);
+		WARN_ON_ONCE(refcount_read(&io->ref));
 	}
 
 	if (ubq->io_cmd_buf)
 		free_pages((unsigned long)ubq->io_cmd_buf, get_order(size));
 }
@@ -2600,11 +2612,10 @@ static int ublk_add_tag_set(struct ublk_device *ub)
 {
 	ub->tag_set.ops = &ublk_mq_ops;
 	ub->tag_set.nr_hw_queues = ub->dev_info.nr_hw_queues;
 	ub->tag_set.queue_depth = ub->dev_info.queue_depth;
 	ub->tag_set.numa_node = NUMA_NO_NODE;
-	ub->tag_set.cmd_size = sizeof(struct ublk_rq_data);
 	ub->tag_set.driver_data = ub;
 	return blk_mq_alloc_tag_set(&ub->tag_set);
 }
 
 static void ublk_remove(struct ublk_device *ub)
-- 
2.45.2

[PATCH v2 03/14] ublk: check cmd_op first

[Caleb Sander Mateos]

In preparation for skipping some of the other checks for certain IO
opcodes, move the cmd_op check earlier.

Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
Reviewed-by: Ming Lei <ming.lei@redhat.com>
---
 drivers/block/ublk_drv.c | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index 55d5435f1ee2..aea2e9200311 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -2177,16 +2177,21 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 	struct task_struct *task;
 	struct ublk_queue *ubq;
 	struct ublk_io *io;
 	u32 cmd_op = cmd->cmd_op;
 	unsigned tag = ub_cmd->tag;
-	int ret = -EINVAL;
+	int ret;
 
 	pr_devel("%s: received: cmd op %d queue %d tag %d result %d\n",
 			__func__, cmd->cmd_op, ub_cmd->q_id, tag,
 			ub_cmd->result);
 
+	ret = ublk_check_cmd_op(cmd_op);
+	if (ret)
+		goto out;
+
+	ret = -EINVAL;
 	if (ub_cmd->q_id >= ub->dev_info.nr_hw_queues)
 		goto out;
 
 	ubq = ublk_get_queue(ub, ub_cmd->q_id);
 
@@ -2215,15 +2220,10 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 	 */
 	if ((!!(io->flags & UBLK_IO_FLAG_NEED_GET_DATA))
 			^ (_IOC_NR(cmd_op) == UBLK_IO_NEED_GET_DATA))
 		goto out;
 
-	ret = ublk_check_cmd_op(cmd_op);
-	if (ret)
-		goto out;
-
-	ret = -EINVAL;
 	switch (_IOC_NR(cmd_op)) {
 	case UBLK_IO_REGISTER_IO_BUF:
 		return ublk_register_io_buf(cmd, ubq, io, ub_cmd->addr, issue_flags);
 	case UBLK_IO_UNREGISTER_IO_BUF:
 		return ublk_unregister_io_buf(cmd, ubq, ub_cmd->addr, issue_flags);
-- 
2.45.2

[PATCH v2 04/14] ublk: handle UBLK_IO_FETCH_REQ earlier

[Caleb Sander Mateos]

Check for UBLK_IO_FETCH_REQ early in __ublk_ch_uring_cmd() and skip the
rest of the checks in this case. This allows removing the checks for
NULL io->task and UBLK_IO_FLAG_OWNED_BY_SRV unset in io->flags, which
are only allowed for FETCH.

Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
Reviewed-by: Ming Lei <ming.lei@redhat.com>
---
 drivers/block/ublk_drv.c | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index aea2e9200311..8df70a5fb129 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -2197,23 +2197,31 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 
 	if (tag >= ubq->q_depth)
 		goto out;
 
 	io = &ubq->ios[tag];
+	/* UBLK_IO_FETCH_REQ can be handled on any task, which sets io->task */
+	if (unlikely(_IOC_NR(cmd_op) == UBLK_IO_FETCH_REQ)) {
+		ret = ublk_fetch(cmd, ubq, io, ub_cmd->addr);
+		if (ret)
+			goto out;
+
+		ublk_prep_cancel(cmd, issue_flags, ubq, tag);
+		return -EIOCBQUEUED;
+	}
+
 	task = READ_ONCE(io->task);
-	if (task && task != current)
+	if (task != current)
 		goto out;
 
 	/* there is pending io cmd, something must be wrong */
 	if (io->flags & UBLK_IO_FLAG_ACTIVE) {
 		ret = -EBUSY;
 		goto out;
 	}
 
-	/* only UBLK_IO_FETCH_REQ is allowed if io is not OWNED_BY_SRV */
-	if (!(io->flags & UBLK_IO_FLAG_OWNED_BY_SRV) &&
-	    _IOC_NR(cmd_op) != UBLK_IO_FETCH_REQ)
+	if (!(io->flags & UBLK_IO_FLAG_OWNED_BY_SRV))
 		goto out;
 
 	/*
 	 * ensure that the user issues UBLK_IO_NEED_GET_DATA
 	 * iff the driver have set the UBLK_IO_FLAG_NEED_GET_DATA.
@@ -2225,15 +2233,10 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 	switch (_IOC_NR(cmd_op)) {
 	case UBLK_IO_REGISTER_IO_BUF:
 		return ublk_register_io_buf(cmd, ubq, io, ub_cmd->addr, issue_flags);
 	case UBLK_IO_UNREGISTER_IO_BUF:
 		return ublk_unregister_io_buf(cmd, ubq, ub_cmd->addr, issue_flags);
-	case UBLK_IO_FETCH_REQ:
-		ret = ublk_fetch(cmd, ubq, io, ub_cmd->addr);
-		if (ret)
-			goto out;
-		break;
 	case UBLK_IO_COMMIT_AND_FETCH_REQ:
 		ret = ublk_commit_and_fetch(ubq, io, cmd, ub_cmd, issue_flags);
 		if (ret)
 			goto out;
 		break;
-- 
2.45.2

[PATCH v2 05/14] ublk: remove task variable from __ublk_ch_uring_cmd()

[Caleb Sander Mateos]

The variable is computed from a simple expression and used once, so just
replace it with the expression.

Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
Reviewed-by: Ming Lei <ming.lei@redhat.com>
---
 drivers/block/ublk_drv.c | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index 8df70a5fb129..7635105c1596 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -2172,11 +2172,10 @@ static bool ublk_get_data(const struct ublk_queue *ubq, struct ublk_io *io)
 static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 			       unsigned int issue_flags,
 			       const struct ublksrv_io_cmd *ub_cmd)
 {
 	struct ublk_device *ub = cmd->file->private_data;
-	struct task_struct *task;
 	struct ublk_queue *ubq;
 	struct ublk_io *io;
 	u32 cmd_op = cmd->cmd_op;
 	unsigned tag = ub_cmd->tag;
 	int ret;
@@ -2207,12 +2206,11 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 
 		ublk_prep_cancel(cmd, issue_flags, ubq, tag);
 		return -EIOCBQUEUED;
 	}
 
-	task = READ_ONCE(io->task);
-	if (task != current)
+	if (READ_ONCE(io->task) != current)
 		goto out;
 
 	/* there is pending io cmd, something must be wrong */
 	if (io->flags & UBLK_IO_FLAG_ACTIVE) {
 		ret = -EBUSY;
-- 
2.45.2

[PATCH v2 06/14] ublk: consolidate UBLK_IO_FLAG_{ACTIVE,OWNED_BY_SRV} checks

[Caleb Sander Mateos]

UBLK_IO_FLAG_ACTIVE and UBLK_IO_FLAG_OWNED_BY_SRV are mutually
exclusive. So just check that UBLK_IO_FLAG_OWNED_BY_SRV is set in
__ublk_ch_uring_cmd(); that implies UBLK_IO_FLAG_ACTIVE is unset.

Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
Reviewed-by: Ming Lei <ming.lei@redhat.com>
---
 drivers/block/ublk_drv.c | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index 7635105c1596..9bfccee3c2b7 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -2210,18 +2210,15 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 
 	if (READ_ONCE(io->task) != current)
 		goto out;
 
 	/* there is pending io cmd, something must be wrong */
-	if (io->flags & UBLK_IO_FLAG_ACTIVE) {
+	if (!(io->flags & UBLK_IO_FLAG_OWNED_BY_SRV)) {
 		ret = -EBUSY;
 		goto out;
 	}
 
-	if (!(io->flags & UBLK_IO_FLAG_OWNED_BY_SRV))
-		goto out;
-
 	/*
 	 * ensure that the user issues UBLK_IO_NEED_GET_DATA
 	 * iff the driver have set the UBLK_IO_FLAG_NEED_GET_DATA.
 	 */
 	if ((!!(io->flags & UBLK_IO_FLAG_NEED_GET_DATA))
-- 
2.45.2

[PATCH v2 07/14] ublk: move ublk_prep_cancel() to case UBLK_IO_COMMIT_AND_FETCH_REQ

[Caleb Sander Mateos]

UBLK_IO_COMMIT_AND_FETCH_REQ is the only one of __ublk_ch_uring_cmd()'s
switch cases that doesn't return or goto. Move the logic following the
switch into this case so it also returns. Drop the now unneeded default
case.

Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
Reviewed-by: Ming Lei <ming.lei@redhat.com>
---
 drivers/block/ublk_drv.c | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index 9bfccee3c2b7..893519f74625 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -2232,22 +2232,20 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 		return ublk_unregister_io_buf(cmd, ubq, ub_cmd->addr, issue_flags);
 	case UBLK_IO_COMMIT_AND_FETCH_REQ:
 		ret = ublk_commit_and_fetch(ubq, io, cmd, ub_cmd, issue_flags);
 		if (ret)
 			goto out;
-		break;
+
+		ublk_prep_cancel(cmd, issue_flags, ubq, tag);
+		return -EIOCBQUEUED;
 	case UBLK_IO_NEED_GET_DATA:
 		io->addr = ub_cmd->addr;
 		if (!ublk_get_data(ubq, io))
 			return -EIOCBQUEUED;
 
 		return UBLK_IO_RES_OK;
-	default:
-		goto out;
 	}
-	ublk_prep_cancel(cmd, issue_flags, ubq, tag);
-	return -EIOCBQUEUED;
 
  out:
 	pr_devel("%s: complete: cmd op %d, tag %d ret %x io_flags %x\n",
 			__func__, cmd_op, tag, ret, io->flags);
 	return ret;
-- 
2.45.2

[PATCH v2 08/14] ublk: don't take ublk_queue in ublk_unregister_io_buf()

[Caleb Sander Mateos]

UBLK_IO_UNREGISTER_IO_BUF currently requires a valid q_id and tag to be
passed in the ublksrv_io_cmd. However, only the addr (registered buffer
index) is actually used to unregister the buffer. There is no check that
the q_id and tag are for the ublk request whose buffer is registered at
the given index. To prepare to allow userspace to omit the q_id and tag,
check the UBLK_F_SUPPORT_ZERO_COPY flag on the ublk_device instead of
the ublk_queue.

Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
---
 drivers/block/ublk_drv.c | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index 893519f74625..237d50490923 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -2023,14 +2023,14 @@ static int ublk_register_io_buf(struct io_uring_cmd *cmd,
 
 	return 0;
 }
 
 static int ublk_unregister_io_buf(struct io_uring_cmd *cmd,
-				  const struct ublk_queue *ubq,
+				  const struct ublk_device *ub,
 				  unsigned int index, unsigned int issue_flags)
 {
-	if (!ublk_support_zero_copy(ubq))
+	if (!(ub->dev_info.flags & UBLK_F_SUPPORT_ZERO_COPY))
 		return -EINVAL;
 
 	return io_buffer_unregister_bvec(cmd, index, issue_flags);
 }
 
@@ -2227,11 +2227,11 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 
 	switch (_IOC_NR(cmd_op)) {
 	case UBLK_IO_REGISTER_IO_BUF:
 		return ublk_register_io_buf(cmd, ubq, io, ub_cmd->addr, issue_flags);
 	case UBLK_IO_UNREGISTER_IO_BUF:
-		return ublk_unregister_io_buf(cmd, ubq, ub_cmd->addr, issue_flags);
+		return ublk_unregister_io_buf(cmd, ub, ub_cmd->addr, issue_flags);
 	case UBLK_IO_COMMIT_AND_FETCH_REQ:
 		ret = ublk_commit_and_fetch(ubq, io, cmd, ub_cmd, issue_flags);
 		if (ret)
 			goto out;
 
-- 
2.45.2

[PATCH v2 09/14] ublk: allow UBLK_IO_(UN)REGISTER_IO_BUF on any task

[Caleb Sander Mateos]

Currently, UBLK_IO_REGISTER_IO_BUF and UBLK_IO_UNREGISTER_IO_BUF are
only permitted on the ublk_io's daemon task. But this restriction is
unnecessary. ublk_register_io_buf() calls __ublk_check_and_get_req() to
look up the request from the tagset and atomically take a reference on
the request without accessing the ublk_io. ublk_unregister_io_buf()
doesn't use the q_id or tag at all.

So allow these opcodes even on tasks other than io->task.

Handle UBLK_IO_UNREGISTER_IO_BUF before obtaining the ubq and io since
the buffer index being unregistered is not necessarily related to the
specified q_id and tag.

Add a feature flag UBLK_F_BUF_REG_OFF_DAEMON that userspace can use to
determine whether the kernel supports off-daemon buffer registration.

Suggested-by: Ming Lei <ming.lei@redhat.com>
Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
---
 drivers/block/ublk_drv.c      | 27 ++++++++++++++++++++++-----
 include/uapi/linux/ublk_cmd.h | 10 ++++++++++
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index 237d50490923..11aa65f36ec9 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -68,11 +68,12 @@
 		| UBLK_F_ZONED \
 		| UBLK_F_USER_RECOVERY_FAIL_IO \
 		| UBLK_F_UPDATE_SIZE \
 		| UBLK_F_AUTO_BUF_REG \
 		| UBLK_F_QUIESCE \
-		| UBLK_F_PER_IO_DAEMON)
+		| UBLK_F_PER_IO_DAEMON \
+		| UBLK_F_BUF_REG_OFF_DAEMON)
 
 #define UBLK_F_ALL_RECOVERY_FLAGS (UBLK_F_USER_RECOVERY \
 		| UBLK_F_USER_RECOVERY_REISSUE \
 		| UBLK_F_USER_RECOVERY_FAIL_IO)
 
@@ -2186,10 +2187,18 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 
 	ret = ublk_check_cmd_op(cmd_op);
 	if (ret)
 		goto out;
 
+	/*
+	 * io_buffer_unregister_bvec() doesn't access the ubq or io,
+	 * so no need to validate the q_id, tag, or task
+	 */
+	if (_IOC_NR(cmd_op) == UBLK_IO_UNREGISTER_IO_BUF)
+		return ublk_unregister_io_buf(cmd, ub, ub_cmd->addr,
+					      issue_flags);
+
 	ret = -EINVAL;
 	if (ub_cmd->q_id >= ub->dev_info.nr_hw_queues)
 		goto out;
 
 	ubq = ublk_get_queue(ub, ub_cmd->q_id);
@@ -2206,12 +2215,21 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 
 		ublk_prep_cancel(cmd, issue_flags, ubq, tag);
 		return -EIOCBQUEUED;
 	}
 
-	if (READ_ONCE(io->task) != current)
+	if (READ_ONCE(io->task) != current) {
+		/*
+		 * ublk_register_io_buf() accesses only the io's refcount,
+		 * so can be handled on any task
+		 */
+		if (_IOC_NR(cmd_op) == UBLK_IO_REGISTER_IO_BUF)
+			return ublk_register_io_buf(cmd, ubq, io, ub_cmd->addr,
+						    issue_flags);
+
 		goto out;
+	}
 
 	/* there is pending io cmd, something must be wrong */
 	if (!(io->flags & UBLK_IO_FLAG_OWNED_BY_SRV)) {
 		ret = -EBUSY;
 		goto out;
@@ -2226,12 +2244,10 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 		goto out;
 
 	switch (_IOC_NR(cmd_op)) {
 	case UBLK_IO_REGISTER_IO_BUF:
 		return ublk_register_io_buf(cmd, ubq, io, ub_cmd->addr, issue_flags);
-	case UBLK_IO_UNREGISTER_IO_BUF:
-		return ublk_unregister_io_buf(cmd, ub, ub_cmd->addr, issue_flags);
 	case UBLK_IO_COMMIT_AND_FETCH_REQ:
 		ret = ublk_commit_and_fetch(ubq, io, cmd, ub_cmd, issue_flags);
 		if (ret)
 			goto out;
 
@@ -2932,11 +2948,12 @@ static int ublk_ctrl_add_dev(const struct ublksrv_ctrl_cmd *header)
 	 */
 	ub->dev_info.flags &= UBLK_F_ALL;
 
 	ub->dev_info.flags |= UBLK_F_CMD_IOCTL_ENCODE |
 		UBLK_F_URING_CMD_COMP_IN_TASK |
-		UBLK_F_PER_IO_DAEMON;
+		UBLK_F_PER_IO_DAEMON |
+		UBLK_F_BUF_REG_OFF_DAEMON;
 
 	/* GET_DATA isn't needed any more with USER_COPY or ZERO COPY */
 	if (ub->dev_info.flags & (UBLK_F_USER_COPY | UBLK_F_SUPPORT_ZERO_COPY |
 				UBLK_F_AUTO_BUF_REG))
 		ub->dev_info.flags &= ~UBLK_F_NEED_GET_DATA;
diff --git a/include/uapi/linux/ublk_cmd.h b/include/uapi/linux/ublk_cmd.h
index 77d9d6af46da..a7b52f6c1ba4 100644
--- a/include/uapi/linux/ublk_cmd.h
+++ b/include/uapi/linux/ublk_cmd.h
@@ -279,10 +279,20 @@
  * (qid1,tag1) and (qid2,tag2), if qid1 == qid2, then the same task must
  * be responsible for handling (qid1,tag1) and (qid2,tag2).
  */
 #define UBLK_F_PER_IO_DAEMON (1ULL << 13)
 
+/*
+ * If this feature is set, UBLK_U_IO_REGISTER_IO_BUF/UBLK_U_IO_UNREGISTER_IO_BUF
+ * can be issued for an I/O on any task. q_id and tag are also ignored in
+ * UBLK_U_IO_UNREGISTER_IO_BUF's ublksrv_io_cmd.
+ * If it is unset, zero-copy buffers can only be registered and unregistered by
+ * the I/O's daemon task. The q_id and tag of the registered buffer are required
+ * in UBLK_U_IO_UNREGISTER_IO_BUF's ublksrv_io_cmd.
+ */
+#define UBLK_F_BUF_REG_OFF_DAEMON (1ULL << 14)
+
 /* device state */
 #define UBLK_S_DEV_DEAD	0
 #define UBLK_S_DEV_LIVE	1
 #define UBLK_S_DEV_QUIESCED	2
 #define UBLK_S_DEV_FAIL_IO 	3
-- 
2.45.2

[PATCH v2 10/14] ublk: return early if blk_should_fake_timeout()

[Caleb Sander Mateos]

Make the unlikely case blk_should_fake_timeout() return early to reduce
the indentation of the successful path.

Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
---
 drivers/block/ublk_drv.c | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index 11aa65f36ec9..f53618391141 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -2143,13 +2143,14 @@ static int ublk_commit_and_fetch(const struct ublk_queue *ubq,
 	io->res = ub_cmd->result;
 
 	if (req_op(req) == REQ_OP_ZONE_APPEND)
 		req->__sector = ub_cmd->zone_append_lba;
 
-	if (likely(!blk_should_fake_timeout(req->q)))
-		ublk_put_req_ref(ubq, io, req);
+	if (unlikely(blk_should_fake_timeout(req->q)))
+		return 0;
 
+	ublk_put_req_ref(ubq, io, req);
 	return 0;
 }
 
 static bool ublk_get_data(const struct ublk_queue *ubq, struct ublk_io *io)
 {
-- 
2.45.2

[PATCH v2 11/14] ublk: optimize UBLK_IO_REGISTER_IO_BUF on daemon task

[Caleb Sander Mateos]

ublk_register_io_buf() performs an expensive atomic refcount increment,
as well as a lot of pointer chasing to look up the struct request.

Create a separate ublk_daemon_register_io_buf() for the daemon task to
call. Initialize ublk_io's reference count to a large number, introduce
a field task_registered_buffers to count the buffers registered on the
daemon task, and atomically subtract the large number minus
task_registered_buffers in ublk_commit_and_fetch().

Also obtain the struct request directly from ublk_io's req field instead
of looking it up on the tagset.

Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
---
 drivers/block/ublk_drv.c | 70 ++++++++++++++++++++++++++++++++++------
 1 file changed, 61 insertions(+), 9 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index f53618391141..b2925e15279a 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -146,10 +146,17 @@ struct ublk_uring_cmd_pdu {
 #define UBLK_IO_FLAG_AUTO_BUF_REG 	0x10
 
 /* atomic RW with ubq->cancel_lock */
 #define UBLK_IO_FLAG_CANCELED	0x80000000
 
+/*
+ * Initialize refcount to a large number to include any registered buffers.
+ * UBLK_IO_COMMIT_AND_FETCH_REQ will release these references minus those for
+ * any buffers registered on the io daemon task.
+ */
+#define UBLK_REFCOUNT_INIT (REFCOUNT_MAX / 2)
+
 struct ublk_io {
 	/* userspace buffer address from io cmd */
 	__u64	addr;
 	unsigned int flags;
 	int res;
@@ -164,18 +171,21 @@ struct ublk_io {
 	struct task_struct *task;
 
 	/*
 	 * The number of uses of this I/O by the ublk server
 	 * if user copy or zero copy are enabled:
-	 * - 1 from dispatch to the server until UBLK_IO_COMMIT_AND_FETCH_REQ
+	 * - UBLK_REFCOUNT_INIT from dispatch to the server
+	 *   until UBLK_IO_COMMIT_AND_FETCH_REQ
 	 * - 1 for each inflight ublk_ch_{read,write}_iter() call
-	 * - 1 for each io_uring registered buffer
+	 * - 1 for each io_uring registered buffer not registered on task
 	 * The I/O can only be completed once all references are dropped.
 	 * User copy and buffer registration operations are only permitted
 	 * if the reference count is nonzero.
 	 */
 	refcount_t ref;
+	/* Count of buffers registered on task and not yet unregistered */
+	unsigned task_registered_buffers;
 
 	/* auto-registered buffer, valid if UBLK_IO_FLAG_AUTO_BUF_REG is set */
 	u16 buf_index;
 	void *buf_ctx_handle;
 };
@@ -684,11 +694,11 @@ static inline bool ublk_need_req_ref(const struct ublk_queue *ubq)
 
 static inline void ublk_init_req_ref(const struct ublk_queue *ubq,
 		struct ublk_io *io)
 {
 	if (ublk_need_req_ref(ubq))
-		refcount_set(&io->ref, 1);
+		refcount_set(&io->ref, UBLK_REFCOUNT_INIT);
 }
 
 static inline bool ublk_get_req_ref(const struct ublk_queue *ubq,
 		struct ublk_io *io)
 {
@@ -707,10 +717,19 @@ static inline void ublk_put_req_ref(const struct ublk_queue *ubq,
 	} else {
 		__ublk_complete_rq(req);
 	}
 }
 
+static inline void ublk_sub_req_ref(struct ublk_io *io, struct request *req)
+{
+	unsigned sub_refs = UBLK_REFCOUNT_INIT - io->task_registered_buffers;
+
+	io->task_registered_buffers = 0;
+	if (refcount_sub_and_test(sub_refs, &io->ref))
+		__ublk_complete_rq(req);
+}
+
 static inline bool ublk_need_get_data(const struct ublk_queue *ubq)
 {
 	return ubq->flags & UBLK_F_NEED_GET_DATA;
 }
 
@@ -1188,11 +1207,10 @@ ublk_auto_buf_reg_fallback(const struct ublk_queue *ubq, struct ublk_io *io)
 {
 	unsigned tag = io - ubq->ios;
 	struct ublksrv_io_desc *iod = ublk_get_iod(ubq, tag);
 
 	iod->op_flags |= UBLK_IO_F_NEED_REG_BUF;
-	refcount_set(&io->ref, 1);
 }
 
 static bool ublk_auto_buf_reg(const struct ublk_queue *ubq, struct request *req,
 			      struct ublk_io *io, unsigned int issue_flags)
 {
@@ -1207,13 +1225,12 @@ static bool ublk_auto_buf_reg(const struct ublk_queue *ubq, struct request *req,
 			return true;
 		}
 		blk_mq_end_request(req, BLK_STS_IOERR);
 		return false;
 	}
-	/* one extra reference is dropped by ublk_io_release */
-	refcount_set(&io->ref, 2);
 
+	io->task_registered_buffers = 1;
 	io->buf_ctx_handle = io_uring_cmd_ctx_handle(io->cmd);
 	/* store buffer index in request payload */
 	io->buf_index = pdu->buf.index;
 	io->flags |= UBLK_IO_FLAG_AUTO_BUF_REG;
 	return true;
@@ -1221,14 +1238,14 @@ static bool ublk_auto_buf_reg(const struct ublk_queue *ubq, struct request *req,
 
 static bool ublk_prep_auto_buf_reg(struct ublk_queue *ubq,
 				   struct request *req, struct ublk_io *io,
 				   unsigned int issue_flags)
 {
+	ublk_init_req_ref(ubq, io);
 	if (ublk_support_auto_buf_reg(ubq) && ublk_rq_has_data(req))
 		return ublk_auto_buf_reg(ubq, req, io, issue_flags);
 
-	ublk_init_req_ref(ubq, io);
 	return true;
 }
 
 static bool ublk_start_io(const struct ublk_queue *ubq, struct request *req,
 			  struct ublk_io *io)
@@ -1488,10 +1505,11 @@ static void ublk_queue_reinit(struct ublk_device *ub, struct ublk_queue *ubq)
 			put_task_struct(io->task);
 			io->task = NULL;
 		}
 
 		WARN_ON_ONCE(refcount_read(&io->ref));
+		WARN_ON_ONCE(io->task_registered_buffers);
 	}
 }
 
 static int ublk_ch_open(struct inode *inode, struct file *filp)
 {
@@ -2023,10 +2041,39 @@ static int ublk_register_io_buf(struct io_uring_cmd *cmd,
 	}
 
 	return 0;
 }
 
+static int
+ublk_daemon_register_io_buf(struct io_uring_cmd *cmd,
+			    const struct ublk_queue *ubq, struct ublk_io *io,
+			    unsigned index, unsigned issue_flags)
+{
+	unsigned new_registered_buffers;
+	struct request *req = io->req;
+	int ret;
+
+	/*
+	 * Ensure there are still references for ublk_sub_req_ref() to release.
+	 * If not, fall back on the thread-safe buffer registration.
+	 */
+	new_registered_buffers = io->task_registered_buffers + 1;
+	if (unlikely(new_registered_buffers >= UBLK_REFCOUNT_INIT))
+		return ublk_register_io_buf(cmd, ubq, io, index, issue_flags);
+
+	if (!ublk_support_zero_copy(ubq) || !ublk_rq_has_data(req))
+		return -EINVAL;
+
+	ret = io_buffer_register_bvec(cmd, req, ublk_io_release, index,
+				      issue_flags);
+	if (ret)
+		return ret;
+
+	io->task_registered_buffers = new_registered_buffers;
+	return 0;
+}
+
 static int ublk_unregister_io_buf(struct io_uring_cmd *cmd,
 				  const struct ublk_device *ub,
 				  unsigned int index, unsigned int issue_flags)
 {
 	if (!(ub->dev_info.flags & UBLK_F_SUPPORT_ZERO_COPY))
@@ -2146,11 +2193,14 @@ static int ublk_commit_and_fetch(const struct ublk_queue *ubq,
 		req->__sector = ub_cmd->zone_append_lba;
 
 	if (unlikely(blk_should_fake_timeout(req->q)))
 		return 0;
 
-	ublk_put_req_ref(ubq, io, req);
+	if (ublk_need_req_ref(ubq))
+		ublk_sub_req_ref(io, req);
+	else
+		__ublk_complete_rq(req);
 	return 0;
 }
 
 static bool ublk_get_data(const struct ublk_queue *ubq, struct ublk_io *io)
 {
@@ -2244,11 +2294,12 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 			^ (_IOC_NR(cmd_op) == UBLK_IO_NEED_GET_DATA))
 		goto out;
 
 	switch (_IOC_NR(cmd_op)) {
 	case UBLK_IO_REGISTER_IO_BUF:
-		return ublk_register_io_buf(cmd, ubq, io, ub_cmd->addr, issue_flags);
+		return ublk_daemon_register_io_buf(cmd, ubq, io, ub_cmd->addr,
+						   issue_flags);
 	case UBLK_IO_COMMIT_AND_FETCH_REQ:
 		ret = ublk_commit_and_fetch(ubq, io, cmd, ub_cmd, issue_flags);
 		if (ret)
 			goto out;
 
@@ -2473,10 +2524,11 @@ static void ublk_deinit_queue(struct ublk_device *ub, int q_id)
 	for (i = 0; i < ubq->q_depth; i++) {
 		struct ublk_io *io = &ubq->ios[i];
 		if (io->task)
 			put_task_struct(io->task);
 		WARN_ON_ONCE(refcount_read(&io->ref));
+		WARN_ON_ONCE(io->task_registered_buffers);
 	}
 
 	if (ubq->io_cmd_buf)
 		free_pages((unsigned long)ubq->io_cmd_buf, get_order(size));
 }
-- 
2.45.2

[PATCH v2 12/14] ublk: optimize UBLK_IO_UNREGISTER_IO_BUF on daemon task

[Caleb Sander Mateos]

ublk_io_release() performs an expensive atomic refcount decrement. This
atomic operation is unnecessary in the common case where the request's
buffer is registered and unregistered on the daemon task before handling
UBLK_IO_COMMIT_AND_FETCH_REQ for the I/O. So if ublk_io_release() is
called on the daemon task and task_registered_buffers is positive, just
decrement task_registered_buffers (nonatomically). ublk_sub_req_ref()
will apply this decrement when it atomically subtracts from io->ref.

Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
---
 drivers/block/ublk_drv.c | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index b2925e15279a..199028f36ec8 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -2012,11 +2012,18 @@ static void ublk_io_release(void *priv)
 {
 	struct request *rq = priv;
 	struct ublk_queue *ubq = rq->mq_hctx->driver_data;
 	struct ublk_io *io = &ubq->ios[rq->tag];
 
-	ublk_put_req_ref(ubq, io, rq);
+	/*
+	 * task_registered_buffers may be 0 if buffers were registered off task
+	 * but unregistered on task. Or after UBLK_IO_COMMIT_AND_FETCH_REQ.
+	 */
+	if (current == io->task && io->task_registered_buffers)
+		io->task_registered_buffers--;
+	else
+		ublk_put_req_ref(ubq, io, rq);
 }
 
 static int ublk_register_io_buf(struct io_uring_cmd *cmd,
 				const struct ublk_queue *ubq,
 				struct ublk_io *io,
-- 
2.45.2

[PATCH v2 13/14] ublk: remove ubq checks from ublk_{get,put}_req_ref()

[Caleb Sander Mateos]

ublk_get_req_ref() and ublk_put_req_ref() currently call
ublk_need_req_ref(ubq) to check whether the ublk device features require
reference counting of its requests. However, all callers already know
that reference counting is required:
- __ublk_check_and_get_req() is only called from
  ublk_check_and_get_req() if user copy is enabled, and from
  ublk_register_io_buf() if zero copy is enabled
- ublk_io_release() is only called for requests registered by
  ublk_register_io_buf(), which requires zero copy
- ublk_ch_read_iter() and ublk_ch_write_iter() only call
  ublk_put_req_ref() if ublk_check_and_get_req() succeeded, which
  requires user copy to be enabled

So drop the ublk_need_req_ref() check and the ubq argument in
ublk_get_req_ref() and ublk_put_req_ref().

Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
---
 drivers/block/ublk_drv.c | 35 +++++++++++------------------------
 1 file changed, 11 insertions(+), 24 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index 199028f36ec8..ebc56681eb68 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -697,28 +697,19 @@ static inline void ublk_init_req_ref(const struct ublk_queue *ubq,
 {
 	if (ublk_need_req_ref(ubq))
 		refcount_set(&io->ref, UBLK_REFCOUNT_INIT);
 }
 
-static inline bool ublk_get_req_ref(const struct ublk_queue *ubq,
-		struct ublk_io *io)
+static inline bool ublk_get_req_ref(struct ublk_io *io)
 {
-	if (ublk_need_req_ref(ubq))
-		return refcount_inc_not_zero(&io->ref);
-
-	return true;
+	return refcount_inc_not_zero(&io->ref);
 }
 
-static inline void ublk_put_req_ref(const struct ublk_queue *ubq,
-		struct ublk_io *io, struct request *req)
+static inline void ublk_put_req_ref(struct ublk_io *io, struct request *req)
 {
-	if (ublk_need_req_ref(ubq)) {
-		if (refcount_dec_and_test(&io->ref))
-			__ublk_complete_rq(req);
-	} else {
+	if (refcount_dec_and_test(&io->ref))
 		__ublk_complete_rq(req);
-	}
 }
 
 static inline void ublk_sub_req_ref(struct ublk_io *io, struct request *req)
 {
 	unsigned sub_refs = UBLK_REFCOUNT_INIT - io->task_registered_buffers;
@@ -2019,11 +2010,11 @@ static void ublk_io_release(void *priv)
 	 * but unregistered on task. Or after UBLK_IO_COMMIT_AND_FETCH_REQ.
 	 */
 	if (current == io->task && io->task_registered_buffers)
 		io->task_registered_buffers--;
 	else
-		ublk_put_req_ref(ubq, io, rq);
+		ublk_put_req_ref(io, rq);
 }
 
 static int ublk_register_io_buf(struct io_uring_cmd *cmd,
 				const struct ublk_queue *ubq,
 				struct ublk_io *io,
@@ -2041,11 +2032,11 @@ static int ublk_register_io_buf(struct io_uring_cmd *cmd,
 		return -EINVAL;
 
 	ret = io_buffer_register_bvec(cmd, req, ublk_io_release, index,
 				      issue_flags);
 	if (ret) {
-		ublk_put_req_ref(ubq, io, req);
+		ublk_put_req_ref(io, req);
 		return ret;
 	}
 
 	return 0;
 }
@@ -2338,11 +2329,11 @@ static inline struct request *__ublk_check_and_get_req(struct ublk_device *ub,
 	 */
 	req = blk_mq_tag_to_rq(ub->tag_set.tags[ubq->q_id], tag);
 	if (!req)
 		return NULL;
 
-	if (!ublk_get_req_ref(ubq, io))
+	if (!ublk_get_req_ref(io))
 		return NULL;
 
 	if (unlikely(!blk_mq_request_started(req) || req->tag != tag))
 		goto fail_put;
 
@@ -2352,11 +2343,11 @@ static inline struct request *__ublk_check_and_get_req(struct ublk_device *ub,
 	if (offset > blk_rq_bytes(req))
 		goto fail_put;
 
 	return req;
 fail_put:
-	ublk_put_req_ref(ubq, io, req);
+	ublk_put_req_ref(io, req);
 	return NULL;
 }
 
 static inline int ublk_ch_uring_cmd_local(struct io_uring_cmd *cmd,
 		unsigned int issue_flags)
@@ -2468,48 +2459,44 @@ static struct request *ublk_check_and_get_req(struct kiocb *iocb,
 		goto fail;
 
 	*off = buf_off;
 	return req;
 fail:
-	ublk_put_req_ref(ubq, *io, req);
+	ublk_put_req_ref(*io, req);
 	return ERR_PTR(-EACCES);
 }
 
 static ssize_t ublk_ch_read_iter(struct kiocb *iocb, struct iov_iter *to)
 {
-	struct ublk_queue *ubq;
 	struct request *req;
 	struct ublk_io *io;
 	size_t buf_off;
 	size_t ret;
 
 	req = ublk_check_and_get_req(iocb, to, &buf_off, ITER_DEST, &io);
 	if (IS_ERR(req))
 		return PTR_ERR(req);
 
 	ret = ublk_copy_user_pages(req, buf_off, to, ITER_DEST);
-	ubq = req->mq_hctx->driver_data;
-	ublk_put_req_ref(ubq, io, req);
+	ublk_put_req_ref(io, req);
 
 	return ret;
 }
 
 static ssize_t ublk_ch_write_iter(struct kiocb *iocb, struct iov_iter *from)
 {
-	struct ublk_queue *ubq;
 	struct request *req;
 	struct ublk_io *io;
 	size_t buf_off;
 	size_t ret;
 
 	req = ublk_check_and_get_req(iocb, from, &buf_off, ITER_SOURCE, &io);
 	if (IS_ERR(req))
 		return PTR_ERR(req);
 
 	ret = ublk_copy_user_pages(req, buf_off, from, ITER_SOURCE);
-	ubq = req->mq_hctx->driver_data;
-	ublk_put_req_ref(ubq, io, req);
+	ublk_put_req_ref(io, req);
 
 	return ret;
 }
 
 static const struct file_operations ublk_ch_fops = {
-- 
2.45.2

[PATCH v2 14/14] ublk: cache-align struct ublk_io

[Caleb Sander Mateos]

struct ublk_io is already 56 bytes on 64-bit architectures, so round it
up to a full cache line (typically 64 bytes). This ensures a single
ublk_io doesn't span multiple cache lines and prevents false sharing if
consecutive ublk_io's are accessed by different daemon tasks.

Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
---
 drivers/block/ublk_drv.c | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index ebc56681eb68..740141c63a93 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -186,11 +186,11 @@ struct ublk_io {
 	unsigned task_registered_buffers;
 
 	/* auto-registered buffer, valid if UBLK_IO_FLAG_AUTO_BUF_REG is set */
 	u16 buf_index;
 	void *buf_ctx_handle;
-};
+} ____cacheline_aligned_in_smp;
 
 struct ublk_queue {
 	int q_id;
 	int q_depth;
 
-- 
2.45.2

Re: [PATCH v2 01/14] ublk: use vmalloc for ublk_device's __queues

[Ming Lei]

On Fri, Jun 20, 2025 at 09:09:55AM -0600, Caleb Sander Mateos wrote:
> struct ublk_device's __queues points to an allocation with up to
> UBLK_MAX_NR_QUEUES (4096) queues, each of which have:
> - struct ublk_queue (48 bytes)
> - Tail array of up to UBLK_MAX_QUEUE_DEPTH (4096) struct ublk_io's,
>   32 bytes each
> This means the full allocation can exceed 512 MB, which may well be
> impossible to service with contiguous physical pages. Switch to
> kvcalloc() and kvfree(), since there is no need for physically
> contiguous memory.
> 
> Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
> Fixes: 71f28f3136af ("ublk_drv: add io_uring based userspace block driver")

Reviewed-by: Ming Lei <ming.lei@redhat.com>

Thanks,
Ming

Re: [PATCH v2 02/14] ublk: remove struct ublk_rq_data

[Ming Lei]

On Fri, Jun 20, 2025 at 09:09:56AM -0600, Caleb Sander Mateos wrote:
> __ublk_check_and_get_req() attempts to atomically look up the struct
> request for a ublk I/O and take a reference on it. However, the request
> can be freed between the lookup on the tagset in blk_mq_tag_to_rq() and
> the increment of its reference count in ublk_get_req_ref(), for example
> if an elevator switch happens concurrently.
> 
> Fix the potential use after free by moving the reference count from
> ublk_rq_data to ublk_io. Move the fields buf_index and buf_ctx_handle
> too to reduce the number of cache lines touched when dispatching and
> completing a ublk I/O, allowing ublk_rq_data to be removed entirely.
> 
> Suggested-by: Ming Lei <ming.lei@redhat.com>
> Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
> Fixes: 62fe99cef94a ("ublk: add read()/write() support for ublk char device")

Reviewed-by: Ming Lei <ming.lei@redhat.com>

Thanks,
Ming

Re: [PATCH v2 08/14] ublk: don't take ublk_queue in ublk_unregister_io_buf()

[Ming Lei]

On Fri, Jun 20, 2025 at 09:10:02AM -0600, Caleb Sander Mateos wrote:
> UBLK_IO_UNREGISTER_IO_BUF currently requires a valid q_id and tag to be
> passed in the ublksrv_io_cmd. However, only the addr (registered buffer
> index) is actually used to unregister the buffer. There is no check that
> the q_id and tag are for the ublk request whose buffer is registered at
> the given index. To prepare to allow userspace to omit the q_id and tag,
> check the UBLK_F_SUPPORT_ZERO_COPY flag on the ublk_device instead of
> the ublk_queue.
> 
> Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>

Reviewed-by: Ming Lei <ming.lei@redhat.com>

Thanks,
Ming

Re: [PATCH v2 09/14] ublk: allow UBLK_IO_(UN)REGISTER_IO_BUF on any task

[Ming Lei]

On Fri, Jun 20, 2025 at 09:10:03AM -0600, Caleb Sander Mateos wrote:
> Currently, UBLK_IO_REGISTER_IO_BUF and UBLK_IO_UNREGISTER_IO_BUF are
> only permitted on the ublk_io's daemon task. But this restriction is
> unnecessary. ublk_register_io_buf() calls __ublk_check_and_get_req() to
> look up the request from the tagset and atomically take a reference on
> the request without accessing the ublk_io. ublk_unregister_io_buf()
> doesn't use the q_id or tag at all.
> 
> So allow these opcodes even on tasks other than io->task.
> 
> Handle UBLK_IO_UNREGISTER_IO_BUF before obtaining the ubq and io since
> the buffer index being unregistered is not necessarily related to the
> specified q_id and tag.
> 
> Add a feature flag UBLK_F_BUF_REG_OFF_DAEMON that userspace can use to
> determine whether the kernel supports off-daemon buffer registration.
> 
> Suggested-by: Ming Lei <ming.lei@redhat.com>
> Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>

Reviewed-by: Ming Lei <ming.lei@redhat.com>

Thanks,
Ming

Re: [PATCH v2 10/14] ublk: return early if blk_should_fake_timeout()

[Ming Lei]

On Fri, Jun 20, 2025 at 09:10:04AM -0600, Caleb Sander Mateos wrote:
> Make the unlikely case blk_should_fake_timeout() return early to reduce
> the indentation of the successful path.
> 
> Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
> ---
>  drivers/block/ublk_drv.c | 5 +++--
>  1 file changed, 3 insertions(+), 2 deletions(-)
> 
> diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
> index 11aa65f36ec9..f53618391141 100644
> --- a/drivers/block/ublk_drv.c
> +++ b/drivers/block/ublk_drv.c
> @@ -2143,13 +2143,14 @@ static int ublk_commit_and_fetch(const struct ublk_queue *ubq,
>  	io->res = ub_cmd->result;
>  
>  	if (req_op(req) == REQ_OP_ZONE_APPEND)
>  		req->__sector = ub_cmd->zone_append_lba;
>  
> -	if (likely(!blk_should_fake_timeout(req->q)))
> -		ublk_put_req_ref(ubq, io, req);
> +	if (unlikely(blk_should_fake_timeout(req->q)))
> +		return 0;
>  
> +	ublk_put_req_ref(ubq, io, req);
>  	return 0;
>  }

Reviewed-by: Ming Lei <ming.lei@redhat.com>

Thanks,
Ming

Re: [PATCH v2 11/14] ublk: optimize UBLK_IO_REGISTER_IO_BUF on daemon task

[Ming Lei]

On Fri, Jun 20, 2025 at 09:10:05AM -0600, Caleb Sander Mateos wrote:
> ublk_register_io_buf() performs an expensive atomic refcount increment,
> as well as a lot of pointer chasing to look up the struct request.
> 
> Create a separate ublk_daemon_register_io_buf() for the daemon task to
> call. Initialize ublk_io's reference count to a large number, introduce
> a field task_registered_buffers to count the buffers registered on the
> daemon task, and atomically subtract the large number minus
> task_registered_buffers in ublk_commit_and_fetch().
> 
> Also obtain the struct request directly from ublk_io's req field instead
> of looking it up on the tagset.
> 
> Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>

Reviewed-by: Ming Lei <ming.lei@redhat.com>

Thanks,
Ming

Re: [PATCH v2 12/14] ublk: optimize UBLK_IO_UNREGISTER_IO_BUF on daemon task

[Ming Lei]

On Fri, Jun 20, 2025 at 09:10:06AM -0600, Caleb Sander Mateos wrote:
> ublk_io_release() performs an expensive atomic refcount decrement. This
> atomic operation is unnecessary in the common case where the request's
> buffer is registered and unregistered on the daemon task before handling
> UBLK_IO_COMMIT_AND_FETCH_REQ for the I/O. So if ublk_io_release() is
> called on the daemon task and task_registered_buffers is positive, just
> decrement task_registered_buffers (nonatomically). ublk_sub_req_ref()
> will apply this decrement when it atomically subtracts from io->ref.
> 
> Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
> ---
>  drivers/block/ublk_drv.c | 9 ++++++++-
>  1 file changed, 8 insertions(+), 1 deletion(-)
> 
> diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
> index b2925e15279a..199028f36ec8 100644
> --- a/drivers/block/ublk_drv.c
> +++ b/drivers/block/ublk_drv.c
> @@ -2012,11 +2012,18 @@ static void ublk_io_release(void *priv)
>  {
>  	struct request *rq = priv;
>  	struct ublk_queue *ubq = rq->mq_hctx->driver_data;
>  	struct ublk_io *io = &ubq->ios[rq->tag];
>  
> -	ublk_put_req_ref(ubq, io, rq);
> +	/*
> +	 * task_registered_buffers may be 0 if buffers were registered off task
> +	 * but unregistered on task. Or after UBLK_IO_COMMIT_AND_FETCH_REQ.
> +	 */
> +	if (current == io->task && io->task_registered_buffers)
> +		io->task_registered_buffers--;
> +	else
> +		ublk_put_req_ref(ubq, io, rq);
>  }

Reviewed-by: Ming Lei <ming.lei@redhat.com>

Thanks,
Ming

Re: [PATCH v2 13/14] ublk: remove ubq checks from ublk_{get,put}_req_ref()

[Ming Lei]

On Fri, Jun 20, 2025 at 09:10:07AM -0600, Caleb Sander Mateos wrote:
> ublk_get_req_ref() and ublk_put_req_ref() currently call
> ublk_need_req_ref(ubq) to check whether the ublk device features require
> reference counting of its requests. However, all callers already know
> that reference counting is required:
> - __ublk_check_and_get_req() is only called from
>   ublk_check_and_get_req() if user copy is enabled, and from
>   ublk_register_io_buf() if zero copy is enabled
> - ublk_io_release() is only called for requests registered by
>   ublk_register_io_buf(), which requires zero copy
> - ublk_ch_read_iter() and ublk_ch_write_iter() only call
>   ublk_put_req_ref() if ublk_check_and_get_req() succeeded, which
>   requires user copy to be enabled
> 
> So drop the ublk_need_req_ref() check and the ubq argument in
> ublk_get_req_ref() and ublk_put_req_ref().
> 
> Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>

Reviewed-by: Ming Lei <ming.lei@redhat.com>

Thanks,
Ming

Re: [PATCH v2 14/14] ublk: cache-align struct ublk_io

[Ming Lei]

On Fri, Jun 20, 2025 at 09:10:08AM -0600, Caleb Sander Mateos wrote:
> struct ublk_io is already 56 bytes on 64-bit architectures, so round it
> up to a full cache line (typically 64 bytes). This ensures a single
> ublk_io doesn't span multiple cache lines and prevents false sharing if
> consecutive ublk_io's are accessed by different daemon tasks.
> 
> Signed-off-by: Caleb Sander Mateos <csander@purestorage.com>
> ---
>  drivers/block/ublk_drv.c | 2 +-
>  1 file changed, 1 insertion(+), 1 deletion(-)
> 
> diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
> index ebc56681eb68..740141c63a93 100644
> --- a/drivers/block/ublk_drv.c
> +++ b/drivers/block/ublk_drv.c
> @@ -186,11 +186,11 @@ struct ublk_io {
>  	unsigned task_registered_buffers;
>  
>  	/* auto-registered buffer, valid if UBLK_IO_FLAG_AUTO_BUF_REG is set */
>  	u16 buf_index;
>  	void *buf_ctx_handle;
> -};
> +} ____cacheline_aligned_in_smp;

Reviewed-by: Ming Lei <ming.lei@redhat.com>

Thanks, 
Ming

Re: [PATCH v2 00/14] ublk: allow off-daemon zero-copy buffer registration

[Jens Axboe]

On Fri, 20 Jun 2025 09:09:54 -0600, Caleb Sander Mateos wrote:
> Currently ublk zero-copy requires ublk request buffers to be registered
> and unregistered by the ublk I/O's daemon task. However, as currently
> implemented, there is no reason for this restriction. Registration looks
> up the request via the ublk device's tagset rather than the daemon-local
> ublk_io structure and takes an atomic reference to prevent racing with
> dispatch or completion of the request. Ming has expressed interest in
> relaxing this restriction[1] so the ublk server can offload the I/O
> operation that uses the zero-copy buffer to another thread.
> 
> [...]

Applied, thanks!

[01/14] ublk: use vmalloc for ublk_device's __queues
        commit: d2bc702a8db21c9819f25737a4343f6dfc238582
[02/14] ublk: remove struct ublk_rq_data
        commit: f053751a413b232cecd5cb3d3d9c1f972a849310
[03/14] ublk: check cmd_op first
        commit: 6762d9c0b76e2a449a818539f352bd14a5799512
[04/14] ublk: handle UBLK_IO_FETCH_REQ earlier
        commit: 1a3402e31fca3dbeb915c34a077ca3d143b42e55
[05/14] ublk: remove task variable from __ublk_ch_uring_cmd()
        commit: e9415d3aee01b51c8633d38ca62f74f7071a74ba
[06/14] ublk: consolidate UBLK_IO_FLAG_{ACTIVE,OWNED_BY_SRV} checks
        commit: 82293b83e42366d15ded8516a4d45f8d466e69ba
[07/14] ublk: move ublk_prep_cancel() to case UBLK_IO_COMMIT_AND_FETCH_REQ
        commit: 624006c145af20924624d4f2846da74f94aa9078
[08/14] ublk: don't take ublk_queue in ublk_unregister_io_buf()
        commit: 3e6b973df269180c22e024606585f8b4ca33ca99
[09/14] ublk: allow UBLK_IO_(UN)REGISTER_IO_BUF on any task
        commit: 73166f433b3c2a6acebb850f5eefea36c137f6c3
[10/14] ublk: return early if blk_should_fake_timeout()
        commit: 358f19d7cadbb5a7626cfd5822260aadead7f37e
[11/14] ublk: optimize UBLK_IO_REGISTER_IO_BUF on daemon task
        commit: 47c894fdaaee8b424cdbcd0fe490a10ea89ab0ee
[12/14] ublk: optimize UBLK_IO_UNREGISTER_IO_BUF on daemon task
        commit: 230c5a4cfc6f98d2432aff05bb3ec092fde519f5
[13/14] ublk: remove ubq checks from ublk_{get,put}_req_ref()
        commit: 21d2d986b5511ac017660720ed0016ce2780396c
[14/14] ublk: cache-align struct ublk_io
        commit: 456ef6804f232f3b2f60147046e05500147b0099

Best regards,
-- 
Jens Axboe

Re: [PATCH v2 00/14] ublk: allow off-daemon zero-copy buffer registration

[Jens Axboe]

On 6/26/25 6:47 PM, Jens Axboe wrote:
> 
> On Fri, 20 Jun 2025 09:09:54 -0600, Caleb Sander Mateos wrote:
>> Currently ublk zero-copy requires ublk request buffers to be registered
>> and unregistered by the ublk I/O's daemon task. However, as currently
>> implemented, there is no reason for this restriction. Registration looks
>> up the request via the ublk device's tagset rather than the daemon-local
>> ublk_io structure and takes an atomic reference to prevent racing with
>> dispatch or completion of the request. Ming has expressed interest in
>> relaxing this restriction[1] so the ublk server can offload the I/O
>> operation that uses the zero-copy buffer to another thread.
>>
>> [...]
> 
> Applied, thanks!
> 
> [01/14] ublk: use vmalloc for ublk_device's __queues
>         commit: d2bc702a8db21c9819f25737a4343f6dfc238582
> [02/14] ublk: remove struct ublk_rq_data
>         commit: f053751a413b232cecd5cb3d3d9c1f972a849310
> [03/14] ublk: check cmd_op first
>         commit: 6762d9c0b76e2a449a818539f352bd14a5799512
> [04/14] ublk: handle UBLK_IO_FETCH_REQ earlier
>         commit: 1a3402e31fca3dbeb915c34a077ca3d143b42e55
> [05/14] ublk: remove task variable from __ublk_ch_uring_cmd()
>         commit: e9415d3aee01b51c8633d38ca62f74f7071a74ba
> [06/14] ublk: consolidate UBLK_IO_FLAG_{ACTIVE,OWNED_BY_SRV} checks
>         commit: 82293b83e42366d15ded8516a4d45f8d466e69ba
> [07/14] ublk: move ublk_prep_cancel() to case UBLK_IO_COMMIT_AND_FETCH_REQ
>         commit: 624006c145af20924624d4f2846da74f94aa9078
> [08/14] ublk: don't take ublk_queue in ublk_unregister_io_buf()
>         commit: 3e6b973df269180c22e024606585f8b4ca33ca99
> [09/14] ublk: allow UBLK_IO_(UN)REGISTER_IO_BUF on any task
>         commit: 73166f433b3c2a6acebb850f5eefea36c137f6c3
> [10/14] ublk: return early if blk_should_fake_timeout()
>         commit: 358f19d7cadbb5a7626cfd5822260aadead7f37e
> [11/14] ublk: optimize UBLK_IO_REGISTER_IO_BUF on daemon task
>         commit: 47c894fdaaee8b424cdbcd0fe490a10ea89ab0ee
> [12/14] ublk: optimize UBLK_IO_UNREGISTER_IO_BUF on daemon task
>         commit: 230c5a4cfc6f98d2432aff05bb3ec092fde519f5
> [13/14] ublk: remove ubq checks from ublk_{get,put}_req_ref()
>         commit: 21d2d986b5511ac017660720ed0016ce2780396c
> [14/14] ublk: cache-align struct ublk_io
>         commit: 456ef6804f232f3b2f60147046e05500147b0099

Caleb, please check the final result in my for-next branch, as
expected this conflicted with changes in master/block-6.16.

-- 
Jens Axboe