feat: bucketized futexes

Cargo.toml

@@ -364,6 +364,7 @@ num_enum = { version = "0.7", default-features = false }
 pci_types = { version = "0.10" }
 pci-ids = { version = "0.2", optional = true }
 rand_chacha = { version = "0.10", default-features = false }
+seahash = "4.1.0"
 shlex = { version = "2", default-features = false }
 simple-shell = { version = "0.0.1", optional = true }
 smallvec = { version = "1", features = ["const_new"] }

src/synch/futex.rs

@@ -1,20 +1,123 @@
+use alloc::collections::LinkedList;
+use alloc::collections::linked_list::CursorMut;
 use core::sync::atomic::AtomicU32;
 use core::sync::atomic::Ordering::SeqCst;
 
-use ahash::RandomState;
-use hashbrown::HashMap;
-use hashbrown::hash_map::Entry;
-use hermit_sync::InterruptTicketMutex;
+use hermit_sync::{InterruptSpinMutex, InterruptSpinMutexGuard};
 
 use crate::arch::kernel::core_local::core_scheduler;
 use crate::arch::kernel::processor::get_timer_ticks;
 use crate::errno::Errno;
 use crate::scheduler::PerCoreSchedulerExt;
-use crate::scheduler::task::TaskHandlePriorityQueue;
+use crate::scheduler::task::{TaskHandle, TaskHandlePriorityQueue};
 
