[PATCH 3/4] ublk: give the command back from COMMIT_AND_FETCH on a canceling queue

[Josef Bacik]

COMMIT_AND_FETCH and NEED_GET_DATA publish the server's next command in
io->cmd without a lock. A cancel which runs at the same time can see
the io active and still read the request pointer io->cmd shares its
storage with, or miss the command altogether. QUIESCE_DEV cancels while
the server still commits, and the command published right after its
cancel pass is never completed, which is what leaves the server waiting
for it forever.

Have the issuer decide instead: read ->canceling before publishing, and
on a canceling queue complete the committed request as usual, but give
the new command back with UBLK_IO_RES_ABORT instead of publishing it,
leaving the io canceled the way ublk_cancel_cmd() does. NEED_GET_DATA
sends its request back the way ublk_queue_rq() does on a canceling
queue.

The read and the publish are one RCU read section, so a cancel which
marks the queue and then calls synchronize_rcu() knows every command
published without seeing the mark is in place before it looks, and
every later one comes back from its issuer. That keeps the commit path
free of locks and barriers. ->canceling is read locklessly now, so write
it with WRITE_ONCE().

Assisted-by: LLM
Signed-off-by: Josef Bacik <josef@xxxxxxxxxxxxxx>
---
 drivers/block/ublk_drv.c | 66 +++++++++++++++++++++++++++++++++++++---
 1 file changed, 61 insertions(+), 5 deletions(-)

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index bd7126dcf92f..6717dabf3a23 100644
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -2492,6 +2492,9 @@ static void ublk_partition_scan_work(struct work_struct *work)
  * - there are no concurrent reads of ubq->canceling from the queue_rq
  *   path. This can be done by quiescing the queue, or through other
  *   means.
+ *
+ * ublk_commit_io_cmd() reads ubq->canceling locklessly, so it is written
+ * with WRITE_ONCE().
  */
 static void ublk_set_canceling(struct ublk_device *ub, bool canceling)
 	__must_hold(&ub->cancel_mutex)
@@ -2500,7 +2503,7 @@ static void ublk_set_canceling(struct ublk_device *ub, bool canceling)
 
 	ub->canceling = canceling;
 	for (i = 0; i < ub->dev_info.nr_hw_queues; i++)
-		ublk_get_queue(ub, i)->canceling = canceling;
+		WRITE_ONCE(ublk_get_queue(ub, i)->canceling, canceling);
 }
 
 static bool ublk_check_and_reset_active_ref(struct ublk_device *ub)
@@ -3109,7 +3112,7 @@ static void ublk_queue_reset_io_flags(struct ublk_device *ub,
 	 */
 	mutex_lock(&ub->cancel_mutex);
 	if (!ub->canceling)
-		ubq->canceling = false;
+		WRITE_ONCE(ubq->canceling, false);
 	mutex_unlock(&ub->cancel_mutex);
 	ubq->fail_io = false;
 	ubq->force_abort = false;
@@ -3227,6 +3230,49 @@ ublk_fill_io_cmd(struct ublk_io *io, struct io_uring_cmd *cmd)
 	return req;
 }
 
