perf(virtio-net): Don't copy packets when num_buffers is 1

Currently, we would copy every single received packet from DeviceAlloc
to a new Vec that is allocated with GlobalAlloc. If we use DeviceAlloc
for the payload of the RxToken, we can avoid copying all packets for the
special (and common) case of num_buffers = 1.

While at it, I refactored the code to make it more readable and avoid
allocating the intermediate Vec that would be collected into a flattened
Vec<u8> anyways.

Signed-off-by: Jens Reidel <[email protected]>

Author: Gelbpunkt

src/drivers/net/gem.rs

@@ -340,7 +340,7 @@ impl NetworkDriver for GEMDriver {
 				};
 				trace!("BUFFER: {buffer:x?}");
 				self.rx_buffer_consumed(index as usize);
-				Some((RxToken::new(buffer.to_vec()), TxToken::new()))
+				Some((RxToken::new(buffer.to_vec_in(DeviceAlloc)), TxToken::new()))
 			}
 			None => None,
 		}

src/drivers/net/rtl8139.rs

@@ -3,6 +3,7 @@
 #![allow(dead_code)]
 
 use alloc::boxed::Box;
+use alloc::vec::Vec;
 use core::mem;
 
 use memory_addresses::VirtAddr;
@@ -17,6 +18,7 @@ use crate::drivers::error::DriverError;
 use crate::drivers::net::NetworkDriver;
 use crate::drivers::pci::PciDevice;
 use crate::executor::device::{RxToken, TxToken};
+use crate::mm::device_alloc::DeviceAlloc;
 
 /// size of the receive buffer
 const RX_BUF_LEN: usize = 8192;
@@ -287,14 +289,18 @@ impl NetworkDriver for RTL8139Driver {
 		let length = self.rx_peek_u16() - 4; // copy packet (but not the CRC)
 		let pos = (self.rxpos + mem::size_of::<u16>()) % RX_BUF_LEN;
 
+		let mut vec_data = Vec::with_capacity_in(length as usize, DeviceAlloc);
+
 		// do we reach the end of the receive buffers?
 		// in this case, we contact the two slices to one vec
-		let vec_data = if pos + length as usize > RX_BUF_LEN {
+		if pos + length as usize > RX_BUF_LEN {
 			let first = &self.rxbuffer[pos..RX_BUF_LEN];
 			let second = &self.rxbuffer[..length as usize - first.len()];
-			[first, second].concat()
+
+			vec_data.extend_from_slice(first);
+			vec_data.extend_from_slice(second);
 		} else {
-			(self.rxbuffer[pos..][..length.into()]).to_vec()
+			vec_data.extend_from_slice(&self.rxbuffer[pos..][..length.into()]);
 		};
 
 		self.consume_current_buffer();

src/drivers/net/virtio/mod.rs

@@ -305,13 +305,17 @@ impl NetworkDriver for VirtioNetDriver {
 	}
 
 	fn receive_packet(&mut self) -> Option<(RxToken, TxToken)> {
-		let mut buffer_tkn = self.recv_vqs.get_next()?;
-		RxQueues::post_processing(&mut buffer_tkn)
-			.inspect_err(|vnet_err| warn!("Post processing failed. Err: {vnet_err:?}"))
-			.ok()?;
-		let first_header = buffer_tkn.used_recv_buff.pop_front_downcast::<Hdr>()?;
-		let first_packet = buffer_tkn.used_recv_buff.pop_front_vec()?;
-		trace!("Header: {first_header:?}");
+		let mut receive_single_packet = || {
+			let mut buffer_tkn = self.recv_vqs.get_next()?;
+			RxQueues::post_processing(&mut buffer_tkn)
+				.inspect_err(|vnet_err| warn!("Post processing failed. Err: {vnet_err:?}"))
+				.ok()?;
+			let header = buffer_tkn.used_recv_buff.pop_front_downcast::<Hdr>()?;
+			let packet = buffer_tkn.used_recv_buff.pop_front_vec()?;
+			Some((header, packet))
+		};
+
+		let (first_header, first_packet) = receive_single_packet()?;
 
 		// According to VIRTIO spec v1.2 sec. 5.1.6.3.2, "num_buffers will always be 1 if VIRTIO_NET_F_MRG_RXBUF is not negotiated."
 		// Unfortunately, NVIDIA MLX5 does not comply with this requirement and we have to manually set the value to the correct one.
@@ -321,17 +325,11 @@ impl NetworkDriver for VirtioNetDriver {
 			1
 		};
 
-		let mut packets = Vec::with_capacity(num_buffers.into());
-		packets.push(first_packet);
+		let mut combined_packets = first_packet;
 
 		for _ in 1..num_buffers {
-			let mut buffer_tkn = self.recv_vqs.get_next().unwrap();
-			RxQueues::post_processing(&mut buffer_tkn)
-				.inspect_err(|vnet_err| warn!("Post processing failed. Err: {vnet_err:?}"))
-				.ok()?;
-			let _header = buffer_tkn.used_recv_buff.pop_front_downcast::<Hdr>()?;
-			let packet = buffer_tkn.used_recv_buff.pop_front_vec()?;
-			packets.push(packet);
+			let (_header, packet) = receive_single_packet()?;
+			combined_packets.extend_from_slice(&packet);
 		}
 
 		fill_queue(
@@ -340,9 +338,7 @@ impl NetworkDriver for VirtioNetDriver {
 			self.recv_vqs.packet_size,
 		);
 
-		let vec_data = packets.into_iter().flatten().collect();
-
-		Some((RxToken::new(vec_data), TxToken::new()))
+		Some((RxToken::new(combined_packets), TxToken::new()))
 	}
 
 	fn set_polling_mode(&mut self, value: bool) {

src/executor/device.rs

@@ -23,6 +23,7 @@ use crate::arch::kernel::mmio as hardware;
 use crate::drivers::net::NetworkDriver;
 #[cfg(feature = "pci")]
 use crate::drivers::pci as hardware;
+use crate::mm::device_alloc::DeviceAlloc;
 
 /// Data type to determine the mac address
 #[derive(Debug, Clone)]
@@ -220,11 +221,11 @@ pub(crate) type RxHandle = usize;
 
 #[doc(hidden)]
 pub(crate) struct RxToken {
-	buffer: Vec<u8>,
+	buffer: Vec<u8, DeviceAlloc>,
 }
 
 impl RxToken {
-	pub(crate) fn new(buffer: Vec<u8>) -> Self {
+	pub(crate) fn new(buffer: Vec<u8, DeviceAlloc>) -> Self {
 		Self { buffer }
 	}
 }