Rename the auto buffer registration functions for clarity: - __ublk_do_auto_buf_reg() -> ublk_auto_buf_register() - ublk_prep_auto_buf_reg_io() -> ublk_auto_buf_io_setup() - ublk_do_auto_buf_reg() -> ublk_auto_buf_dispatch() Add comments documenting the locking requirements for each function. No functional change. Signed-off-by: Ming Lei --- drivers/block/ublk_drv.c | 39 ++++++++++++++++++++++++++------------- 1 file changed, 26 insertions(+), 13 deletions(-) diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c index 31fda782c47c..7981decd1cee 100644 --- a/drivers/block/ublk_drv.c +++ b/drivers/block/ublk_drv.c @@ -1509,10 +1509,16 @@ enum auto_buf_reg_res { 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) +/* + * Setup io state after auto buffer registration. + * + * Must be called after ublk_auto_buf_register() is done. + * Caller must hold io->lock in batch context. + */ +static void ublk_auto_buf_io_setup(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; @@ -1523,8 +1529,9 @@ static void ublk_prep_auto_buf_reg_io(const struct ublk_queue *ubq, __ublk_prep_compl_io_cmd(io, req); } +/* Register request bvec to io_uring for auto buffer registration. */ static enum auto_buf_reg_res -__ublk_do_auto_buf_reg(const struct ublk_queue *ubq, struct request *req, +ublk_auto_buf_register(const struct ublk_queue *ubq, struct request *req, struct ublk_io *io, struct io_uring_cmd *cmd, unsigned int issue_flags) { @@ -1544,15 +1551,21 @@ __ublk_do_auto_buf_reg(const struct ublk_queue *ubq, struct request *req, return AUTO_BUF_REG_OK; } -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) +/* + * Dispatch IO to userspace with auto buffer registration. + * + * Only called in non-batch context from task work, io->lock not held. + */ +static void ublk_auto_buf_dispatch(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 res = __ublk_do_auto_buf_reg(ubq, req, io, cmd, + enum auto_buf_reg_res res = ublk_auto_buf_register(ubq, req, io, cmd, issue_flags); if (res != AUTO_BUF_REG_FAIL) { - ublk_prep_auto_buf_reg_io(ubq, req, io, cmd, res); + ublk_auto_buf_io_setup(ubq, req, io, cmd, res); io_uring_cmd_done(cmd, UBLK_IO_RES_OK, issue_flags); } } @@ -1627,7 +1640,7 @@ static void ublk_dispatch_req(struct ublk_queue *ubq, struct request *req) return; if (ublk_support_auto_buf_reg(ubq) && ublk_rq_has_data(req)) { - ublk_do_auto_buf_reg(ubq, req, io, io->cmd, issue_flags); + ublk_auto_buf_dispatch(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); @@ -1648,7 +1661,7 @@ static bool __ublk_batch_prep_dispatch(struct ublk_queue *ubq, return false; if (ublk_support_auto_buf_reg(ubq) && ublk_rq_has_data(req)) { - res = __ublk_do_auto_buf_reg(ubq, req, io, cmd, + res = ublk_auto_buf_register(ubq, req, io, cmd, data->issue_flags); if (res == AUTO_BUF_REG_FAIL) @@ -1656,7 +1669,7 @@ static bool __ublk_batch_prep_dispatch(struct ublk_queue *ubq, } ublk_io_lock(io); - ublk_prep_auto_buf_reg_io(ubq, req, io, cmd, res); + ublk_auto_buf_io_setup(ubq, req, io, cmd, res); ublk_io_unlock(io); return true; -- 2.47.0 Add comprehensive documentation for ublk's split reference counting model (io->ref + io->task_registered_buffers) above ublk_init_req_ref() given this model isn't very straightforward. Signed-off-by: Ming Lei --- drivers/block/ublk_drv.c | 64 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/drivers/block/ublk_drv.c b/drivers/block/ublk_drv.c index 7981decd1cee..91218b78e711 100644 --- a/drivers/block/ublk_drv.c +++ b/drivers/block/ublk_drv.c @@ -985,6 +985,70 @@ static inline bool ublk_dev_need_req_ref(const struct ublk_device *ub) ublk_dev_support_auto_buf_reg(ub); } +/* + * ublk IO Reference Counting Design + * ================================== + * + * For user-copy and zero-copy modes, ublk uses a split reference model with + * two counters that together track IO lifetime: + * + * - io->ref: refcount for off-task buffer registrations and user-copy ops + * - io->task_registered_buffers: count of buffers registered on the IO task + * + * Key Invariant: + * -------------- + * The sum (io->ref + io->task_registered_buffers) must equal UBLK_REFCOUNT_INIT + * when no active references exist. This invariant is checked by + * ublk_check_and_reset_active_ref() during daemon exit to determine if all + * references have been released. + * + * Why Split Counters: + * ------------------- + * Buffers registered on the IO daemon task can use the lightweight + * task_registered_buffers counter (simple increment/decrement) instead of + * atomic refcount operations. The ublk_io_release() callback checks if + * current == io->task to decide which counter to update. + * + * Reference Lifecycle: + * -------------------- + * 1. ublk_init_req_ref(): Sets io->ref = UBLK_REFCOUNT_INIT at IO dispatch + * + * 2. During IO processing: + * - On-task buffer reg: task_registered_buffers++ (no ref change) + * - Off-task buffer reg: ref++ via ublk_get_req_ref() + * - Buffer unregister callback (ublk_io_release): + * * If on-task: task_registered_buffers-- + * * If off-task: ref-- via ublk_put_req_ref() + * + * 3. ublk_sub_req_ref() at IO completion: + * - Computes: sub_refs = UBLK_REFCOUNT_INIT - task_registered_buffers + * - Subtracts sub_refs from ref + * - This accounts for the initial UBLK_REFCOUNT_INIT minus any on-task + * buffers that were already counted in task_registered_buffers + * + * Example (zero-copy, register on-task, unregister off-task): + * - Dispatch: ref = UBLK_REFCOUNT_INIT, task_registered_buffers = 0 + * - Register buffer on-task: task_registered_buffers = 1 + * - Unregister off-task: ref-- (UBLK_REFCOUNT_INIT - 1), task_registered_buffers stays 1 + * - Completion via ublk_sub_req_ref(): + * sub_refs = UBLK_REFCOUNT_INIT - 1, ref = (UBLK_REFCOUNT_INIT - 1) - (UBLK_REFCOUNT_INIT - 1) = 0 + * + * Example (auto buffer registration): + * Auto buffer registration sets task_registered_buffers = 1 at dispatch. + * + * - Dispatch: ref = UBLK_REFCOUNT_INIT, task_registered_buffers = 1 + * - Buffer unregister: task_registered_buffers-- (becomes 0) + * - Completion via ublk_sub_req_ref(): sub_refs = UBLK_REFCOUNT_INIT - 0, ref becomes 0 + * - Daemon exit check: sum = ref + task_registered_buffers = UBLK_REFCOUNT_INIT + * - Sum equals UBLK_REFCOUNT_INIT, so no active reference exists + * + * Batch IO Special Case: + * ---------------------- + * In batch IO mode, io->task is NULL. This means ublk_io_release() always + * takes the off-task path (ublk_put_req_ref), decrementing io->ref. The + * task_registered_buffers counter still tracks registered buffers for the + * invariant check, even though the callback doesn't decrement it. + */ static inline void ublk_init_req_ref(const struct ublk_queue *ubq, struct ublk_io *io) { -- 2.47.0