[PATCH v4 0/4] ublk: decouple server threads from hctxs

[PATCH v4 0/4] ublk: decouple server threads from hctxs

[Uday Shankar]

This patch set aims to allow ublk server threads to better balance load
amongst themselves by decoupling server threads from ublk queues/hctxs,
so that multiple threads can service I/Os from a single hctx.

Signed-off-by: Uday Shankar <ushankar@purestorage.com>
---
Changes in v4:
- Drop "ublk: properly serialize all FETCH_REQs" since Ming is taking it
  in another set
- Prevent data races by marking data structures which should be
  read-only in the I/O path as const (Ming Lei)
- Link to v3: https://lore.kernel.org/r/20250410-ublk_task_per_io-v3-0-b811e8f4554a@purestorage.com

Changes in v3:
- Check for UBLK_IO_FLAG_ACTIVE on I/O again after taking lock to ensure
  that two concurrent FETCH_REQs on the same I/O can't succeed (Caleb
  Sander Mateos)
- Link to v2: https://lore.kernel.org/r/20250408-ublk_task_per_io-v2-0-b97877e6fd50@purestorage.com

Changes in v2:
- Remove changes split into other patches
- To ease error handling/synchronization, associate each I/O (instead of
  each queue) to the last task that issues a FETCH_REQ against it. Only
  that task is allowed to operate on the I/O.
- Link to v1: https://lore.kernel.org/r/20241002224437.3088981-1-ushankar@purestorage.com

---
Uday Shankar (4):
      ublk: require unique task per io instead of unique task per hctx
      ublk: mark ublk_queue as const for ublk_commit_and_fetch
      ublk: mark ublk_queue as const for ublk_register_io_buf
      ublk: mark ublk_queue as const for ublk_handle_need_get_data

 drivers/block/ublk_drv.c | 205 +++++++++++++++++++++++------------------------
 1 file changed, 100 insertions(+), 105 deletions(-)
---
base-commit: d3b4b25e363e4ce193e6103e64f7de12b96668b9
change-id: 20250408-ublk_task_per_io-c693cf608d7a

Best regards,
-- 
Uday Shankar <ushankar@purestorage.com>

[PATCH v4 1/4] ublk: require unique task per io instead of unique task per hctx

[Uday Shankar]

