ublk: allow off-daemon zero-copy buffer registration

.github/workflows/kernel_build.yml

@@ -0,0 +1,28 @@
+name: blktests-ci
+
+on:
+  pull_request:
+
+jobs:
+  build-kernel:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Configure git
+        run: |
+          git config --global --add safe.directory '*'
+      - name: Checkout git
+        run: |
+          sudo apt-get install -y libelf-dev
+          mkdir -p linux
+          cd linux
+          git init
+          git remote add origin https://github.com/${{ github.repository }}
+          git fetch origin --depth=5 ${{ github.event.pull_request.head.sha }}
+          git reset --hard ${{ github.event.pull_request.head.sha }}
+          git log -1
+      - name: Build kernel
+        run: |
+          cd linux
+          make defconfig
+          make -j 8
+

drivers/block/ublk_drv.c

@@ -70,7 +70,8 @@
 		| 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 \
@@ -82,8 +83,16 @@
 	 UBLK_PARAM_TYPE_DEVT | UBLK_PARAM_TYPE_ZONED |    \
 	 UBLK_PARAM_TYPE_DMA_ALIGN | UBLK_PARAM_TYPE_SEGMENT)
 
+/*
+ * 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_rq_data {
 	refcount_t ref;
+	unsigned buffers_registered;
 
 	/* for auto-unregister buffer in case of UBLK_F_AUTO_BUF_REG */
 	u16 buf_index;
@@ -678,33 +687,33 @@ static inline void ublk_init_req_ref(const struct ublk_queue *ubq,
 	if (ublk_need_req_ref(ubq)) {
 		struct ublk_rq_data *data = blk_mq_rq_to_pdu(req);
 
-		refcount_set(&data->ref, 1);
+		refcount_set(&data->ref, UBLK_REFCOUNT_INIT);
+		data->buffers_registered = 0;
 	}
 }
 
-static inline bool ublk_get_req_ref(const struct ublk_queue *ubq,
-		struct request *req)
+static inline bool ublk_get_req_ref(struct request *req)
 {
-	if (ublk_need_req_ref(ubq)) {
-		struct ublk_rq_data *data = blk_mq_rq_to_pdu(req);
+	struct ublk_rq_data *data = blk_mq_rq_to_pdu(req);
 
-		return refcount_inc_not_zero(&data->ref);
-	}
+	return refcount_inc_not_zero(&data->ref);
+}
 
-	return true;
+static inline void ublk_put_req_ref(struct request *req)
+{
+	struct ublk_rq_data *data = blk_mq_rq_to_pdu(req);
+
+	if (refcount_dec_and_test(&data->ref))
+		__ublk_complete_rq(req);
 }
 
-static inline void ublk_put_req_ref(const struct ublk_queue *ubq,
-		struct request *req)
+static inline void ublk_sub_req_ref(struct request *req)
 {
-	if (ublk_need_req_ref(ubq)) {
-		struct ublk_rq_data *data = blk_mq_rq_to_pdu(req);
+	struct ublk_rq_data *data = blk_mq_rq_to_pdu(req);
+	unsigned sub_refs = UBLK_REFCOUNT_INIT - data->buffers_registered;
 
-		if (refcount_dec_and_test(&data->ref))
-			__ublk_complete_rq(req);
-	} else {
+	if (refcount_sub_and_test(sub_refs, &data->ref))
 		__ublk_complete_rq(req);
-	}
 }
 
 static inline bool ublk_need_get_data(const struct ublk_queue *ubq)
@@ -1185,10 +1194,8 @@ static void ublk_auto_buf_reg_fallback(struct request *req)
 {
 	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);
 
 	iod->op_flags |= UBLK_IO_F_NEED_REG_BUF;
-	refcount_set(&data->ref, 1);
 }
 
 static bool ublk_auto_buf_reg(struct request *req, struct ublk_io *io,
@@ -1208,9 +1215,8 @@ static bool ublk_auto_buf_reg(struct request *req, struct ublk_io *io,
 		blk_mq_end_request(req, BLK_STS_IOERR);
 		return false;
 	}
-	/* one extra reference is dropped by ublk_io_release */
-	refcount_set(&data->ref, 2);
 
+	data->buffers_registered = 1;
 	data->buf_ctx_handle = io_uring_cmd_ctx_handle(io->cmd);
 	/* store buffer index in request payload */
 	data->buf_index = pdu->buf.index;
@@ -1222,10 +1228,10 @@ 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, req);
 	if (ublk_support_auto_buf_reg(ubq) && ublk_rq_has_data(req))
 		return ublk_auto_buf_reg(req, io, issue_flags);
 
-	ublk_init_req_ref(ubq, req);
 	return true;
 }
 
@@ -1990,9 +1996,8 @@ 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;
 
