ublk: refactor auto buffer register in ublk_dispatch_req()

Refactor auto buffer register code and prepare for supporting batch IO
feature, and the main motivation is to put 'ublk_io' operation code
together, so that per-io lock can be applied for the code block.

The key changes are:
- Rename ublk_auto_buf_reg() as ublk_do_auto_buf_reg()
- Introduce an enum `auto_buf_reg_res` to represent the result of
  the buffer registration attempt (FAIL, FALLBACK, OK).
- Split the existing `ublk_do_auto_buf_reg` function into two:
  - `__ublk_do_auto_buf_reg`: Performs the actual buffer registration
    and returns the `auto_buf_reg_res` status.
  - `ublk_do_auto_buf_reg`: A wrapper that calls the internal function
    and handles the I/O preparation based on the result.
- Introduce `ublk_prep_auto_buf_reg_io` to encapsulate the logic for
  preparing the I/O for completion after buffer registration.
- Pass the `tag` directly to `ublk_auto_buf_reg_fallback` to avoid
  recalculating it.

This refactoring makes the control flow clearer and isolates the different
stages of the auto buffer registration process.

Reviewed-by: Caleb Sander Mateos <csander@purestorage.com>
Signed-off-by: Ming Lei <ming.lei@redhat.com>
Signed-off-by: Jens Axboe <axboe@kernel.dk>

diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c
index 7217146781fa3556c0faf841e2b8698e794fd4d0..4d2dc655683278e546bf5fcac0cc8cedb80e96ed 100644 (file)
--- a/drivers/block/ublk_drv.c
+++ b/drivers/block/ublk_drv.c
@@ -1168,17 +1168,37 @@ static inline void __ublk_abort_rq(struct ublk_queue *ubq,
 }
 
 static void
-ublk_auto_buf_reg_fallback(const struct ublk_queue *ubq, struct ublk_io *io)
+ublk_auto_buf_reg_fallback(const struct ublk_queue *ubq, unsigned tag)
 {
-       unsigned tag = io - ubq->ios;
        struct ublksrv_io_desc *iod = ublk_get_iod(ubq, tag);
 
        iod->op_flags |= UBLK_IO_F_NEED_REG_BUF;
 }
 
-static bool ublk_auto_buf_reg(const struct ublk_queue *ubq, struct request *req,
-                             struct ublk_io *io, struct io_uring_cmd *cmd,
-                             unsigned int issue_flags)
+enum auto_buf_reg_res {
+       AUTO_BUF_REG_FAIL,
+       AUTO_BUF_REG_FALLBACK,
+       AUTO_BUF_REG_OK,
+};
+
+static void ublk_prep_auto_buf_reg_io(const struct ublk_queue *ubq,
+                                     struct request *req, struct ublk_io *io,
+                                     struct io_uring_cmd *cmd,
+                                     enum auto_buf_reg_res res)
+{
+       if (res == AUTO_BUF_REG_OK) {
+               io->task_registered_buffers = 1;
+               io->buf_ctx_handle = io_uring_cmd_ctx_handle(cmd);
+               io->flags |= UBLK_IO_FLAG_AUTO_BUF_REG;
+       }
+       ublk_init_req_ref(ubq, io);
+       __ublk_prep_compl_io_cmd(io, req);
+}
+
+static enum auto_buf_reg_res
+__ublk_do_auto_buf_reg(const struct ublk_queue *ubq, struct request *req,
+                      struct ublk_io *io, struct io_uring_cmd *cmd,
+                      unsigned int issue_flags)
 {
        int ret;
 
@@ -1186,29 +1206,27 @@ static bool ublk_auto_buf_reg(const struct ublk_queue *ubq, struct request *req,
                                      io->buf.auto_reg.index, issue_flags);
        if (ret) {
                if (io->buf.auto_reg.flags & UBLK_AUTO_BUF_REG_FALLBACK) {
-                       ublk_auto_buf_reg_fallback(ubq, io);
-                       return true;
+                       ublk_auto_buf_reg_fallback(ubq, req->tag);
+                       return AUTO_BUF_REG_FALLBACK;
                }
                blk_mq_end_request(req, BLK_STS_IOERR);
-               return false;
+               return AUTO_BUF_REG_FAIL;
        }
 
-       io->task_registered_buffers = 1;
-       io->buf_ctx_handle = io_uring_cmd_ctx_handle(cmd);
-       io->flags |= UBLK_IO_FLAG_AUTO_BUF_REG;
-       return true;
+       return AUTO_BUF_REG_OK;
 }
 
-static bool ublk_prep_auto_buf_reg(struct ublk_queue *ubq,
-                                  struct request *req, struct ublk_io *io,
-                                  struct io_uring_cmd *cmd,
-                                  unsigned int issue_flags)
+static void ublk_do_auto_buf_reg(const struct ublk_queue *ubq, struct request *req,
+                                struct ublk_io *io, struct io_uring_cmd *cmd,
+                                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, cmd, issue_flags);
+       enum auto_buf_reg_res res = __ublk_do_auto_buf_reg(ubq, req, io, cmd,
+                       issue_flags);
 
-       return true;
+       if (res != AUTO_BUF_REG_FAIL) {
+               ublk_prep_auto_buf_reg_io(ubq, req, io, cmd, res);
+               io_uring_cmd_done(cmd, UBLK_IO_RES_OK, issue_flags);
+       }
 }
 
 static bool ublk_start_io(const struct ublk_queue *ubq, struct request *req,
@@ -1281,8 +1299,12 @@ static void ublk_dispatch_req(struct ublk_queue *ubq,
        if (!ublk_start_io(ubq, req, io))
                return;
 
-       if (ublk_prep_auto_buf_reg(ubq, req, io, io->cmd, issue_flags))
+       if (ublk_support_auto_buf_reg(ubq) && ublk_rq_has_data(req)) {
+               ublk_do_auto_buf_reg(ubq, req, io, io->cmd, issue_flags);
+       } else {
+               ublk_init_req_ref(ubq, io);
                ublk_complete_io_cmd(io, req, UBLK_IO_RES_OK, issue_flags);
+       }
 }
 
 static void ublk_cmd_tw_cb(struct io_uring_cmd *cmd,