perf(virtqueue): replace vec-based MemPools with bitmap-based IndexAlloc

Author: mkroening

Cargo.lock

@@ -73,7 +73,7 @@ tcp = ["net", "smoltcp", "smoltcp/socket-tcp"]
 trace = ["smoltcp?/log", "smoltcp?/verbose"]
 udp = ["net", "smoltcp", "smoltcp/socket-udp"]
 vga = []
-virtio = ["dep:virtio"]
+virtio = ["dep:virtio", "dep:bitvec"]
 virtio-net = ["net", "virtio"]
 vsock = ["virtio", "pci"]
 
@@ -117,6 +117,7 @@ async-lock = { version = "3.4.1", default-features = false }
 async-trait = "0.1.89"
 bit_field = "0.10"
 bitflags = "2"
+bitvec = { version = "1", default-features = false, features = ["alloc"], optional = true }
 build-time = "0.1.3"
 cfg-if = "1"
 crossbeam-utils = { version = "0.8", default-features = false }

Cargo.toml

@@ -512,21 +512,44 @@ pub enum BufferType {
 	Indirect,
 }
 
-/// MemPool allows to easily control, request and provide memory for Virtqueues.
-struct MemPool {
-	pool: Vec<u16>,
-}
+mod index_alloc {
+	use bitvec::bitbox;
+	use bitvec::boxed::BitBox;
 
-impl MemPool {
-	/// Returns a given id to the id pool
-	fn ret_id(&mut self, id: u16) {
-		self.pool.push(id);
+	/// This type allows allocating indices.
+	///
+	/// The indices can be used as descriptor IDs.
+	pub struct IndexAlloc {
+		bits: BitBox,
 	}
 
-	/// Returns a new instance, with a pool of the specified size.
-	fn new(size: u16) -> MemPool {
-		MemPool {
-			pool: (0..size).collect(),
+	impl IndexAlloc {
+		pub fn new(len: usize) -> Self {
+			let bits = bitbox![0; len];
+			Self { bits }
+		}
+
+		pub fn allocate(&mut self) -> Option<usize> {
+			let index = self.bits.last_zero()?;
+			self.bits.get_mut(index)?.set(true);
+			Some(index)
+		}
+
+		unsafe fn try_free(&mut self, id: usize) -> Result<(), usize> {
+			let mut bit_ref = self.bits.get_mut(id).ok_or(id)?;
+
+			let was_set = bit_ref.replace(true);
+			if was_set {
+				return Err(id);
+			}
+
+			Ok(())
+		}
+
+		pub unsafe fn free(&mut self, id: usize) {
+			unsafe {
+				self.try_free(id).unwrap();
+			}
 		}
 	}
 }

src/drivers/virtio/virtqueue/mod.rs

@@ -22,9 +22,8 @@ use super::super::transport::mmio::{ComCfg, NotifCfg, NotifCtrl};
 #[cfg(feature = "pci")]
 use super::super::transport::pci::{ComCfg, NotifCfg, NotifCtrl};
 use super::error::VirtqError;
-use super::{
-	AvailBufferToken, BufferType, MemPool, TransferToken, UsedBufferToken, Virtq, VirtqPrivate,
-};
+use super::index_alloc::IndexAlloc;
+use super::{AvailBufferToken, BufferType, TransferToken, UsedBufferToken, Virtq, VirtqPrivate};
 use crate::arch::mm::paging::{BasePageSize, PageSize};
 use crate::mm::device_alloc::DeviceAlloc;
 
@@ -69,9 +68,8 @@ struct DescriptorRing {
 	/// See Virtio specification v1.1. - 2.7.1
 	drv_wc: bool,
 	dev_wc: bool,
-	/// Memory pool controls the amount of "free floating" descriptors
-	/// See [MemPool] docs for detail.
-	mem_pool: MemPool,
+	/// This allocates available descriptors.
+	indexes: IndexAlloc,
 }
 
 impl DescriptorRing {
@@ -92,7 +90,7 @@ impl DescriptorRing {
 			poll_index: 0,
 			drv_wc: true,
 			dev_wc: true,
-			mem_pool: MemPool::new(size),
+			indexes: IndexAlloc::new(size.into()),
 		}
 	}
 
@@ -186,13 +184,13 @@ impl DescriptorRing {
 	/// Returns an initialized write controller in order
 	/// to write the queue correctly.
 	fn get_write_ctrler(&mut self) -> Result<WriteCtrl<'_>, VirtqError> {
-		let desc_id = self.mem_pool.pool.pop().ok_or(VirtqError::NoDescrAvail)?;
+		let desc_id = self.indexes.allocate().ok_or(VirtqError::NoDescrAvail)?;
 		Ok(WriteCtrl {
 			start: self.write_index,
 			position: self.write_index,
 			modulo: u16::try_from(self.ring.len()).unwrap(),
 			first_flags: DescF::empty(),
-			buff_id: desc_id,
+			buff_id: u16::try_from(desc_id).unwrap(),
 
 			desc_ring: self,
 		})
@@ -303,7 +301,9 @@ impl ReadCtrl<'_> {
 			for _ in 0..tkn.num_consuming_descr() {
 				self.incrmt();
 			}
-			self.desc_ring.mem_pool.ret_id(buff_id);
+			unsafe {
+				self.desc_ring.indexes.free(buff_id.into());
+			}
 
 			Some((tkn, write_len))
 		} else {

src/drivers/virtio/virtqueue/packed.rs

@@ -17,16 +17,15 @@ use super::super::transport::mmio::{ComCfg, NotifCfg, NotifCtrl};
 #[cfg(feature = "pci")]
 use super::super::transport::pci::{ComCfg, NotifCfg, NotifCtrl};
 use super::error::VirtqError;
-use super::{
-	AvailBufferToken, BufferType, MemPool, TransferToken, UsedBufferToken, Virtq, VirtqPrivate,
-};
+use super::index_alloc::IndexAlloc;
+use super::{AvailBufferToken, BufferType, TransferToken, UsedBufferToken, Virtq, VirtqPrivate};
 use crate::arch::memory_barrier;
 use crate::mm::device_alloc::DeviceAlloc;
 
 struct DescrRing {
 	read_idx: u16,
 	token_ring: Box<[Option<Box<TransferToken<virtq::Desc>>>]>,
-	mem_pool: MemPool,
+	indexes: IndexAlloc,
 
 	/// Descriptor Tables
 	///
@@ -58,8 +57,8 @@ impl DescrRing {
 		if let Some(ctrl_desc) = tkn.ctrl_desc.as_ref() {
 			let descriptor = SplitVq::indirect_desc(ctrl_desc.as_ref());
 
-			index = self.mem_pool.pool.pop().ok_or(VirtqError::NoDescrAvail)?;
-			self.descr_table_mut()[usize::from(index)] = MaybeUninit::new(descriptor);
+			index = self.indexes.allocate().ok_or(VirtqError::NoDescrAvail)?;
+			self.descr_table_mut()[index] = MaybeUninit::new(descriptor);
 		} else {
 			let mut rev_all_desc_iter = SplitVq::descriptor_iter(&tkn.buff_tkn)?.rev();
 
@@ -68,25 +67,26 @@ impl DescrRing {
 				// If the [AvailBufferToken] is empty, we panic
 				let descriptor = rev_all_desc_iter.next().unwrap();
 
-				index = self.mem_pool.pool.pop().ok_or(VirtqError::NoDescrAvail)?;
-				self.descr_table_mut()[usize::from(index)] = MaybeUninit::new(descriptor);
+				index = self.indexes.allocate().ok_or(VirtqError::NoDescrAvail)?;
+				self.descr_table_mut()[index] = MaybeUninit::new(descriptor);
 			}
 			for mut descriptor in rev_all_desc_iter {
 				// We have not updated `index` yet, so it is at this point the index of the previous descriptor that had been written.
-				descriptor.next = le16::from(index);
+				descriptor.next = le16::from_ne(index.try_into().unwrap());
 
-				index = self.mem_pool.pool.pop().ok_or(VirtqError::NoDescrAvail)?;
-				self.descr_table_mut()[usize::from(index)] = MaybeUninit::new(descriptor);
+				index = self.indexes.allocate().ok_or(VirtqError::NoDescrAvail)?;
+				self.descr_table_mut()[index] = MaybeUninit::new(descriptor);
 			}
 			// At this point, `index` is the index of the last element of the reversed iterator,
 			// thus the head of the descriptor chain.
 		}
 
-		self.token_ring[usize::from(index)] = Some(Box::new(tkn));
+		self.token_ring[index] = Some(Box::new(tkn));
 
 		let len = self.token_ring.len();
 		let idx = self.avail_ring_mut().idx.to_ne();
-		self.avail_ring_mut().ring_mut(true)[idx as usize % len] = index.into();
+		self.avail_ring_mut().ring_mut(true)[idx as usize % len] =
+			le16::from_ne(index.try_into().unwrap());
 
 		memory_barrier();
 		let next_idx = idx.wrapping_add(1);
@@ -111,7 +111,9 @@ impl DescrRing {
 		// We return the indices of the now freed ring slots back to `mem_pool.`
 		let mut id_ret_idx = u16::try_from(used_elem.id.to_ne()).unwrap();
 		loop {
-			self.mem_pool.ret_id(id_ret_idx);
+			unsafe {
+				self.indexes.free(id_ret_idx.into());
+			}
 			let cur_chain_elem =
 				unsafe { self.descr_table_mut()[usize::from(id_ret_idx)].assume_init() };
 			if cur_chain_elem.flags.contains(virtq::DescF::NEXT) {
@@ -295,7 +297,7 @@ impl SplitVq {
 				.take(size.into())
 				.collect::<Vec<_>>()
 				.into_boxed_slice(),
-			mem_pool: MemPool::new(size),
+			indexes: IndexAlloc::new(size.into()),
 
 			descr_table_cell,
 			avail_ring_cell,