-	ublk_put_req_ref(ubq, rq);
+	ublk_put_req_ref(rq);
 }
 
 static int ublk_register_io_buf(struct io_uring_cmd *cmd,
@@ -2013,21 +2018,32 @@ static int ublk_register_io_buf(struct io_uring_cmd *cmd,
 	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(req);
 		return ret;
 	}
 
 	return 0;
 }
 
-static int ublk_unregister_io_buf(struct io_uring_cmd *cmd,
-				  const struct ublk_queue *ubq,
-				  unsigned int index, unsigned int issue_flags)
+static int ublk_daemon_register_io_buf(struct io_uring_cmd *cmd,
+				       const struct ublk_queue *ubq,
+				       const struct ublk_io *io,
+				       unsigned index, unsigned issue_flags)
 {
-	if (!ublk_support_zero_copy(ubq))
+	struct request *req = io->req;
+	struct ublk_rq_data *data = blk_mq_rq_to_pdu(req);
+	int ret;
+
+	if (!ublk_support_zero_copy(ubq) || !ublk_rq_has_data(req))
 		return -EINVAL;
 
-	return io_buffer_unregister_bvec(cmd, index, issue_flags);
+	ret = io_buffer_register_bvec(cmd, req, ublk_io_release, index,
+				      issue_flags);
+	if (ret)
+		return ret;
+
+	data->buffers_registered++;
+	return 0;
 }
 
 static int ublk_fetch(struct io_uring_cmd *cmd, struct ublk_queue *ubq,
@@ -2142,9 +2158,13 @@ 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);
+	if (unlikely(blk_should_fake_timeout(req->q)))
+		return 0;
 
+	if (ublk_need_req_ref(ubq))
+		ublk_sub_req_ref(req);
+	else
+		__ublk_complete_rq(req);
 	return 0;
 }
 
@@ -2172,17 +2192,29 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 			       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 = -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;
+
+	/*
+	 * 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 io_buffer_unregister_bvec(cmd, ub_cmd->addr,
+						 issue_flags);
+
+	ret = -EINVAL;
 	if (ub_cmd->q_id >= ub->dev_info.nr_hw_queues)
 		goto out;
 
@@ -2192,21 +2224,34 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 		goto out;
 
 	io = &ubq->ios[tag];
-	task = READ_ONCE(io->task);
-	if (task && task != current)
+	/* 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;
+	}
+
+	if (READ_ONCE(io->task) != current) {
+		/*
+		 * ublk_register_io_buf() accesses only the request, not io,
+		 * so can be handled on any task
+		 */
+		if (_IOC_NR(cmd_op) == UBLK_IO_REGISTER_IO_BUF)
+			return ublk_register_io_buf(cmd, ubq, tag, ub_cmd->addr,
+						    issue_flags);
+
 		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;
 	}
 
-	/* 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)
-		goto out;
-
 	/*
 	 * ensure that the user issues UBLK_IO_NEED_GET_DATA
 	 * iff the driver have set the UBLK_IO_FLAG_NEED_GET_DATA.
@@ -2215,37 +2260,24 @@ static int __ublk_ch_uring_cmd(struct io_uring_cmd *cmd,
 			^ (_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, tag, 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;
+		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;
-		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",
@@ -2262,7 +2294,7 @@ static inline struct request *__ublk_check_and_get_req(struct ublk_device *ub,
 	if (!req)
 		return NULL;
 
-	if (!ublk_get_req_ref(ubq, req))
+	if (!ublk_get_req_ref(req))
 		return NULL;
 
 	if (unlikely(!blk_mq_request_started(req) || req->tag != tag))
@@ -2276,7 +2308,7 @@ static inline struct request *__ublk_check_and_get_req(struct ublk_device *ub,
 
 	return req;
 fail_put:
-	ublk_put_req_ref(ubq, req);
+	ublk_put_req_ref(req);
 	return NULL;
 }
 
@@ -2390,13 +2422,12 @@ static struct request *ublk_check_and_get_req(struct kiocb *iocb,
 	*off = buf_off;
 	return req;
 fail:
-	ublk_put_req_ref(ubq, req);
+	ublk_put_req_ref(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;
 	size_t buf_off;
 	size_t ret;
@@ -2406,15 +2437,13 @@ static ssize_t ublk_ch_read_iter(struct kiocb *iocb, struct iov_iter *to)
 		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(req);
 
 	return ret;
 }
 
 static ssize_t ublk_ch_write_iter(struct kiocb *iocb, struct iov_iter *from)
 {
-	struct ublk_queue *ubq;
 	struct request *req;
 	size_t buf_off;
 	size_t ret;
@@ -2424,8 +2453,7 @@ static ssize_t ublk_ch_write_iter(struct kiocb *iocb, struct iov_iter *from)
 		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(req);
 
 	return ret;
 }
@@ -2924,7 +2952,8 @@ static int ublk_ctrl_add_dev(const struct ublksrv_ctrl_cmd *header)
 
 	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 |