+/*
+ * Instead of ublk_fill_io_cmd() on a canceling queue: the server's new
+ * command is not published, it goes back with UBLK_IO_RES_ABORT, and the
+ * io is left canceled the way ublk_cancel_cmd() leaves it.  Returns the
+ * request the server owned.
+ */
+static struct request *ublk_cancel_io_cmd(struct ublk_queue *ubq,
+					  struct ublk_io *io)
+{
+	struct request *req = io->req;
+
+	spin_lock(&ubq->cancel_lock);
+	io->flags &= ~(UBLK_IO_FLAG_OWNED_BY_SRV | UBLK_IO_FLAG_NEED_GET_DATA);
+	io->flags |= UBLK_IO_FLAG_CANCELED;
+	spin_unlock(&ubq->cancel_lock);
+
+	return req;
+}
+
+/*
+ * Publish the command a COMMIT_AND_FETCH or NEED_GET_DATA brings, unless
+ * the queue is canceling, see ublk_quiesce_cancel().  The read of
+ * ->canceling and the publish are one RCU read section: a cancel which
+ * marks the queue after the read waits in synchronize_rcu() until the
+ * command is published, and claims it then.  Returns whether the command
+ * goes back to the server instead.
+ */
+static bool ublk_commit_io_cmd(struct ublk_queue *ubq, struct ublk_io *io,
+			       struct io_uring_cmd *cmd, struct request **req)
+{
+	bool canceling;
+
+	rcu_read_lock();
+	canceling = READ_ONCE(ubq->canceling);
+	if (likely(!canceling))
+		*req = ublk_fill_io_cmd(io, cmd);
+	else
+		*req = ublk_cancel_io_cmd(ubq, io);
+	rcu_read_unlock();
+
+	return canceling;
+}
+
 /*
  * Call before ublk_fill_io_cmd() publishes @cmd in io->cmd: a control-path
  * cancel may complete any command found there, and io_uring_cmd_done() only
@@ -3465,6 +3511,7 @@ static int ublk_ch_uring_cmd_local(struct io_uring_cmd *cmd,
 	u64 addr = READ_ONCE(ub_src->addr); /* unioned with zone_append_lba */
 	struct request *req;
 	int ret;
+	bool canceled;
 	bool compl;
 
 	WARN_ON_ONCE(issue_flags & IO_URING_F_UNLOCKED);
@@ -3547,7 +3594,7 @@ static int ublk_ch_uring_cmd_local(struct io_uring_cmd *cmd,
 			goto out;
 		io->res = result;
 		ublk_prep_cancel(cmd, issue_flags, ubq, tag);
-		req = ublk_fill_io_cmd(io, cmd);
+		canceled = ublk_commit_io_cmd(ubq, io, cmd, &req);
 		ublk_apply_io_buf(ub, io, cmd, addr, &auto_buf, &buf_idx);
 		if (buf_idx != UBLK_INVALID_BUF_IDX)
 			io_buffer_unregister(cmd, buf_idx, issue_flags);
@@ -3557,6 +3604,10 @@ static int ublk_ch_uring_cmd_local(struct io_uring_cmd *cmd,
 			req->__sector = addr;
 		if (compl)
 			__ublk_complete_rq(req, io, ublk_dev_need_map_io(ub), NULL);
+		if (unlikely(canceled)) {
+			ret = UBLK_IO_RES_ABORT;
+			goto out_done;
+		}
 		break;
 	}
 	case UBLK_IO_NEED_GET_DATA:
@@ -3566,7 +3617,12 @@ static int ublk_ch_uring_cmd_local(struct io_uring_cmd *cmd,
 		 * request
 		 */
 		ublk_prep_cancel(cmd, issue_flags, ubq, tag);
-		req = ublk_fill_io_cmd(io, cmd);
+		if (unlikely(ublk_commit_io_cmd(ubq, io, cmd, &req))) {
+			/* as ublk_queue_rq() does on a canceling queue */
+			__ublk_abort_rq(ubq, req);
+			ret = UBLK_IO_RES_ABORT;
+			goto out_done;
+		}
 		io->buf.addr = addr;
 		if (likely(ublk_get_data(ubq, io, req))) {
 			__ublk_prep_compl_io_cmd(io, req);
@@ -3765,7 +3821,7 @@ static int ublk_batch_unprep_io(struct ublk_queue *ubq,
 	if (ublk_queue_ready(ubq)) {
 		data->ub->nr_queue_ready--;
 		spin_lock(&ubq->cancel_lock);
-		ubq->canceling = true;
+		WRITE_ONCE(ubq->canceling, true);
 		spin_unlock(&ubq->cancel_lock);
 	}
 	ubq->nr_io_ready--;
-- 
2.55.0