Currently, ublk_drv associates to each hardware queue (hctx) a unique
task (called the queue's ubq_daemon) which is allowed to issue
COMMIT_AND_FETCH commands against the hctx. If any other task attempts
to do so, the command fails immediately with EINVAL. When considered
together with the block layer architecture, the result is that for each
CPU C on the system, there is a unique ublk server thread which is
allowed to handle I/O submitted on CPU C. This can lead to suboptimal
performance under imbalanced load generation. For an extreme example,
suppose all the load is generated on CPUs mapping to a single ublk
server thread. Then that thread may be fully utilized and become the
bottleneck in the system, while other ublk server threads are totally
idle.

This issue can also be addressed directly in the ublk server without
kernel support by having threads dequeue I/Os and pass them around to
ensure even load. But this solution requires inter-thread communication
at least twice for each I/O (submission and completion), which is
generally a bad pattern for performance. The problem gets even worse
with zero copy, as more inter-thread communication would be required to
have the buffer register/unregister calls to come from the correct
thread.

Therefore, address this issue in ublk_drv by requiring a unique task per
I/O instead of per queue/hctx. Imbalanced load can then be balanced
across all ublk server threads by having threads issue FETCH_REQs in a
round-robin manner. As a small toy example, consider a system with a
single ublk device having 2 queues, each of queue depth 4. A ublk server
having 4 threads could issue its FETCH_REQs against this device as
follows (where each entry is the qid,tag pair that the FETCH_REQ
targets):

poller thread:	T0	T1	T2	T3
		0,0	0,1	0,2	0,3
		1,3	1,0	1,1	1,2

Since tags appear to be allocated in sequential chunks, this setup
provides a rough approximation to distributing I/Os round-robin across
all ublk server threads, while letting I/Os stay fully thread-local.

Signed-off-by: Uday Shankar <ushankar@purestorage.com>
Reviewed-by: Caleb Sander Mateos <csander@purestorage.com>
---
 drivers/block/ublk_drv.c | 75 ++++++++++++++++++++++--------------------------
 1 file changed, 34 insertions(+), 41 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index cdb1543fa4a9817aa2ca2fca66720f589cf222be..9a0d2547512fc8119460739230599d48d2c2a306 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -150,6 +150,7 @@ struct ublk_io {
 	int res;
 
 	struct io_uring_cmd *cmd;
+	struct task_struct *task;
 };
 
 struct ublk_queue {
@@ -157,11 +158,9 @@ struct ublk_queue {
 	int q_depth;
 
 	unsigned long flags;
-	struct task_struct	*ubq_daemon;
 	struct ublksrv_io_desc *io_cmd_buf;
 
 	bool force_abort;
-	bool timeout;
 	bool canceling;
 	bool fail_io; /* copy of dev->state == UBLK_S_DEV_FAIL_IO */
 	unsigned short nr_io_ready;	/* how many ios setup */
@@ -1072,11 +1071,6 @@ static inline struct ublk_uring_cmd_pdu *ublk_get_uring_cmd_pdu(
 	return io_uring_cmd_to_pdu(ioucmd, struct ublk_uring_cmd_pdu);
 }
 
-static inline bool ubq_daemon_is_dying(struct ublk_queue *ubq)
-{
-	return ubq->ubq_daemon->flags & PF_EXITING;
-}
-
 /* todo: handle partial completion */
 static inline void __ublk_complete_rq(struct request *req)
 {
@@ -1224,13 +1218,13 @@ static void ublk_dispatch_req(struct ublk_queue *ubq,
 	/*
 	 * Task is exiting if either:
 	 *
-	 * (1) current != ubq_daemon.
+	 * (1) current != io->task.
 	 * io_uring_cmd_complete_in_task() tries to run task_work
-	 * in a workqueue if ubq_daemon(cmd's task) is PF_EXITING.
+	 * in a workqueue if cmd's task is PF_EXITING.
 	 *
 	 * (2) current->flags & PF_EXITING.
 	 */
-	if (unlikely(current != ubq->ubq_daemon || current->flags & PF_EXITING)) {
+	if (unlikely(current != io->task || current->flags & PF_EXITING)) {
 		__ublk_abort_rq(ubq, req);
 		return;
 	}
@@ -1336,23 +1330,20 @@ static void ublk_queue_cmd_list(struct ublk_queue *ubq, struct rq_list *l)
 static enum blk_eh_timer_return ublk_timeout(struct request *rq)
 {
 	struct ublk_queue *ubq = rq->mq_hctx->driver_data;
+	struct ublk_io *io = &ubq->ios[rq->tag];
 	unsigned int nr_inflight = 0;
 	int i;
 
 	if (ubq->flags & UBLK_F_UNPRIVILEGED_DEV) {
-		if (!ubq->timeout) {
-			send_sig(SIGKILL, ubq->ubq_daemon, 0);
-			ubq->timeout = true;
-		}
-
+		send_sig(SIGKILL, io->task, 0);
 		return BLK_EH_DONE;
 	}
 
-	if (!ubq_daemon_is_dying(ubq))
+	if (!(io->task->flags & PF_EXITING))
 		return BLK_EH_RESET_TIMER;
 
 	for (i = 0; i < ubq->q_depth; i++) {
-		struct ublk_io *io = &ubq->ios[i];
+		io = &ubq->ios[i];
 
 		if (!(io->flags & UBLK_IO_FLAG_ACTIVE))
 			nr_inflight++;
@@ -1552,8 +1543,8 @@ static void ublk_commit_completion(struct ublk_device *ub,
 }
 
 /*
- * Called from ubq_daemon context via cancel fn, meantime quiesce ublk
- * blk-mq queue, so we are called exclusively with blk-mq and ubq_daemon
+ * Called from io task context via cancel fn, meantime quiesce ublk
+ * blk-mq queue, so we are called exclusively with blk-mq and io task
  * context, so everything is serialized.
  */
 static void ublk_abort_queue(struct ublk_device *ub, struct ublk_queue *ubq)
@@ -1669,13 +1660,13 @@ static void ublk_uring_cmd_cancel_fn(struct io_uring_cmd *cmd,
 		return;
 
 	task = io_uring_cmd_get_task(cmd);
-	if (WARN_ON_ONCE(task && task != ubq->ubq_daemon))
+	io = &ubq->ios[pdu->tag];
+	if (WARN_ON_ONCE(task && task != io->task))
 		return;
 
 	ub = ubq->dev;
 	need_schedule = ublk_abort_requests(ub, ubq);
 
-	io = &ubq->ios[pdu->tag];
 	WARN_ON_ONCE(io->cmd != cmd);
 	ublk_cancel_cmd(ubq, io, issue_flags);
 
@@ -1836,8 +1827,6 @@ static void ublk_mark_io_ready(struct ublk_device *ub, struct ublk_queue *ubq)
 	mutex_lock(&ub->mutex);
 	ubq->nr_io_ready++;
 	if (ublk_queue_ready(ubq)) {
-		ubq->ubq_daemon = current;
-		get_task_struct(ubq->ubq_daemon);
 		ub->nr_queues_ready++;
 
 		if (capable(CAP_SYS_ADMIN))
@@ -1952,14 +1941,14 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 	if (!ubq || ub_cmd->q_id != ubq->q_id)
 		goto out;
 
-	if (ubq->ubq_daemon && ubq->ubq_daemon != current)
-		goto out;
-
 	if (tag >= ubq->q_depth)
 		goto out;
 
 	io = &ubq->ios[tag];
 
+	if (io->task && io->task != current)
+		goto out;
+
 	/* there is pending io cmd, something must be wrong */
 	if (io->flags & UBLK_IO_FLAG_ACTIVE) {
 		ret = -EBUSY;
@@ -2012,6 +2001,7 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 
 		ublk_fill_io_cmd(io, cmd, ub_cmd->addr);
 		ublk_mark_io_ready(ub, ubq);
+		io->task = get_task_struct(current);
 		break;
 	case UBLK_IO_COMMIT_AND_FETCH_REQ:
 		req = blk_mq_tag_to_rq(ub->tag_set.tags[ub_cmd->q_id], tag);
@@ -2248,9 +2238,15 @@ static void ublk_deinit_queue(struct ublk_device *ub, int q_id)
 {
 	int size = ublk_queue_cmd_buf_size(ub, q_id);
 	struct ublk_queue *ubq = ublk_get_queue(ub, q_id);
+	struct ublk_io *io;
+	int i;
+
+	for (i = 0; i < ubq->q_depth; i++) {
+		io = &ubq->ios[i];
+		if (io->task)
+			put_task_struct(io->task);
+	}
 
-	if (ubq->ubq_daemon)
-		put_task_struct(ubq->ubq_daemon);
 	if (ubq->io_cmd_buf)
 		free_pages((unsigned long)ubq->io_cmd_buf, get_order(size));
 }
@@ -2936,15 +2932,8 @@ static void ublk_queue_reinit(struct ublk_device *ub, struct ublk_queue *ubq)
 {
 	int i;
 
-	WARN_ON_ONCE(!(ubq->ubq_daemon && ubq_daemon_is_dying(ubq)));
-
 	/* All old ioucmds have to be completed */
 	ubq->nr_io_ready = 0;
-	/* old daemon is PF_EXITING, put it now */
-	put_task_struct(ubq->ubq_daemon);
-	/* We have to reset it to NULL, otherwise ub won't accept new FETCH_REQ */
-	ubq->ubq_daemon = NULL;
-	ubq->timeout = false;
 	ubq->canceling = false;
 
 	for (i = 0; i < ubq->q_depth; i++) {
@@ -2954,6 +2943,10 @@ static void ublk_queue_reinit(struct ublk_device *ub, struct ublk_queue *ubq)
 		io->flags = 0;
 		io->cmd = NULL;
 		io->addr = 0;
+
+		WARN_ON_ONCE(!(io->task && (io->task->flags & PF_EXITING)));
+		put_task_struct(io->task);
+		io->task = NULL;
 	}
 }
 
@@ -2993,7 +2986,7 @@ static int ublk_ctrl_start_recovery(struct ublk_device *ub,
 	pr_devel("%s: start recovery for dev id %d.\n", __func__, header->dev_id);
 	for (i = 0; i < ub->dev_info.nr_hw_queues; i++)
 		ublk_queue_reinit(ub, ublk_get_queue(ub, i));
-	/* set to NULL, otherwise new ubq_daemon cannot mmap the io_cmd_buf */
+	/* set to NULL, otherwise new tasks cannot mmap the io_cmd_buf */
 	ub->mm = NULL;
 	ub->nr_queues_ready = 0;
 	ub->nr_privileged_daemon = 0;
@@ -3011,14 +3004,14 @@ static int ublk_ctrl_end_recovery(struct ublk_device *ub,
 	int ret = -EINVAL;
 	int i;
 
-	pr_devel("%s: Waiting for new ubq_daemons(nr: %d) are ready, dev id %d...\n",
-			__func__, ub->dev_info.nr_hw_queues, header->dev_id);
-	/* wait until new ubq_daemon sending all FETCH_REQ */
+	pr_devel("%s: Waiting for all FETCH_REQs, dev id %d...\n", __func__,
+		 header->dev_id);
+
 	if (wait_for_completion_interruptible(&ub->completion))
 		return -EINTR;
 
-	pr_devel("%s: All new ubq_daemons(nr: %d) are ready, dev id %d\n",
-			__func__, ub->dev_info.nr_hw_queues, header->dev_id);
+	pr_devel("%s: All FETCH_REQs received, dev id %d\n", __func__,
+		 header->dev_id);
 
 	mutex_lock(&ub->mutex);
 	if (ublk_nosrv_should_stop_dev(ub))

-- 
2.34.1

[PATCH v4 2/4] ublk: mark ublk_queue as const for ublk_commit_and_fetch

[Uday Shankar]

We now allow multiple tasks to operate on I/Os belonging to the same
queue concurrently. This means that any writes to ublk_queue in the I/O
path are potential sources of data races. Try to prevent these by
marking ublk_queue pointers as const when handling COMMIT_AND_FETCH.
Move the logic for this command into its own function
ublk_commit_and_fetch. Also open code ublk_commit_completion in
ublk_commit_and_fetch to reduce the number of parameters/avoid a
redundant lookup.

Suggested-by: Ming Lei <ming.lei@redhat.com>
Signed-off-by: Uday Shankar <ushankar@purestorage.com>
---
 drivers/block/ublk_drv.c | 91 +++++++++++++++++++++++-------------------------
 1 file changed, 43 insertions(+), 48 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index 9a0d2547512fc8119460739230599d48d2c2a306..153f67d92248ad45bddd2437b1306bb23df7d1ae 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -1518,30 +1518,6 @@ static int ublk_ch_mmap(struct file *filp, struct vm_area_struct *vma)
 	return remap_pfn_range(vma, vma->vm_start, pfn, sz, vma->vm_page_prot);
 }
 
-static void ublk_commit_completion(struct ublk_device *ub,
-		const struct ublksrv_io_cmd *ub_cmd)
-{
-	u32 qid = ub_cmd->q_id, tag = ub_cmd->tag;
-	struct ublk_queue *ubq = ublk_get_queue(ub, qid);
-	struct ublk_io *io = &ubq->ios[tag];
-	struct request *req;
-
-	/* now this cmd slot is owned by nbd driver */
-	io->flags &= ~UBLK_IO_FLAG_OWNED_BY_SRV;
-	io->res = ub_cmd->result;
-
-	/* find the io request and complete */
-	req = blk_mq_tag_to_rq(ub->tag_set.tags[qid], tag);
-	if (WARN_ON_ONCE(unlikely(!req)))
-		return;
-
-	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);
-}
-
 /*
  * Called from io task context via cancel fn, meantime quiesce ublk
  * blk-mq queue, so we are called exclusively with blk-mq and io task
@@ -1918,6 +1894,45 @@ static int ublk_unregister_io_buf(struct io_uring_cmd *cmd,
 	return io_buffer_unregister_bvec(cmd, index, issue_flags);
 }
 
+static int ublk_commit_and_fetch(const struct ublk_queue *ubq,
+				 struct ublk_io *io, struct io_uring_cmd *cmd,
+				 const struct ublksrv_io_cmd *ub_cmd,
+				 struct request *req)
+{
+	if (!(io->flags & UBLK_IO_FLAG_OWNED_BY_SRV))
+		return -EINVAL;
+
+	if (ublk_need_map_io(ubq)) {
+		/*
+		 * COMMIT_AND_FETCH_REQ has to provide IO buffer if
+		 * NEED GET DATA is not enabled or it is Read IO.
+		 */
+		if (!ub_cmd->addr && (!ublk_need_get_data(ubq) ||
+					req_op(req) == REQ_OP_READ))
+			return -EINVAL;
+	} else if (req_op(req) != REQ_OP_ZONE_APPEND && ub_cmd->addr) {
+		/*
+		 * User copy requires addr to be unset when command is
+		 * not zone append
+		 */
+		return -EINVAL;
+	}
+
+	ublk_fill_io_cmd(io, cmd, ub_cmd->addr);
+
+	/* now this cmd slot is owned by ublk driver */
+	io->flags &= ~UBLK_IO_FLAG_OWNED_BY_SRV;
+	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, req);
+
+	return -EIOCBQUEUED;
+}
+
 static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 			       unsigned int issue_flags,
 			       const struct ublksrv_io_cmd *ub_cmd)
@@ -1928,7 +1943,6 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 	u32 cmd_op = cmd->cmd_op;
 	unsigned tag = ub_cmd->tag;
 	int ret = -EINVAL;
-	struct request *req;
 
 	pr_devel("%s: received: cmd op %d queue %d tag %d result %d\n",
 			__func__, cmd->cmd_op, ub_cmd->q_id, tag,
@@ -2004,30 +2018,11 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 		io->task = get_task_struct(current);
 		break;
 	case UBLK_IO_COMMIT_AND_FETCH_REQ:
-		req = blk_mq_tag_to_rq(ub->tag_set.tags[ub_cmd->q_id], tag);
-
-		if (!(io->flags & UBLK_IO_FLAG_OWNED_BY_SRV))
+		ret = ublk_commit_and_fetch(
+			ubq, io, cmd, ub_cmd,
+			blk_mq_tag_to_rq(ub->tag_set.tags[ub_cmd->q_id], tag));
+		if (ret != -EIOCBQUEUED)
 			goto out;
-
-		if (ublk_need_map_io(ubq)) {
-			/*
-			 * COMMIT_AND_FETCH_REQ has to provide IO buffer if
-			 * NEED GET DATA is not enabled or it is Read IO.
-			 */
-			if (!ub_cmd->addr && (!ublk_need_get_data(ubq) ||
-						req_op(req) == REQ_OP_READ))
-				goto out;
-		} else if (req_op(req) != REQ_OP_ZONE_APPEND && ub_cmd->addr) {
-			/*
-			 * User copy requires addr to be unset when command is
-			 * not zone append
-			 */
-			ret = -EINVAL;
-			goto out;
-		}
-
-		ublk_fill_io_cmd(io, cmd, ub_cmd->addr);
-		ublk_commit_completion(ub, ub_cmd);
 		break;
 	case UBLK_IO_NEED_GET_DATA:
 		if (!(io->flags & UBLK_IO_FLAG_OWNED_BY_SRV))

-- 
2.34.1

[PATCH v4 3/4] ublk: mark ublk_queue as const for ublk_register_io_buf

[Uday Shankar]

We now allow multiple tasks to operate on I/Os belonging to the same
queue concurrently. This means that any writes to ublk_queue in the I/O
path are potential sources of data races. Try to prevent these by
marking ublk_queue pointers as const in ublk_register_io_buf.

Suggested-by: Ming Lei <ming.lei@redhat.com>
Signed-off-by: Uday Shankar <ushankar@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 153f67d92248ad45bddd2437b1306bb23df7d1ae..e2cb54895481aebaa91ab23ba05cf26a950a642f 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -211,7 +211,7 @@ struct ublk_params_header {
 static bool ublk_abort_requests(struct ublk_device *ub, struct ublk_queue *ubq);
 
 static inline struct request *__ublk_check_and_get_req(struct ublk_device *ub,
-		struct ublk_queue *ubq, int tag, size_t offset);
+		const struct ublk_queue *ubq, int tag, size_t offset);
 static inline unsigned int ublk_req_build_flags(struct request *req);
 static inline struct ublksrv_io_desc *ublk_get_iod(struct ublk_queue *ubq,
 						   int tag);
@@ -1867,7 +1867,7 @@ static void ublk_io_release(void *priv)
 }
 
 static int ublk_register_io_buf(struct io_uring_cmd *cmd,
-				struct ublk_queue *ubq, unsigned int tag,
+				const struct ublk_queue *ubq, unsigned int tag,
 				unsigned int index, unsigned int issue_flags)
 {
 	struct ublk_device *ub = cmd->file->private_data;
@@ -2043,7 +2043,7 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 }
 
 static inline struct request *__ublk_check_and_get_req(struct ublk_device *ub,
-		struct ublk_queue *ubq, int tag, size_t offset)
+		const struct ublk_queue *ubq, int tag, size_t offset)
 {
 	struct request *req;
 

-- 
2.34.1

[PATCH v4 4/4] ublk: mark ublk_queue as const for ublk_handle_need_get_data

[Uday Shankar]

We now allow multiple tasks to operate on I/Os belonging to the same
queue concurrently. This means that any writes to ublk_queue in the I/O
path are potential sources of data races. Try to prevent these by
marking ublk_queue pointers as const in ublk_handle_need_get_data. Also
move a bit more of the NEED_GET_DATA-specific logic into
ublk_handle_need_get_data, to make the pattern in __ublk_ch_uring_cmd
more uniform.

Suggested-by: Ming Lei <ming.lei@redhat.com>
Signed-off-by: Uday Shankar <ushankar@purestorage.com>
---
 drivers/block/ublk_drv.c | 33 ++++++++++++++++++++-------------
 1 file changed, 20 insertions(+), 13 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index e2cb54895481aebaa91ab23ba05cf26a950a642f..c8ce9349ca280b8b16040a1242a62b895ee01b5d 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -1291,7 +1291,7 @@ static void ublk_cmd_tw_cb(struct io_uring_cmd *cmd,
 	ublk_dispatch_req(ubq, pdu->req, issue_flags);
 }
 
-static void ublk_queue_cmd(struct ublk_queue *ubq, struct request *rq)
+static void ublk_queue_cmd(const struct ublk_queue *ubq, struct request *rq)
 {
 	struct io_uring_cmd *cmd = ubq->ios[rq->tag].cmd;
 	struct ublk_uring_cmd_pdu *pdu = ublk_get_uring_cmd_pdu(cmd);
@@ -1813,15 +1813,6 @@ static void ublk_mark_io_ready(struct ublk_device *ub, struct ublk_queue *ubq)
 	mutex_unlock(&ub->mutex);
 }
 
-static void ublk_handle_need_get_data(struct ublk_device *ub, int q_id,
-		int tag)
-{
-	struct ublk_queue *ubq = ublk_get_queue(ub, q_id);
-	struct request *req = blk_mq_tag_to_rq(ub->tag_set.tags[q_id], tag);
-
-	ublk_queue_cmd(ubq, req);
-}
-
 static inline int ublk_check_cmd_op(u32 cmd_op)
 {
 	u32 ioc_type = _IOC_TYPE(cmd_op);
@@ -1933,6 +1924,21 @@ static int ublk_commit_and_fetch(const struct ublk_queue *ubq,
 	return -EIOCBQUEUED;
 }
 
+static int ublk_handle_need_get_data(const struct ublk_queue *ubq,
+				     struct ublk_io *io,
+				     struct io_uring_cmd *cmd,
+				     const struct ublksrv_io_cmd *ub_cmd,
+				     struct request *req)
+{
+	if (!(io->flags & UBLK_IO_FLAG_OWNED_BY_SRV))
+		return -EINVAL;
+
+	ublk_fill_io_cmd(io, cmd, ub_cmd->addr);
+	ublk_queue_cmd(ubq, req);
+
+	return -EIOCBQUEUED;
+}
+
 static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 			       unsigned int issue_flags,
 			       const struct ublksrv_io_cmd *ub_cmd)
@@ -2025,10 +2031,11 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 			goto out;
 		break;
 	case UBLK_IO_NEED_GET_DATA:
-		if (!(io->flags & UBLK_IO_FLAG_OWNED_BY_SRV))
+		ret = ublk_handle_need_get_data(
+			ubq, io, cmd, ub_cmd,
+			blk_mq_tag_to_rq(ub->tag_set.tags[ub_cmd->q_id], tag));
+		if (ret != -EIOCBQUEUED)
 			goto out;
-		ublk_fill_io_cmd(io, cmd, ub_cmd->addr);
-		ublk_handle_need_get_data(ub, ub_cmd->q_id, ub_cmd->tag);
 		break;
 	default:
 		goto out;

-- 
2.34.1

Re: [PATCH v4 1/4] ublk: require unique task per io instead of unique task per hctx

[Caleb Sander Mateos]

On Tue, Apr 15, 2025 at 6:00 PM Uday Shankar <ushankar@purestorage.com> wrote:
>
> Currently, ublk_drv associates to each hardware queue (hctx) a unique
> task (called the queue's ubq_daemon) which is allowed to issue
> COMMIT_AND_FETCH commands against the hctx. If any other task attempts
> to do so, the command fails immediately with EINVAL. When considered
> together with the block layer architecture, the result is that for each
> CPU C on the system, there is a unique ublk server thread which is
> allowed to handle I/O submitted on CPU C. This can lead to suboptimal
> performance under imbalanced load generation. For an extreme example,
> suppose all the load is generated on CPUs mapping to a single ublk
> server thread. Then that thread may be fully utilized and become the
> bottleneck in the system, while other ublk server threads are totally
> idle.
>
> This issue can also be addressed directly in the ublk server without
> kernel support by having threads dequeue I/Os and pass them around to
> ensure even load. But this solution requires inter-thread communication
> at least twice for each I/O (submission and completion), which is
> generally a bad pattern for performance. The problem gets even worse
> with zero copy, as more inter-thread communication would be required to
> have the buffer register/unregister calls to come from the correct
> thread.
>
> Therefore, address this issue in ublk_drv by requiring a unique task per
> I/O instead of per queue/hctx. Imbalanced load can then be balanced
> across all ublk server threads by having threads issue FETCH_REQs in a
> round-robin manner. As a small toy example, consider a system with a
> single ublk device having 2 queues, each of queue depth 4. A ublk server
> having 4 threads could issue its FETCH_REQs against this device as
> follows (where each entry is the qid,tag pair that the FETCH_REQ
> targets):
>
> poller thread:  T0      T1      T2      T3
>                 0,0     0,1     0,2     0,3
>                 1,3     1,0     1,1     1,2
>
> Since tags appear to be allocated in sequential chunks, this setup
> provides a rough approximation to distributing I/Os round-robin across
> all ublk server threads, while letting I/Os stay fully thread-local.
>
> Signed-off-by: Uday Shankar <ushankar@purestorage.com>
> Reviewed-by: Caleb Sander Mateos <csander@purestorage.com>
> ---
>  drivers/block/ublk_drv.c | 75 ++++++++++++++++++++++--------------------------
>  1 file changed, 34 insertions(+), 41 deletions(-)
>
> diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
> index cdb1543fa4a9817aa2ca2fca66720f589cf222be..9a0d2547512fc8119460739230599d48d2c2a306 100644
> --- a/drivers/block/ublk_drv.c
> +++ b/drivers/block/ublk_drv.c
> @@ -150,6 +150,7 @@ struct ublk_io {
>         int res;
>
>         struct io_uring_cmd *cmd;
> +       struct task_struct *task;
>  };
>
>  struct ublk_queue {
> @@ -157,11 +158,9 @@ struct ublk_queue {
>         int q_depth;
>
>         unsigned long flags;
> -       struct task_struct      *ubq_daemon;
>         struct ublksrv_io_desc *io_cmd_buf;
>
>         bool force_abort;
> -       bool timeout;
>         bool canceling;
>         bool fail_io; /* copy of dev->state == UBLK_S_DEV_FAIL_IO */
>         unsigned short nr_io_ready;     /* how many ios setup */
> @@ -1072,11 +1071,6 @@ static inline struct ublk_uring_cmd_pdu *ublk_get_uring_cmd_pdu(
>         return io_uring_cmd_to_pdu(ioucmd, struct ublk_uring_cmd_pdu);
>  }
>
> -static inline bool ubq_daemon_is_dying(struct ublk_queue *ubq)
> -{
> -       return ubq->ubq_daemon->flags & PF_EXITING;
> -}
> -
>  /* todo: handle partial completion */
>  static inline void __ublk_complete_rq(struct request *req)
>  {
> @@ -1224,13 +1218,13 @@ static void ublk_dispatch_req(struct ublk_queue *ubq,
>         /*
>          * Task is exiting if either:
>          *
> -        * (1) current != ubq_daemon.
> +        * (1) current != io->task.
>          * io_uring_cmd_complete_in_task() tries to run task_work
> -        * in a workqueue if ubq_daemon(cmd's task) is PF_EXITING.
> +        * in a workqueue if cmd's task is PF_EXITING.
>          *
>          * (2) current->flags & PF_EXITING.
>          */
> -       if (unlikely(current != ubq->ubq_daemon || current->flags & PF_EXITING)) {
> +       if (unlikely(current != io->task || current->flags & PF_EXITING)) {
>                 __ublk_abort_rq(ubq, req);
>                 return;
>         }
> @@ -1336,23 +1330,20 @@ static void ublk_queue_cmd_list(struct ublk_queue *ubq, struct rq_list *l)
>  static enum blk_eh_timer_return ublk_timeout(struct request *rq)
>  {
>         struct ublk_queue *ubq = rq->mq_hctx->driver_data;
> +       struct ublk_io *io = &ubq->ios[rq->tag];
>         unsigned int nr_inflight = 0;
>         int i;
>
>         if (ubq->flags & UBLK_F_UNPRIVILEGED_DEV) {
> -               if (!ubq->timeout) {
> -                       send_sig(SIGKILL, ubq->ubq_daemon, 0);
> -                       ubq->timeout = true;
> -               }
> -
> +               send_sig(SIGKILL, io->task, 0);
>                 return BLK_EH_DONE;
>         }
>
> -       if (!ubq_daemon_is_dying(ubq))
> +       if (!(io->task->flags & PF_EXITING))
>                 return BLK_EH_RESET_TIMER;
>
>         for (i = 0; i < ubq->q_depth; i++) {
> -               struct ublk_io *io = &ubq->ios[i];
> +               io = &ubq->ios[i];
>
>                 if (!(io->flags & UBLK_IO_FLAG_ACTIVE))
>                         nr_inflight++;
> @@ -1552,8 +1543,8 @@ static void ublk_commit_completion(struct ublk_device *ub,
>  }
>
>  /*
> - * Called from ubq_daemon context via cancel fn, meantime quiesce ublk
> - * blk-mq queue, so we are called exclusively with blk-mq and ubq_daemon
> + * Called from io task context via cancel fn, meantime quiesce ublk
> + * blk-mq queue, so we are called exclusively with blk-mq and io task
>   * context, so everything is serialized.
>   */
>  static void ublk_abort_queue(struct ublk_device *ub, struct ublk_queue *ubq)
> @@ -1669,13 +1660,13 @@ static void ublk_uring_cmd_cancel_fn(struct io_uring_cmd *cmd,
>                 return;
>
>         task = io_uring_cmd_get_task(cmd);
> -       if (WARN_ON_ONCE(task && task != ubq->ubq_daemon))
> +       io = &ubq->ios[pdu->tag];
> +       if (WARN_ON_ONCE(task && task != io->task))
>                 return;
>
>         ub = ubq->dev;
>         need_schedule = ublk_abort_requests(ub, ubq);
>
> -       io = &ubq->ios[pdu->tag];
>         WARN_ON_ONCE(io->cmd != cmd);
>         ublk_cancel_cmd(ubq, io, issue_flags);
>
> @@ -1836,8 +1827,6 @@ static void ublk_mark_io_ready(struct ublk_device *ub, struct ublk_queue *ubq)
>         mutex_lock(&ub->mutex);
>         ubq->nr_io_ready++;
>         if (ublk_queue_ready(ubq)) {
> -               ubq->ubq_daemon = current;
> -               get_task_struct(ubq->ubq_daemon);
>                 ub->nr_queues_ready++;
>
>                 if (capable(CAP_SYS_ADMIN))
> @@ -1952,14 +1941,14 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
>         if (!ubq || ub_cmd->q_id != ubq->q_id)
>                 goto out;
>
> -       if (ubq->ubq_daemon && ubq->ubq_daemon != current)
> -               goto out;
> -
>         if (tag >= ubq->q_depth)
>                 goto out;
>
>         io = &ubq->ios[tag];
>
> +       if (io->task && io->task != current)
> +               goto out;
> +
>         /* there is pending io cmd, something must be wrong */
>         if (io->flags & UBLK_IO_FLAG_ACTIVE) {
>                 ret = -EBUSY;
> @@ -2012,6 +2001,7 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
>
>                 ublk_fill_io_cmd(io, cmd, ub_cmd->addr);
>                 ublk_mark_io_ready(ub, ubq);
> +               io->task = get_task_struct(current);

Should io->task be set before ublk_mark_io_ready()? I worry that once
the ublk device is marked ready, it may be assumed that io->task is
constant.

Best,
Caleb

Re: [PATCH v4 2/4] ublk: mark ublk_queue as const for ublk_commit_and_fetch

[Caleb Sander Mateos]

On Tue, Apr 15, 2025 at 6:00 PM Uday Shankar <ushankar@purestorage.com> wrote:
>
> We now allow multiple tasks to operate on I/Os belonging to the same
> queue concurrently. This means that any writes to ublk_queue in the I/O
> path are potential sources of data races. Try to prevent these by
> marking ublk_queue pointers as const when handling COMMIT_AND_FETCH.
> Move the logic for this command into its own function
> ublk_commit_and_fetch. Also open code ublk_commit_completion in
> ublk_commit_and_fetch to reduce the number of parameters/avoid a
> redundant lookup.
>
> Suggested-by: Ming Lei <ming.lei@redhat.com>
> Signed-off-by: Uday Shankar <ushankar@purestorage.com>
> ---
>  drivers/block/ublk_drv.c | 91 +++++++++++++++++++++++-------------------------
>  1 file changed, 43 insertions(+), 48 deletions(-)
>
> diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
> index 9a0d2547512fc8119460739230599d48d2c2a306..153f67d92248ad45bddd2437b1306bb23df7d1ae 100644
> --- a/drivers/block/ublk_drv.c
> +++ b/drivers/block/ublk_drv.c
> @@ -1518,30 +1518,6 @@ static int ublk_ch_mmap(struct file *filp, struct vm_area_struct *vma)
>         return remap_pfn_range(vma, vma->vm_start, pfn, sz, vma->vm_page_prot);
>  }
>
> -static void ublk_commit_completion(struct ublk_device *ub,
> -               const struct ublksrv_io_cmd *ub_cmd)
> -{
> -       u32 qid = ub_cmd->q_id, tag = ub_cmd->tag;
> -       struct ublk_queue *ubq = ublk_get_queue(ub, qid);
> -       struct ublk_io *io = &ubq->ios[tag];
> -       struct request *req;
> -
> -       /* now this cmd slot is owned by nbd driver */
> -       io->flags &= ~UBLK_IO_FLAG_OWNED_BY_SRV;
> -       io->res = ub_cmd->result;
> -
> -       /* find the io request and complete */
> -       req = blk_mq_tag_to_rq(ub->tag_set.tags[qid], tag);
> -       if (WARN_ON_ONCE(unlikely(!req)))
> -               return;
> -
> -       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);
> -}
> -
>  /*
>   * Called from io task context via cancel fn, meantime quiesce ublk
>   * blk-mq queue, so we are called exclusively with blk-mq and io task
> @@ -1918,6 +1894,45 @@ static int ublk_unregister_io_buf(struct io_uring_cmd *cmd,
>         return io_buffer_unregister_bvec(cmd, index, issue_flags);
>  }
>
> +static int ublk_commit_and_fetch(const struct ublk_queue *ubq,
> +                                struct ublk_io *io, struct io_uring_cmd *cmd,
> +                                const struct ublksrv_io_cmd *ub_cmd,
> +                                struct request *req)
> +{
> +       if (!(io->flags & UBLK_IO_FLAG_OWNED_BY_SRV))
> +               return -EINVAL;
> +
> +       if (ublk_need_map_io(ubq)) {
> +               /*
> +                * COMMIT_AND_FETCH_REQ has to provide IO buffer if
> +                * NEED GET DATA is not enabled or it is Read IO.
> +                */
> +               if (!ub_cmd->addr && (!ublk_need_get_data(ubq) ||
> +                                       req_op(req) == REQ_OP_READ))
> +                       return -EINVAL;
> +       } else if (req_op(req) != REQ_OP_ZONE_APPEND && ub_cmd->addr) {
> +               /*
> +                * User copy requires addr to be unset when command is
> +                * not zone append
> +                */
> +               return -EINVAL;
> +       }
> +
> +       ublk_fill_io_cmd(io, cmd, ub_cmd->addr);
> +
> +       /* now this cmd slot is owned by ublk driver */
> +       io->flags &= ~UBLK_IO_FLAG_OWNED_BY_SRV;
> +       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, req);
> +
> +       return -EIOCBQUEUED;

I think it would be clearer to just return 0. __ublk_ch_uring_cmd()
already takes care of returning -EIOCBQUEUED in the successful case.
Aside from that,

Reviewed-by: Caleb Sander Mateos <csander@purestorage.com>

Re: [PATCH v4 3/4] ublk: mark ublk_queue as const for ublk_register_io_buf

[Caleb Sander Mateos]

On Tue, Apr 15, 2025 at 6:00 PM Uday Shankar <ushankar@purestorage.com> wrote:
>
> We now allow multiple tasks to operate on I/Os belonging to the same
> queue concurrently. This means that any writes to ublk_queue in the I/O
> path are potential sources of data races. Try to prevent these by
> marking ublk_queue pointers as const in ublk_register_io_buf.
>
> Suggested-by: Ming Lei <ming.lei@redhat.com>
> Signed-off-by: Uday Shankar <ushankar@purestorage.com>

Reviewed-by: Caleb Sander Mateos <csander@purestorage.com>

Re: [PATCH v4 4/4] ublk: mark ublk_queue as const for ublk_handle_need_get_data

[Caleb Sander Mateos]

On Tue, Apr 15, 2025 at 6:00 PM Uday Shankar <ushankar@purestorage.com> wrote:
>
> We now allow multiple tasks to operate on I/Os belonging to the same
> queue concurrently. This means that any writes to ublk_queue in the I/O
> path are potential sources of data races. Try to prevent these by
> marking ublk_queue pointers as const in ublk_handle_need_get_data. Also
> move a bit more of the NEED_GET_DATA-specific logic into
> ublk_handle_need_get_data, to make the pattern in __ublk_ch_uring_cmd
> more uniform.
>
> Suggested-by: Ming Lei <ming.lei@redhat.com>
> Signed-off-by: Uday Shankar <ushankar@purestorage.com>
> ---
>  drivers/block/ublk_drv.c | 33 ++++++++++++++++++++-------------
>  1 file changed, 20 insertions(+), 13 deletions(-)
>
> diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
> index e2cb54895481aebaa91ab23ba05cf26a950a642f..c8ce9349ca280b8b16040a1242a62b895ee01b5d 100644
> --- a/drivers/block/ublk_drv.c
> +++ b/drivers/block/ublk_drv.c
> @@ -1291,7 +1291,7 @@ static void ublk_cmd_tw_cb(struct io_uring_cmd *cmd,
>         ublk_dispatch_req(ubq, pdu->req, issue_flags);
>  }
>
> -static void ublk_queue_cmd(struct ublk_queue *ubq, struct request *rq)
> +static void ublk_queue_cmd(const struct ublk_queue *ubq, struct request *rq)
>  {
>         struct io_uring_cmd *cmd = ubq->ios[rq->tag].cmd;
>         struct ublk_uring_cmd_pdu *pdu = ublk_get_uring_cmd_pdu(cmd);
> @@ -1813,15 +1813,6 @@ static void ublk_mark_io_ready(struct ublk_device *ub, struct ublk_queue *ubq)
>         mutex_unlock(&ub->mutex);
>  }
>
> -static void ublk_handle_need_get_data(struct ublk_device *ub, int q_id,
> -               int tag)
> -{
> -       struct ublk_queue *ubq = ublk_get_queue(ub, q_id);
> -       struct request *req = blk_mq_tag_to_rq(ub->tag_set.tags[q_id], tag);
> -
> -       ublk_queue_cmd(ubq, req);
> -}
> -
>  static inline int ublk_check_cmd_op(u32 cmd_op)
>  {
>         u32 ioc_type = _IOC_TYPE(cmd_op);
> @@ -1933,6 +1924,21 @@ static int ublk_commit_and_fetch(const struct ublk_queue *ubq,
>         return -EIOCBQUEUED;
>  }
>
> +static int ublk_handle_need_get_data(const struct ublk_queue *ubq,
> +                                    struct ublk_io *io,
> +                                    struct io_uring_cmd *cmd,
> +                                    const struct ublksrv_io_cmd *ub_cmd,
> +                                    struct request *req)

nit: I see this is matching the name of the opcode (I am not sure why
it has "need" in its name) and there is already a function named
"ublk_need_get_data". But maybe naming this function "ublk_get_data"
would be clearer?

> +{
> +       if (!(io->flags & UBLK_IO_FLAG_OWNED_BY_SRV))
> +               return -EINVAL;
> +
> +       ublk_fill_io_cmd(io, cmd, ub_cmd->addr);
> +       ublk_queue_cmd(ubq, req);
> +
> +       return -EIOCBQUEUED;

Here too, I think a return value of 0 would be clearer.

Best,
Caleb

Re: [PATCH v4 1/4] ublk: require unique task per io instead of unique task per hctx

[Uday Shankar]

On Wed, Apr 16, 2025 at 11:48:45AM -0700, Caleb Sander Mateos wrote:
> On Tue, Apr 15, 2025 at 6:00 PM Uday Shankar <ushankar@purestorage.com> wrote:
> >
> > Currently, ublk_drv associates to each hardware queue (hctx) a unique
> > task (called the queue's ubq_daemon) which is allowed to issue
> > COMMIT_AND_FETCH commands against the hctx. If any other task attempts
> > to do so, the command fails immediately with EINVAL. When considered
> > together with the block layer architecture, the result is that for each
> > CPU C on the system, there is a unique ublk server thread which is
> > allowed to handle I/O submitted on CPU C. This can lead to suboptimal
> > performance under imbalanced load generation. For an extreme example,
> > suppose all the load is generated on CPUs mapping to a single ublk
> > server thread. Then that thread may be fully utilized and become the
> > bottleneck in the system, while other ublk server threads are totally
> > idle.
> >
> > This issue can also be addressed directly in the ublk server without
> > kernel support by having threads dequeue I/Os and pass them around to
> > ensure even load. But this solution requires inter-thread communication
> > at least twice for each I/O (submission and completion), which is
> > generally a bad pattern for performance. The problem gets even worse
> > with zero copy, as more inter-thread communication would be required to
> > have the buffer register/unregister calls to come from the correct
> > thread.
> >
> > Therefore, address this issue in ublk_drv by requiring a unique task per
> > I/O instead of per queue/hctx. Imbalanced load can then be balanced
> > across all ublk server threads by having threads issue FETCH_REQs in a
> > round-robin manner. As a small toy example, consider a system with a
> > single ublk device having 2 queues, each of queue depth 4. A ublk server
> > having 4 threads could issue its FETCH_REQs against this device as
> > follows (where each entry is the qid,tag pair that the FETCH_REQ
> > targets):
> >
> > poller thread:  T0      T1      T2      T3
> >                 0,0     0,1     0,2     0,3
> >                 1,3     1,0     1,1     1,2
> >
> > Since tags appear to be allocated in sequential chunks, this setup
> > provides a rough approximation to distributing I/Os round-robin across
> > all ublk server threads, while letting I/Os stay fully thread-local.
> >
> > Signed-off-by: Uday Shankar <ushankar@purestorage.com>
> > Reviewed-by: Caleb Sander Mateos <csander@purestorage.com>
> > ---
> >  drivers/block/ublk_drv.c | 75 ++++++++++++++++++++++--------------------------
> >  1 file changed, 34 insertions(+), 41 deletions(-)
> >
> > diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
> > index cdb1543fa4a9817aa2ca2fca66720f589cf222be..9a0d2547512fc8119460739230599d48d2c2a306 100644
> > --- a/drivers/block/ublk_drv.c
> > +++ b/drivers/block/ublk_drv.c
> > @@ -150,6 +150,7 @@ struct ublk_io {
> >         int res;
> >
> >         struct io_uring_cmd *cmd;
> > +       struct task_struct *task;
> >  };
> >
> >  struct ublk_queue {
> > @@ -157,11 +158,9 @@ struct ublk_queue {
> >         int q_depth;
> >
> >         unsigned long flags;
> > -       struct task_struct      *ubq_daemon;
> >         struct ublksrv_io_desc *io_cmd_buf;
> >
> >         bool force_abort;
> > -       bool timeout;
> >         bool canceling;
> >         bool fail_io; /* copy of dev->state == UBLK_S_DEV_FAIL_IO */
> >         unsigned short nr_io_ready;     /* how many ios setup */
> > @@ -1072,11 +1071,6 @@ static inline struct ublk_uring_cmd_pdu *ublk_get_uring_cmd_pdu(
> >         return io_uring_cmd_to_pdu(ioucmd, struct ublk_uring_cmd_pdu);
> >  }
> >
> > -static inline bool ubq_daemon_is_dying(struct ublk_queue *ubq)
> > -{
> > -       return ubq->ubq_daemon->flags & PF_EXITING;
> > -}
> > -
> >  /* todo: handle partial completion */
> >  static inline void __ublk_complete_rq(struct request *req)
> >  {
> > @@ -1224,13 +1218,13 @@ static void ublk_dispatch_req(struct ublk_queue *ubq,
> >         /*
> >          * Task is exiting if either:
> >          *
> > -        * (1) current != ubq_daemon.
> > +        * (1) current != io->task.
> >          * io_uring_cmd_complete_in_task() tries to run task_work
> > -        * in a workqueue if ubq_daemon(cmd's task) is PF_EXITING.
> > +        * in a workqueue if cmd's task is PF_EXITING.
> >          *
> >          * (2) current->flags & PF_EXITING.
> >          */
> > -       if (unlikely(current != ubq->ubq_daemon || current->flags & PF_EXITING)) {
> > +       if (unlikely(current != io->task || current->flags & PF_EXITING)) {
> >                 __ublk_abort_rq(ubq, req);
> >                 return;
> >         }
> > @@ -1336,23 +1330,20 @@ static void ublk_queue_cmd_list(struct ublk_queue *ubq, struct rq_list *l)
> >  static enum blk_eh_timer_return ublk_timeout(struct request *rq)
> >  {
> >         struct ublk_queue *ubq = rq->mq_hctx->driver_data;
> > +       struct ublk_io *io = &ubq->ios[rq->tag];
> >         unsigned int nr_inflight = 0;
> >         int i;
> >
> >         if (ubq->flags & UBLK_F_UNPRIVILEGED_DEV) {
> > -               if (!ubq->timeout) {
> > -                       send_sig(SIGKILL, ubq->ubq_daemon, 0);
> > -                       ubq->timeout = true;
> > -               }
> > -
> > +               send_sig(SIGKILL, io->task, 0);
> >                 return BLK_EH_DONE;
> >         }
> >
> > -       if (!ubq_daemon_is_dying(ubq))
> > +       if (!(io->task->flags & PF_EXITING))
> >                 return BLK_EH_RESET_TIMER;
> >
> >         for (i = 0; i < ubq->q_depth; i++) {
> > -               struct ublk_io *io = &ubq->ios[i];
> > +               io = &ubq->ios[i];
> >
> >                 if (!(io->flags & UBLK_IO_FLAG_ACTIVE))
> >                         nr_inflight++;
> > @@ -1552,8 +1543,8 @@ static void ublk_commit_completion(struct ublk_device *ub,
> >  }
> >
> >  /*
> > - * Called from ubq_daemon context via cancel fn, meantime quiesce ublk
> > - * blk-mq queue, so we are called exclusively with blk-mq and ubq_daemon
> > + * Called from io task context via cancel fn, meantime quiesce ublk
> > + * blk-mq queue, so we are called exclusively with blk-mq and io task
> >   * context, so everything is serialized.
> >   */
> >  static void ublk_abort_queue(struct ublk_device *ub, struct ublk_queue *ubq)
> > @@ -1669,13 +1660,13 @@ static void ublk_uring_cmd_cancel_fn(struct io_uring_cmd *cmd,
> >                 return;
> >
> >         task = io_uring_cmd_get_task(cmd);
> > -       if (WARN_ON_ONCE(task && task != ubq->ubq_daemon))
> > +       io = &ubq->ios[pdu->tag];
> > +       if (WARN_ON_ONCE(task && task != io->task))
> >                 return;
> >
> >         ub = ubq->dev;
> >         need_schedule = ublk_abort_requests(ub, ubq);
> >
> > -       io = &ubq->ios[pdu->tag];
> >         WARN_ON_ONCE(io->cmd != cmd);
> >         ublk_cancel_cmd(ubq, io, issue_flags);
> >
> > @@ -1836,8 +1827,6 @@ static void ublk_mark_io_ready(struct ublk_device *ub, struct ublk_queue *ubq)
> >         mutex_lock(&ub->mutex);
> >         ubq->nr_io_ready++;
> >         if (ublk_queue_ready(ubq)) {
> > -               ubq->ubq_daemon = current;
> > -               get_task_struct(ubq->ubq_daemon);
> >                 ub->nr_queues_ready++;
> >
> >                 if (capable(CAP_SYS_ADMIN))
> > @@ -1952,14 +1941,14 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
> >         if (!ubq || ub_cmd->q_id != ubq->q_id)
> >                 goto out;
> >
> > -       if (ubq->ubq_daemon && ubq->ubq_daemon != current)
> > -               goto out;
> > -
> >         if (tag >= ubq->q_depth)
> >                 goto out;
> >
> >         io = &ubq->ios[tag];
> >
> > +       if (io->task && io->task != current)
> > +               goto out;
> > +
> >         /* there is pending io cmd, something must be wrong */
> >         if (io->flags & UBLK_IO_FLAG_ACTIVE) {
> >                 ret = -EBUSY;
> > @@ -2012,6 +2001,7 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
> >
> >                 ublk_fill_io_cmd(io, cmd, ub_cmd->addr);
> >                 ublk_mark_io_ready(ub, ubq);
> > +               io->task = get_task_struct(current);
> 
> Should io->task be set before ublk_mark_io_ready()? I worry that once
> the ublk device is marked ready, it may be assumed that io->task is
> constant.

Not sure if we're dependent on that assumption anywhere, but it does
make sense. I'll move it.

I also think we might need to set atomically here, and read atomically
at the top of __ublk_ch_uring_cmd, as I can't see anything preventing
those accesses from being concurrent.

Re: [PATCH v4 2/4] ublk: mark ublk_queue as const for ublk_commit_and_fetch

[Uday Shankar]

On Wed, Apr 16, 2025 at 11:59:25AM -0700, Caleb Sander Mateos wrote:
> On Tue, Apr 15, 2025 at 6:00 PM Uday Shankar <ushankar@purestorage.com> wrote:
> >
> > We now allow multiple tasks to operate on I/Os belonging to the same
> > queue concurrently. This means that any writes to ublk_queue in the I/O
> > path are potential sources of data races. Try to prevent these by
> > marking ublk_queue pointers as const when handling COMMIT_AND_FETCH.
> > Move the logic for this command into its own function
> > ublk_commit_and_fetch. Also open code ublk_commit_completion in
> > ublk_commit_and_fetch to reduce the number of parameters/avoid a
> > redundant lookup.
> >
> > Suggested-by: Ming Lei <ming.lei@redhat.com>
> > Signed-off-by: Uday Shankar <ushankar@purestorage.com>
> > ---
> >  drivers/block/ublk_drv.c | 91 +++++++++++++++++++++++-------------------------
> >  1 file changed, 43 insertions(+), 48 deletions(-)
> >
> > diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
> > index 9a0d2547512fc8119460739230599d48d2c2a306..153f67d92248ad45bddd2437b1306bb23df7d1ae 100644
> > --- a/drivers/block/ublk_drv.c
> > +++ b/drivers/block/ublk_drv.c
> > @@ -1518,30 +1518,6 @@ static int ublk_ch_mmap(struct file *filp, struct vm_area_struct *vma)
> >         return remap_pfn_range(vma, vma->vm_start, pfn, sz, vma->vm_page_prot);
> >  }
> >
> > -static void ublk_commit_completion(struct ublk_device *ub,
> > -               const struct ublksrv_io_cmd *ub_cmd)
> > -{
> > -       u32 qid = ub_cmd->q_id, tag = ub_cmd->tag;
> > -       struct ublk_queue *ubq = ublk_get_queue(ub, qid);
> > -       struct ublk_io *io = &ubq->ios[tag];
> > -       struct request *req;
> > -
> > -       /* now this cmd slot is owned by nbd driver */
> > -       io->flags &= ~UBLK_IO_FLAG_OWNED_BY_SRV;
> > -       io->res = ub_cmd->result;
> > -
> > -       /* find the io request and complete */
> > -       req = blk_mq_tag_to_rq(ub->tag_set.tags[qid], tag);
> > -       if (WARN_ON_ONCE(unlikely(!req)))
> > -               return;
> > -
> > -       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);
> > -}
> > -
> >  /*
> >   * Called from io task context via cancel fn, meantime quiesce ublk
> >   * blk-mq queue, so we are called exclusively with blk-mq and io task
> > @@ -1918,6 +1894,45 @@ static int ublk_unregister_io_buf(struct io_uring_cmd *cmd,
> >         return io_buffer_unregister_bvec(cmd, index, issue_flags);
> >  }
> >
> > +static int ublk_commit_and_fetch(const struct ublk_queue *ubq,
> > +                                struct ublk_io *io, struct io_uring_cmd *cmd,
> > +                                const struct ublksrv_io_cmd *ub_cmd,
> > +                                struct request *req)
> > +{
> > +       if (!(io->flags & UBLK_IO_FLAG_OWNED_BY_SRV))
> > +               return -EINVAL;
> > +
> > +       if (ublk_need_map_io(ubq)) {
> > +               /*
> > +                * COMMIT_AND_FETCH_REQ has to provide IO buffer if
> > +                * NEED GET DATA is not enabled or it is Read IO.
> > +                */
> > +               if (!ub_cmd->addr && (!ublk_need_get_data(ubq) ||
> > +                                       req_op(req) == REQ_OP_READ))
> > +                       return -EINVAL;
> > +       } else if (req_op(req) != REQ_OP_ZONE_APPEND && ub_cmd->addr) {
> > +               /*
> > +                * User copy requires addr to be unset when command is
> > +                * not zone append
> > +                */
> > +               return -EINVAL;
> > +       }
> > +
> > +       ublk_fill_io_cmd(io, cmd, ub_cmd->addr);
> > +
> > +       /* now this cmd slot is owned by ublk driver */
> > +       io->flags &= ~UBLK_IO_FLAG_OWNED_BY_SRV;
> > +       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, req);
> > +
> > +       return -EIOCBQUEUED;
> 
> I think it would be clearer to just return 0. __ublk_ch_uring_cmd()
> already takes care of returning -EIOCBQUEUED in the successful case.

Sounds good - your recommendation is also in line with the convention
Ming is using in
https://lore.kernel.org/linux-block/20250416035444.99569-2-ming.lei@redhat.com/