[RFC PATCH 0/2] io_uring/net: Add bundle support for multishot recvmsg

[RFC PATCH 0/2] io_uring/net: Add bundle support for multishot recvmsg

[Clément Léger]

IORING_RECVSEND_BUNDLE is currently supported for IORING_OP_RECV but
rejected with -EINVAL for IORING_OP_RECVMSG, since the recvmsg output
envelope (struct io_uring_recvmsg_out, name and control data) was never
accounted for when bundling a receive across multiple provided buffers.

This series enables bundles for multishot recvmsg. Patch 1 is a
preparatory refactor with no functional change. Patch 2 adds the bundle
path for multishot recvmsg.

Note that contrary to plain recv which handles retry and can append more
data to a bundle, recvmsg does not support that since each payload is
associated to an envelope. trying to enqueue more data in an existing
bundle would leave an envelope in the middle of some payload buffers
which is not expected at all. The envelope must remain at the start of
the buffer.

Clément Léger (2):
  io_uring/net: Factor out receive buffer selection
  io_uring/net: Add bundle support for multishot recvmsg

 io_uring/net.c | 240 ++++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 196 insertions(+), 44 deletions(-)

-- 
Clément Léger

[RFC PATCH 1/2] io_uring/net: Factor out receive buffer selection

[Clément Léger]

In preparation of bundle support for recvmsg, split out the provided
buffer ring peek into io_recv_bundle_peek() and the existing recvmsg
buffer setup into io_recvmsg_buf_select(). No functional change
intended.

Signed-off-by: Clément Léger <cleger@meta.com>
---
 io_uring/net.c | 108 ++++++++++++++++++++++++++++++++-----------------
 1 file changed, 71 insertions(+), 37 deletions(-)

diff --git a/io_uring/net.c b/io_uring/net.c
index f5f4f3ce6ef3..1606824971ec 100644
--- a/io_uring/net.c
+++ b/io_uring/net.c
@@ -1018,6 +1018,67 @@ static int io_recvmsg_multishot(struct socket *sock, struct io_sr_msg *io,
 			kmsg->controllen + err;
 }
 
