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

["Clément Léger" <cleger@meta.com>]

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