Flip iter operations to keep associated address/header/packets together (#22245)

Flip iter operations to keep associated address/header/packets together

Before this change, if cast_socket_addr() returned a None for any
address/header pair, the subsequent zip() would misalign the
address/header pair and packet. So, this change zips all three together,
then does filter_map() so keep things aligned.

Additionally, compute total_size inline to avoid running through packets
a second time.
This commit is contained in:
steviez 2022-01-03 18:15:50 -05:00 committed by GitHub
parent 73e6038986
commit 20b61e28b6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 9 deletions

View File

@ -99,21 +99,20 @@ pub fn recv_mmsg(sock: &UdpSocket, packets: &mut [Packet]) -> io::Result<(usize,
return Err(io::Error::last_os_error());
}
let mut npkts = 0;
addrs
.iter()
.zip(hdrs)
let mut total_size = 0;
izip!(addrs, hdrs, packets.iter_mut())
.take(nrecv as usize)
.filter_map(|(addr, hdr)| {
let addr = cast_socket_addr(addr, &hdr)?.to_std();
Some((addr, hdr))
.filter_map(|(addr, hdr, pkt)| {
let addr = cast_socket_addr(&addr, &hdr)?.to_std();
Some((addr, hdr, pkt))
})
.zip(packets.iter_mut())
.for_each(|((addr, hdr), pkt)| {
.for_each(|(addr, hdr, pkt)| {
pkt.meta.size = hdr.msg_len as usize;
pkt.meta.set_addr(&addr);
npkts += 1;
total_size += pkt.meta.size;
});
let total_size = packets.iter().take(npkts).map(|pkt| pkt.meta.size).sum();
Ok((total_size, npkts))
}