+static int io_recv_bundle_peek(struct io_kiocb *req,
+			       struct io_async_msghdr *kmsg,
+			       struct io_br_sel *sel, struct buf_sel_arg *arg,
+			       size_t max_len)
+{
+	struct io_sr_msg *sr = io_kiocb_to_cmd(req, struct io_sr_msg);
+	int nr_iovs;
+
+	*arg = (struct buf_sel_arg) {
+		.iovs = &kmsg->fast_iov,
+		.nr_iovs = 1,
+		.mode = KBUF_MODE_EXPAND,
+		.buf_group = sr->buf_group,
+		.max_len = max_len,
+	};
+
+	if (kmsg->vec.iovec) {
+		arg->nr_iovs = kmsg->vec.nr;
+		arg->iovs = kmsg->vec.iovec;
+		arg->mode |= KBUF_MODE_FREE;
+	}
+
+	nr_iovs = io_buffers_peek(req, arg, sel);
+	if (unlikely(nr_iovs < 0))
+		return nr_iovs;
+
+	if (arg->iovs != &kmsg->fast_iov && arg->iovs != kmsg->vec.iovec) {
+		kmsg->vec.nr = nr_iovs;
+		kmsg->vec.iovec = arg->iovs;
+		req->flags |= REQ_F_NEED_CLEANUP;
+	}
+	if (arg->partial_map)
+		sr->flags |= IORING_RECV_PARTIAL_MAP;
+	return nr_iovs;
+}
+
+static int io_recvmsg_buf_select(struct io_kiocb *req,
+				 struct io_async_msghdr *kmsg,
+				 struct io_br_sel *sel,
+				 unsigned int issue_flags)
+{
+	struct io_sr_msg *sr = io_kiocb_to_cmd(req, struct io_sr_msg);
+	size_t len = sr->len;
+	int ret;
+
+	*sel = io_buffer_select(req, &len, sr->buf_group, issue_flags);
+	if (!sel->addr)
+		return -ENOBUFS;
+
+	if (req->flags & REQ_F_APOLL_MULTISHOT) {
+		ret = io_recvmsg_prep_multishot(kmsg, sr, &sel->addr, &len);
+		if (ret) {
+			io_kbuf_recycle(req, sel->buf_list, issue_flags);
+			return ret;
+		}
+	}
+
+	iov_iter_ubuf(&kmsg->msg.msg_iter, ITER_DEST, sel->addr, len);
+	return 0;
+}
+
 int io_recvmsg(struct io_kiocb *req, unsigned int issue_flags)
 {
 	struct io_sr_msg *sr = io_kiocb_to_cmd(req, struct io_sr_msg);
@@ -1047,19 +1108,10 @@ int io_recvmsg(struct io_kiocb *req, unsigned int issue_flags)
 	sel.buf_list = NULL;
 	len = sr->len;
 	if (io_do_buffer_select(req)) {
-		sel = io_buffer_select(req, &len, sr->buf_group, issue_flags);
-		if (!sel.addr)
-			return -ENOBUFS;
-
-		if (req->flags & REQ_F_APOLL_MULTISHOT) {
-			ret = io_recvmsg_prep_multishot(kmsg, sr, &sel.addr, &len);
-			if (ret) {
-				io_kbuf_recycle(req, sel.buf_list, issue_flags);
-				return ret;
-			}
-		}
-
-		iov_iter_ubuf(&kmsg->msg.msg_iter, ITER_DEST, sel.addr, len);
+		ret = io_recvmsg_buf_select(req, kmsg, &sel, issue_flags);
+		if (unlikely(ret < 0))
+			return ret;
+		len = iov_iter_count(&kmsg->msg.msg_iter);
 	}
 
 	kmsg->msg.msg_get_inq = 1;
@@ -1130,39 +1182,21 @@ static int io_recv_buf_select(struct io_kiocb *req, struct io_async_msghdr *kmsg
 	 */
 	if (!(issue_flags & IO_URING_F_UNLOCKED) &&
 	    sr->flags & IORING_RECVSEND_BUNDLE) {
-		struct buf_sel_arg arg = {
-			.iovs = &kmsg->fast_iov,
-			.nr_iovs = 1,
-			.mode = KBUF_MODE_EXPAND,
-			.buf_group = sr->buf_group,
-		};
-
-		if (kmsg->vec.iovec) {
-			arg.nr_iovs = kmsg->vec.nr;
-			arg.iovs = kmsg->vec.iovec;
-			arg.mode |= KBUF_MODE_FREE;
-		}
+		struct buf_sel_arg arg;
+		size_t max_len = 0;
 
 		if (sel->val)
-			arg.max_len = sel->val;
+			max_len = sel->val;
 		else if (kmsg->msg.msg_inq > 1)
-			arg.max_len = min_not_zero(sel->val, (ssize_t) kmsg->msg.msg_inq);
+			max_len = min_not_zero(sel->val, (ssize_t)kmsg->msg.msg_inq);
 
 		/* if mshot limited, ensure we don't go over */
 		if (sr->flags & IORING_RECV_MSHOT_LIM)
-			arg.max_len = min_not_zero(arg.max_len, sr->mshot_total_len);
-		ret = io_buffers_peek(req, &arg, sel);
+			max_len = min_not_zero(max_len, sr->mshot_total_len);
+		ret = io_recv_bundle_peek(req, kmsg, sel, &arg, max_len);
 		if (unlikely(ret < 0))
 			return ret;
 
-		if (arg.iovs != &kmsg->fast_iov && arg.iovs != kmsg->vec.iovec) {
-			kmsg->vec.nr = ret;
-			kmsg->vec.iovec = arg.iovs;
-			req->flags |= REQ_F_NEED_CLEANUP;
-		}
-		if (arg.partial_map)
-			sr->flags |= IORING_RECV_PARTIAL_MAP;
-
 		/* special case 1 vec, can be a fast path */
 		if (ret == 1) {
 			sr->buf = arg.iovs[0].iov_base;
-- 
2.53.0-Meta

[RFC PATCH 2/2] io_uring/net: Add bundle support for multishot recvmsg

[Clément Léger]

IORING_RECVSEND_BUNDLE was disabled for IORING_OP_RECVMSG because the
envelope and control-message sizing was not implemented. Enable
bundled multishot recvmsg by reserving its metadata in the first buffer
and accounting for payload spread across the selected buffer run.

Signed-off-by: Clément Léger <cleger@meta.com>
---
 io_uring/net.c | 150 +++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 134 insertions(+), 16 deletions(-)

diff --git a/io_uring/net.c b/io_uring/net.c
index 1606824971ec..3ae2ae405ce3 100644
--- a/io_uring/net.c
+++ b/io_uring/net.c
@@ -735,6 +735,8 @@ static int io_recvmsg_copy_hdr(struct io_kiocb *req,
 	if (unlikely(ret))
 		return ret;
 
+	iomsg->msg.msg_inq = 0;
+
 	if (!(req->flags & REQ_F_BUFFER_SELECT)) {
 		ret = io_net_import_vec(req, iomsg, msg.msg_iov, msg.msg_iovlen,
 					ITER_DEST);
@@ -830,7 +832,8 @@ int io_recvmsg_prep(struct io_kiocb *req, const struct io_uring_sqe *sqe)
 	}
 
 	if (sr->flags & IORING_RECVSEND_BUNDLE) {
-		if (req->opcode == IORING_OP_RECVMSG)
+		if (req->opcode == IORING_OP_RECVMSG &&
+		    !(sr->flags & IORING_RECV_MULTISHOT))
 			return -EINVAL;
 	}
 
@@ -845,6 +848,37 @@ int io_recvmsg_prep(struct io_kiocb *req, const struct io_uring_sqe *sqe)
 #define CQE_F_MASK	(IORING_CQE_F_SOCK_NONEMPTY|IORING_CQE_F_MORE|\
 			 IORING_CQE_F_BUF_MORE)
 
+static int io_recv_payload_len(struct io_kiocb *req,
+			       struct io_async_msghdr *kmsg, int consumed)
+{
+	size_t hdr;
+
+	if (req->opcode != IORING_OP_RECVMSG ||
+	    !(req->flags & REQ_F_APOLL_MULTISHOT) || consumed <= 0)
+		return consumed;
+
+	if (check_add_overflow(sizeof(struct io_uring_recvmsg_out),
+			       (size_t)kmsg->namelen, &hdr) ||
+	    check_add_overflow(hdr, kmsg->controllen, &hdr) ||
+	    (size_t)consumed <= hdr)
+		return 0;
+
+	return consumed - (int)hdr;
+}
+
+static bool io_recvmsg_bundle_bufs_left(struct io_async_msghdr *kmsg,
+					int nbufs)
+{
+	struct iov_iter *iter = &kmsg->msg.msg_iter;
+	const struct iovec *iov;
+
+	if (!iter_is_iovec(iter))
+		return false;
+
+	iov = kmsg->vec.iovec ?: &kmsg->fast_iov;
+	return iter_iov(iter) - iov + iter->nr_segs > nbufs;
+}
+
 /*
  * Finishes io_recv and io_recvmsg.
  *
@@ -877,20 +911,37 @@ static inline bool io_recv_finish(struct io_kiocb *req,
 
 	if (sr->flags & IORING_RECVSEND_BUNDLE) {
 		size_t this_ret = sel->val - sr->done_io;
+		int nbufs;
+		int payload = io_recv_payload_len(req, kmsg, consumed);
+
+		/*
+		 * A zero-length plain recv recycles its selection before this point.
+		 * recvmsg still consumes the first buffer for its output envelope.
+		 */
+		if (req->opcode == IORING_OP_RECVMSG && consumed > 0 && !payload)
+			nbufs = 1;
+		else
+			nbufs = io_bundle_nbufs(kmsg, payload);
 
-		cflags |= io_put_kbufs(req, consumed, sel->buf_list, io_bundle_nbufs(kmsg, consumed));
+		cflags |= io_put_kbufs(req, consumed, sel->buf_list, nbufs);
 		if (sr->flags & IORING_RECV_RETRY)
 			cflags = req->cqe.flags | (cflags & CQE_F_MASK);
 		if (sr->mshot_len && sel->val >= sr->mshot_len)
 			sr->flags |= IORING_RECV_MSHOT_CAP;
+		if (req->opcode == IORING_OP_RECVMSG &&
+		    io_recvmsg_bundle_bufs_left(kmsg, nbufs))
+			req->flags &= ~REQ_F_BL_EMPTY;
 		/* bundle with no more immediate buffers, we're done */
 		if (req->flags & REQ_F_BL_EMPTY)
 			goto finish;
 		/*
 		 * If more is available AND it was a full transfer, retry and
-		 * append to this one
+		 * append to this one. This is not supported for recvmsg: each
+		 * CQE must carry exactly one envelope, so appending a second
+		 * recvmsg return into the same completion is not possible.
 		 */
-		if (!(sr->flags & IORING_RECV_NO_RETRY) &&
+		if (req->opcode != IORING_OP_RECVMSG &&
+		    !(sr->flags & IORING_RECV_NO_RETRY) &&
 		    kmsg->msg.msg_inq > 1 && this_ret > 0 &&
 		    !iov_iter_count(&kmsg->msg.msg_iter)) {
 			req->cqe.flags = cflags & ~CQE_F_MASK;
@@ -936,14 +987,14 @@ static inline bool io_recv_finish(struct io_kiocb *req,
 
 static int io_recvmsg_prep_multishot(struct io_async_msghdr *kmsg,
 				     struct io_sr_msg *sr, void __user **buf,
-				     size_t *len)
+				     size_t *len, size_t total)
 {
 	unsigned long ubuf = (unsigned long) *buf;
 	unsigned long hdr;
 
 	hdr = sizeof(struct io_uring_recvmsg_out) + kmsg->namelen +
 		kmsg->controllen;
-	if (*len < hdr)
+	if (*len < hdr || total < hdr)
 		return -EFAULT;
 
 	if (kmsg->controllen) {
@@ -955,7 +1006,8 @@ static int io_recvmsg_prep_multishot(struct io_async_msghdr *kmsg,
 
 	sr->buf = *buf; /* stash for later copy */
 	*buf = (void __user *) (ubuf + hdr);
-	kmsg->payloadlen = *len = *len - hdr;
+	*len = *len - hdr;
+	kmsg->payloadlen = total - hdr;
 	return 0;
 }
 
@@ -1054,28 +1106,94 @@ static int io_recv_bundle_peek(struct io_kiocb *req,
 	return nr_iovs;
 }
 
+static int io_recvmsg_bundle_select(struct io_kiocb *req,
+				    struct io_async_msghdr *kmsg,
+				    struct io_br_sel *sel)
+{
+	struct io_sr_msg *sr = io_kiocb_to_cmd(req, struct io_sr_msg);
+	struct buf_sel_arg arg;
+	size_t iov0_len, max_len;
+	int nr_iovs;
+	int ret;
+
+	if (sr->len)
+		max_len = sr->len;
+	else if (kmsg->msg.msg_inq > 1)
+		/* the envelope + cmsg also live in the mapped buffers */
+		max_len = kmsg->msg.msg_inq +
+			  sizeof(struct io_uring_recvmsg_out) +
+			  kmsg->namelen + kmsg->controllen;
+	else
+		/*
+		 * No size hint (datagram sockets don't report inq): map all
+		 * available buffers so a single datagram can span them. A
+		 * recvmsg returns at most one datagram and only the consumed
+		 * buffers are committed.
+		 */
+		max_len = MAX_RW_COUNT;
+
+	nr_iovs = io_recv_bundle_peek(req, kmsg, sel, &arg, max_len);
+	if (unlikely(nr_iovs <= 0))
+		return !nr_iovs ? -ENOBUFS : nr_iovs;
+	/* legacy provided buffers map a single buffer and don't set out_len */
+	if (!sel->buf_list)
+		arg.out_len = arg.iovs[0].iov_len;
+
+	if (sel->buf_list && sel->buf_list->flags & IOBL_INC)
+		return -EOPNOTSUPP;
+
+	/* reserve the envelope + cmsg at the front of the first buffer only */
+	iov0_len = arg.iovs[0].iov_len;
+	ret = io_recvmsg_prep_multishot(kmsg, sr, &arg.iovs[0].iov_base,
+					&iov0_len, arg.out_len);
+	if (unlikely(ret))
+		return ret;
+	arg.iovs[0].iov_len = iov0_len;
+
+	iov_iter_init(&kmsg->msg.msg_iter, ITER_DEST, arg.iovs, nr_iovs,
+		      kmsg->payloadlen);
+	return 0;
+}
+
 static int io_recvmsg_buf_select(struct io_kiocb *req,
 				 struct io_async_msghdr *kmsg,
 				 struct io_br_sel *sel,
 				 unsigned int issue_flags)
 {
 	struct io_sr_msg *sr = io_kiocb_to_cmd(req, struct io_sr_msg);
-	size_t len = sr->len;
 	int ret;
 
-	*sel = io_buffer_select(req, &len, sr->buf_group, issue_flags);
-	if (!sel->addr)
-		return -ENOBUFS;
-
-	if (req->flags & REQ_F_APOLL_MULTISHOT) {
-		ret = io_recvmsg_prep_multishot(kmsg, sr, &sel->addr, &len);
-		if (ret) {
+	/*
+	 * Bundle recvmsg spans multiple provided buffers, but the peek
+	 * interface can only be used when the ring is locked. If we came
+	 * in unlocked, fall back to a single-buffer envelope carve.
+	 */
+	if (!(issue_flags & IO_URING_F_UNLOCKED) &&
+	    sr->flags & IORING_RECVSEND_BUNDLE) {
+		ret = io_recvmsg_bundle_select(req, kmsg, sel);
+		if (unlikely(ret < 0)) {
 			io_kbuf_recycle(req, sel->buf_list, issue_flags);
 			return ret;
 		}
+	} else {
+		size_t len = sr->len;
+
+		*sel = io_buffer_select(req, &len, sr->buf_group, issue_flags);
+		if (!sel->addr)
+			return -ENOBUFS;
+
+		if (req->flags & REQ_F_APOLL_MULTISHOT) {
+			ret = io_recvmsg_prep_multishot(kmsg, sr, &sel->addr,
+							&len, len);
+			if (ret) {
+				io_kbuf_recycle(req, sel->buf_list, issue_flags);
+				return ret;
+			}
+		}
+
+		iov_iter_ubuf(&kmsg->msg.msg_iter, ITER_DEST, sel->addr, len);
 	}
 
-	iov_iter_ubuf(&kmsg->msg.msg_iter, ITER_DEST, sel->addr, len);
 	return 0;
 }
 
-- 
2.53.0-Meta