-// TODO: Replace with a concurrent hashmap.
-static PARKING_LOT: InterruptTicketMutex<HashMap<usize, TaskHandlePriorityQueue, RandomState>> =
-	InterruptTicketMutex::new(HashMap::with_hasher(RandomState::with_seeds(0, 0, 0, 0)));
+struct BucketElem(usize, TaskHandlePriorityQueue);
+
+type Bucket = InterruptSpinMutex<TaskListBucket>;
+
+#[repr(transparent)]
+struct TaskListBucket(LinkedList<BucketElem>);
+
+struct WrappedCursor<'a>(CursorMut<'a, BucketElem>);
+
+impl WrappedCursor<'_> {
+	pub fn pop(&mut self) -> Option<TaskHandle> {
+		match self.0.current() {
+			None => None,
+			Some(task_list) => {
+				let task = task_list.1.pop();
+
+				if task_list.1.is_empty() {
+					self.0.remove_current();
+				}
+
+				task
+			}
+		}
+	}
+}
+
+impl TaskListBucket {
+	pub fn insert_task(&mut self, address: usize, handle: TaskHandle) {
+		for elem in self.0.iter_mut() {
+			if elem.0 == address {
+				elem.1.push(handle);
+				return;
+			}
+		}
+
+		let mut task_list = TaskHandlePriorityQueue::new();
+		task_list.push(handle);
+		self.0.push_front(BucketElem(address, task_list));
+	}
+
+	pub fn contains_task(&self, address: usize, handle: TaskHandle) -> bool {
+		for elem in self.0.iter() {
+			if elem.0 == address {
+				return elem.1.contains(handle);
+			}
+		}
+		false
+	}
+
+	/// Removes a task from this bucket, and returns a boolean indicating if it was present.
+	pub fn remove_task(&mut self, address: usize, task: TaskHandle) -> bool {
+		let mut cursor = self.0.cursor_front_mut();
+		while let Some(elem) = cursor.current() {
+			if elem.0 == address {
+				let was_present = elem.1.remove(task);
+
+				if elem.1.is_empty() {
+					cursor.remove_current();
+				}
+
+				return was_present;
+			}
+			cursor.move_next();
+		}
+
+		false
+	}
+
+	fn get_pop_list(&mut self, address: usize) -> Option<WrappedCursor<'_>> {
+		let mut cursor = self.0.cursor_front_mut();
+		while let Some(elem) = cursor.current() {
+			if elem.0 == address {
+				return Some(WrappedCursor(cursor));
+			}
+			cursor.move_next();
+		}
+		None
+	}
+}
+
+struct BucketList<const N: usize>([Bucket; N]);
+
+impl<const N: usize> BucketList<N> {
+	pub const fn new() -> Self {
+		Self([const { InterruptSpinMutex::new(TaskListBucket(LinkedList::new())) }; N])
+	}
+
+	fn hash_key(v: usize) -> usize {
+		let v = (v >> 3).to_be_bytes();
+		let hashed = seahash::hash(&v) as usize;
+		hashed % N
+	}
+
+	pub fn lock_bucket(&self, address: usize) -> InterruptSpinMutexGuard<'_, TaskListBucket> {
+		if N == 1 {
+			return self.0[0].lock();
+		}
+		let bucket = Self::hash_key(address);
+		self.0[bucket].lock()
+	}
+}
+
+#[cfg(feature = "smp")]
+static PARKING_LOT: BucketList<64> = BucketList::new();
+
+#[cfg(not(feature = "smp"))]
+static PARKING_LOT: BucketList<1> = BucketList::new();
 
 bitflags! {
 	pub struct Flags: u32 {
@@ -23,6 +126,7 @@ bitflags! {
 	}
 }
 
+#[inline(always)]
 fn addr(addr: &AtomicU32) -> usize {
 	let ptr: *const _ = addr;
 	ptr.addr()
@@ -40,7 +144,8 @@ pub(crate) fn futex_wait(
 	timeout: Option<u64>,
 	flags: Flags,
 ) -> i32 {
-	let mut parking_lot = PARKING_LOT.lock();
+	let address_usize = addr(address);
+	let mut parking_lot = PARKING_LOT.lock_bucket(address_usize);
 	// Check the futex value after locking the parking lot so that all changes are observed.
 	if address.load(SeqCst) != expected {
 		return -i32::from(Errno::Again);
@@ -55,40 +160,34 @@ pub(crate) fn futex_wait(
 	let scheduler = core_scheduler();
 	scheduler.block_current_task(wakeup_time);
 	let handle = scheduler.get_current_task_handle();
-	parking_lot.entry(addr(address)).or_default().push(handle);
+	parking_lot.insert_task(address_usize, handle);
 	drop(parking_lot);
 
 	loop {
 		scheduler.reschedule();
+		// Assume this will return immediately (no other task on core!)
 
-		let mut parking_lot = PARKING_LOT.lock();
+		let mut parking_lot = PARKING_LOT.lock_bucket(address_usize);
 		if matches!(wakeup_time, Some(t) if t <= get_timer_ticks()) {
-			let mut wakeup = true;
 			// Timeout occurred, try to remove ourselves from the waiting queue.
-			if let Entry::Occupied(mut queue) = parking_lot.entry(addr(address)) {
-				// If we are not in the waking queue, this must have been a wakeup.
-				wakeup = !queue.get_mut().remove(handle);
-				if queue.get().is_empty() {
-					queue.remove();
-				}
-			}
+			let was_present = parking_lot.remove_task(address_usize, handle);
 
-			if wakeup {
-				return 0;
+			return if was_present {
+				-i32::from(Errno::Timedout)
 			} else {
-				return -i32::from(Errno::Timedout);
-			}
+				// If we are not in the waking queue, this must have been a wakeup.
+				0
+			};
 		} else {
-			// If we are not in the waking queue, this must have been a wakeup.
-			let wakeup = !matches!(parking_lot
-				.get(&addr(address)), Some(queue) if queue.contains(handle));
+			let is_in_queue = parking_lot.contains_task(address_usize, handle);
 
-			if wakeup {
-				return 0;
-			} else {
+			if is_in_queue {
 				// A spurious wakeup occurred, sleep again.
 				// Tasks do not change core, so the handle in the parking lot is still current.
 				scheduler.block_current_task(wakeup_time);
+			} else {
+				// If we are not in the waking queue, this must have been a wakeup.
+				return 0;
 			}
 		}
 		drop(parking_lot);
@@ -109,7 +208,8 @@ pub(crate) fn futex_wait_and_set(
 	flags: Flags,
 	new_value: u32,
 ) -> i32 {
-	let mut parking_lot = PARKING_LOT.lock();
+	let address_usize = addr(address);
+	let mut parking_lot = PARKING_LOT.lock_bucket(address_usize);
 	// Check the futex value after locking the parking lot so that all changes are observed.
 	if address.swap(new_value, SeqCst) != expected {
 		return -i32::from(Errno::Again);
@@ -124,40 +224,33 @@ pub(crate) fn futex_wait_and_set(
 	let scheduler = core_scheduler();
 	scheduler.block_current_task(wakeup_time);
 	let handle = scheduler.get_current_task_handle();
-	parking_lot.entry(addr(address)).or_default().push(handle);
+	parking_lot.insert_task(address_usize, handle);
 	drop(parking_lot);
 
 	loop {
 		scheduler.reschedule();
 
-		let mut parking_lot = PARKING_LOT.lock();
+		let mut parking_lot = PARKING_LOT.lock_bucket(address_usize);
 		if matches!(wakeup_time, Some(t) if t <= get_timer_ticks()) {
-			let mut wakeup = true;
 			// Timeout occurred, try to remove ourselves from the waiting queue.
-			if let Entry::Occupied(mut queue) = parking_lot.entry(addr(address)) {
-				// If we are not in the waking queue, this must have been a wakeup.
-				wakeup = !queue.get_mut().remove(handle);
-				if queue.get().is_empty() {
-					queue.remove();
-				}
-			}
+			let was_present = parking_lot.remove_task(address_usize, handle);
 
-			if wakeup {
-				return 0;
+			return if was_present {
+				-i32::from(Errno::Timedout)
 			} else {
-				return -i32::from(Errno::Timedout);
-			}
+				// If we are not in the waking queue, this must have been a wakeup.
+				0
+			};
 		} else {
-			// If we are not in the waking queue, this must have been a wakeup.
-			let wakeup = !matches!(parking_lot
-				.get(&addr(address)), Some(queue) if queue.contains(handle));
+			let is_in_queue = parking_lot.contains_task(address_usize, handle);
 
-			if wakeup {
-				return 0;
-			} else {
+			if is_in_queue {
 				// A spurious wakeup occurred, sleep again.
 				// Tasks do not change core, so the handle in the parking lot is still current.
 				scheduler.block_current_task(wakeup_time);
+			} else {
+				// If we are not in the waking queue, this must have been a wakeup.
+				return 0;
 			}
 		}
 		drop(parking_lot);
@@ -174,26 +267,22 @@ pub(crate) fn futex_wake(address: *const AtomicU32, count: i32) -> i32 {
 		return -i32::from(Errno::Inval);
 	}
 
-	let mut parking_lot = PARKING_LOT.lock();
-	let mut queue = match parking_lot.entry(address.addr()) {
-		Entry::Occupied(entry) => entry,
-		Entry::Vacant(_) => return 0,
+	let address_usize = address.addr();
+	let mut parking_lot = PARKING_LOT.lock_bucket(address_usize);
+	let Some(mut queue) = parking_lot.get_pop_list(address_usize) else {
+		return 0;
 	};
 
 	let scheduler = core_scheduler();
 	let mut woken = 0;
 	while woken != count || count == i32::MAX {
-		match queue.get_mut().pop() {
+		match queue.pop() {
 			Some(handle) => scheduler.custom_wakeup(handle),
 			None => break,
 		}
 		woken = woken.saturating_add(1);
 	}
 
-	if queue.get().is_empty() {
-		queue.remove();
-	}
-
 	woken
 }
 
@@ -206,29 +295,23 @@ pub(crate) fn futex_wake_or_set(address: &AtomicU32, count: i32, new_value: u32)
 		return -i32::from(Errno::Inval);
 	}
 
-	let mut parking_lot = PARKING_LOT.lock();
-	let mut queue = match parking_lot.entry(addr(address)) {
-		Entry::Occupied(entry) => entry,
-		Entry::Vacant(_) => {
-			address.store(new_value, SeqCst);
-			return 0;
-		}
+	let address_usize = addr(address);
+	let mut parking_lot = PARKING_LOT.lock_bucket(address_usize);
+	let Some(mut queue) = parking_lot.get_pop_list(address_usize) else {
+		address.store(new_value, SeqCst);
+		return 0;
 	};
 
 	let scheduler = core_scheduler();
 	let mut woken = 0;
 	while woken != count || count == i32::MAX {
-		match queue.get_mut().pop() {
+		match queue.pop() {
 			Some(handle) => scheduler.custom_wakeup(handle),
 			None => break,
 		}
 		woken = woken.saturating_add(1);
 	}
 
-	if queue.get().is_empty() {
-		queue.remove();
-	}
-
 	if woken == 0 {
 		address.store(new_value, SeqCst);
 	}