Skip to content

Commit

Permalink
For recvmmsg split lifetime 'a
Browse files Browse the repository at this point in the history
into separate lifetimes for 'data, headers ('hdr) and related iterator
  • Loading branch information
pacak committed Feb 10, 2025
1 parent 1b0dba6 commit 52c049f
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 27 deletions.
23 changes: 13 additions & 10 deletions src/sys/socket/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1765,7 +1765,7 @@ pub fn sendmmsg<'a, XS, AS, C, I, S>(
// shared across all the messages
cmsgs: C,
flags: MsgFlags
) -> crate::Result<MultiResults<'a, S>>
) -> crate::Result<MultiResults<'a, 'a, S>>
where
XS: IntoIterator<Item = &'a I>,
AS: AsRef<[Option<S>]>,
Expand Down Expand Up @@ -1820,6 +1820,7 @@ pub fn sendmmsg<'a, XS, AS, C, I, S>(
Ok(MultiResults {
rmm: data,
current_index: 0,
slices: std::marker::PhantomData,
received: sent
})

Expand Down Expand Up @@ -1908,16 +1909,16 @@ impl<S> MultiHeaders<S> {
// always produce the desired results - see https://github.com/nix-rust/nix/pull/1744 for more
// details
#[cfg(any(linux_android, target_os = "freebsd", target_os = "netbsd"))]
pub fn recvmmsg<'a, XS, S, I>(
pub fn recvmmsg<'hdr, 'iter, 'data, XS, S, I>(
fd: RawFd,
data: &'a mut MultiHeaders<S>,
data: &'hdr mut MultiHeaders<S>,
slices: XS,
flags: MsgFlags,
mut timeout: Option<crate::sys::time::TimeSpec>,
) -> crate::Result<MultiResults<'a, S>>
) -> crate::Result<MultiResults<'hdr, 'data, S>>
where
XS: IntoIterator<Item = &'a mut I>,
I: AsMut<[IoSliceMut<'a>]> + 'a,
XS: IntoIterator<Item = &'iter mut I>,
I: AsMut<[IoSliceMut<'data>]> + 'iter,
{
let mut count = 0;
for (i, (slice, mmsghdr)) in slices.into_iter().zip(data.items.iter_mut()).enumerate() {
Expand Down Expand Up @@ -1949,6 +1950,7 @@ where
})? as usize;

Ok(MultiResults {
slices: std::marker::PhantomData,
rmm: data,
current_index: 0,
received,
Expand All @@ -1958,19 +1960,20 @@ where
/// Iterator over results of [`recvmmsg`]/[`sendmmsg`]
#[cfg(any(linux_android, target_os = "freebsd", target_os = "netbsd"))]
#[derive(Debug)]
pub struct MultiResults<'a, S> {
pub struct MultiResults<'hdrs, 'data, S> {
// preallocated structures
rmm: &'a MultiHeaders<S>,
rmm: &'hdrs MultiHeaders<S>,
slices: std::marker::PhantomData<&'data ()>,
current_index: usize,
received: usize,
}

#[cfg(any(linux_android, target_os = "freebsd", target_os = "netbsd"))]
impl<'a, S> Iterator for MultiResults<'a, S>
impl<'hdrs, 'data, S> Iterator for MultiResults<'hdrs, 'data, S>
where
S: Copy + SockaddrLike,
{
type Item = RecvMsg<'a, 'a, S>;
type Item = RecvMsg<'hdrs, 'data, S>;

// The cast is not unnecessary on all platforms.
#[allow(clippy::unnecessary_cast)]
Expand Down
40 changes: 23 additions & 17 deletions test/sys/test_socket.rs
Original file line number Diff line number Diff line change
Expand Up @@ -646,8 +646,9 @@ mod recvfrom {
)
.expect("send socket failed");

const LOOPS: usize = 3;
let send_thread = thread::spawn(move || {
for _ in 0..NUM_MESSAGES_SENT {
for _ in 0..NUM_MESSAGES_SENT * LOOPS {
sendto(
ssock.as_raw_fd(),
&DATA[..],
Expand All @@ -671,24 +672,29 @@ mod recvfrom {
let mut data =
MultiHeaders::<SockaddrIn>::preallocate(msgs.len(), None);

let res: Vec<RecvMsg<SockaddrIn>> = recvmmsg(
rsock.as_raw_fd(),
&mut data,
msgs.iter_mut(),
MsgFlags::empty(),
None,
)
.expect("recvmmsg")
.collect();
assert_eq!(res.len(), DATA.len());
for _ in 0..LOOPS {
let res: Vec<RecvMsg<SockaddrIn>> = recvmmsg(
rsock.as_raw_fd(),
&mut data,
msgs.iter_mut(),
MsgFlags::empty(),
None,
)
.expect("recvmmsg")
.collect();
assert_eq!(res.len(), DATA.len());

for RecvMsg { address, bytes, .. } in res.into_iter() {
assert_eq!(AddressFamily::Inet, address.unwrap().family().unwrap());
assert_eq!(DATA.len(), bytes);
}
for RecvMsg { address, bytes, .. } in res.into_iter() {
assert_eq!(
AddressFamily::Inet,
address.unwrap().family().unwrap()
);
assert_eq!(DATA.len(), bytes);
}

for buf in &receive_buffers {
assert_eq!(&buf[..DATA.len()], DATA);
for buf in msgs.iter().flatten() {
assert_eq!(&buf[..DATA.len()], DATA);
}
}

send_thread.join().unwrap();
Expand Down

0 comments on commit 52c049f

Please sign in to comment.