diff --git a/lib/std/os/linux/io_uring.zig b/lib/std/os/linux/io_uring.zig index 74e6416102..ad78f7bfab 100644 --- a/lib/std/os/linux/io_uring.zig +++ b/lib/std/os/linux/io_uring.zig @@ -528,17 +528,39 @@ pub const IO_Uring = struct { return sqe; } + /// Used to select how the recv call should be handled. + pub const RecvBuffer = union(enum) { + /// io_uring will recv directly into this buffer + buffer: []u8, + + /// io_uring will select a buffer that has previously been provided with `provide_buffers`. + /// The buffer group referenced by `group_id` must contain at least one buffer for the recv call to work. + /// `len` controls the number of bytes to read into the selected buffer. + buffer_selection: struct { + group_id: u16, + len: usize, + }, + }; + /// Queues (but does not submit) an SQE to perform a `recv(2)`. /// Returns a pointer to the SQE. pub fn recv( self: *IO_Uring, user_data: u64, fd: os.fd_t, - buffer: []u8, + buffer: RecvBuffer, flags: u32, ) !*io_uring_sqe { const sqe = try self.get_sqe(); - io_uring_prep_recv(sqe, fd, buffer, flags); + switch (buffer) { + .buffer => |slice| io_uring_prep_recv(sqe, fd, slice, flags), + .buffer_selection => |selection| { + io_uring_prep_rw(.RECV, sqe, fd, 0, selection.len, 0); + sqe.rw_flags = flags; + sqe.flags |= linux.IOSQE_BUFFER_SELECT; + sqe.buf_index = selection.group_id; + }, + } sqe.user_data = user_data; return sqe; } @@ -2014,7 +2036,7 @@ test "accept/connect/send/recv" { const send = try ring.send(0xeeeeeeee, client, buffer_send[0..], 0); send.flags |= linux.IOSQE_IO_LINK; - _ = try ring.recv(0xffffffff, cqe_accept.res, buffer_recv[0..], 0); + _ = try ring.recv(0xffffffff, cqe_accept.res, .{ .buffer = buffer_recv[0..] }, 0); try testing.expectEqual(@as(u32, 2), try ring.submit()); const cqe_send = try ring.copy_cqe(); @@ -2282,7 +2304,7 @@ test "accept/connect/recv/link_timeout" { .flags = 0, }, cqe_connect); - const sqe_recv = try ring.recv(0xffffffff, cqe_accept.res, buffer_recv[0..], 0); + const sqe_recv = try ring.recv(0xffffffff, cqe_accept.res, .{ .buffer = buffer_recv[0..] }, 0); sqe_recv.flags |= linux.IOSQE_IO_LINK; const ts = os.linux.kernel_timespec{ .tv_sec = 0, .tv_nsec = 1000000 }; @@ -2469,7 +2491,7 @@ test "accept/connect/recv/cancel" { .flags = 0, }, cqe_connect); - _ = try ring.recv(0xffffffff, cqe_accept.res, buffer_recv[0..], 0); + _ = try ring.recv(0xffffffff, cqe_accept.res, .{ .buffer = buffer_recv[0..] }, 0); try testing.expectEqual(@as(u32, 1), try ring.submit()); const sqe_cancel = try ring.cancel(0x99999999, 0xffffffff, 0);