|
| 1 | +use ahash::RandomState; |
| 2 | +use core::sync::atomic::{AtomicU32, Ordering::SeqCst}; |
| 3 | +use hashbrown::{hash_map::Entry, HashMap}; |
| 4 | + |
| 5 | +use crate::{ |
| 6 | + arch::kernel::{percore::core_scheduler, processor::get_timer_ticks}, |
| 7 | + errno::{EAGAIN, EINVAL, ETIMEDOUT}, |
| 8 | + scheduler::task::TaskHandlePriorityQueue, |
| 9 | +}; |
| 10 | + |
| 11 | +use super::spinlock::SpinlockIrqSave; |
| 12 | + |
| 13 | +// TODO: Replace with a concurrent hashmap. |
| 14 | +static PARKING_LOT: SpinlockIrqSave<HashMap<usize, TaskHandlePriorityQueue, RandomState>> = |
| 15 | + SpinlockIrqSave::new(HashMap::with_hasher(RandomState::with_seeds(0, 0, 0, 0))); |
| 16 | + |
| 17 | +bitflags! { |
| 18 | + pub struct Flags: u32 { |
| 19 | + /// Use a relative timeout |
| 20 | + const RELATIVE = 0b01; |
| 21 | + } |
| 22 | +} |
| 23 | + |
| 24 | +/// If the value at address matches the expected value, park the current thread until it is either |
| 25 | +/// woken up with `futex_wake` (returns 0) or the specified timeout elapses (returns -ETIMEDOUT). |
| 26 | +/// |
| 27 | +/// The timeout is given in microseconds. If [`Flags::RELATIVE`] is given, it is interpreted as |
| 28 | +/// relative to the current time. Otherwise it is understood to be an absolute time |
| 29 | +/// (see `get_timer_ticks`). |
| 30 | +pub fn futex_wait(address: &AtomicU32, expected: u32, timeout: Option<u64>, flags: Flags) -> i32 { |
| 31 | + let mut parking_lot = PARKING_LOT.lock(); |
| 32 | + // Check the futex value after locking the parking lot so that all changes are observed. |
| 33 | + if address.load(SeqCst) != expected { |
| 34 | + return -EAGAIN; |
| 35 | + } |
| 36 | + |
| 37 | + let wakeup_time = if flags.contains(Flags::RELATIVE) { |
| 38 | + timeout.and_then(|t| get_timer_ticks().checked_add(t)) |
| 39 | + } else { |
| 40 | + timeout |
| 41 | + }; |
| 42 | + |
| 43 | + let scheduler = core_scheduler(); |
| 44 | + scheduler.block_current_task(wakeup_time); |
| 45 | + let handle = scheduler.get_current_task_handle(); |
| 46 | + parking_lot |
| 47 | + .entry(address.as_mut_ptr().addr()) |
| 48 | + .or_default() |
| 49 | + .push(handle); |
| 50 | + drop(parking_lot); |
| 51 | + |
| 52 | + loop { |
| 53 | + scheduler.reschedule(); |
| 54 | + |
| 55 | + // Try to remove ourselves from the waiting queue. |
| 56 | + let mut parking_lot = PARKING_LOT.lock(); |
| 57 | + let mut wakeup = true; |
| 58 | + if let Entry::Occupied(mut queue) = parking_lot.entry(address.as_mut_ptr().addr()) { |
| 59 | + // If we are not in the waking queue, this must have been a wakeup. |
| 60 | + wakeup = !queue.get_mut().remove(handle); |
| 61 | + if queue.get().is_empty() { |
| 62 | + queue.remove(); |
| 63 | + } |
| 64 | + }; |
| 65 | + |
| 66 | + if wakeup { |
| 67 | + return 0; |
| 68 | + } else if wakeup_time.is_some_and(|&t| t <= get_timer_ticks()) { |
| 69 | + // If the current time is past the wakeup time, the operation timed out. |
| 70 | + return -ETIMEDOUT; |
| 71 | + } |
| 72 | + |
| 73 | + // A spurious wakeup occurred, sleep again. |
| 74 | + scheduler.block_current_task(wakeup_time); |
| 75 | + } |
| 76 | +} |
| 77 | + |
| 78 | +/// Wake `count` threads waiting on the futex at address. Returns the number of threads |
| 79 | +/// woken up (saturates to `i32::MAX`). If `count` is `i32::MAX`, wake up all matching |
| 80 | +/// waiting threads. If `count` is negative, returns -EINVAL. |
| 81 | +pub fn futex_wake(address: &AtomicU32, count: i32) -> i32 { |
| 82 | + if count < 0 { |
| 83 | + return -EINVAL; |
| 84 | + } |
| 85 | + |
| 86 | + let mut parking_lot = PARKING_LOT.lock(); |
| 87 | + let mut queue = match parking_lot.entry(address.as_mut_ptr().addr()) { |
| 88 | + Entry::Occupied(entry) => entry, |
| 89 | + Entry::Vacant(_) => return 0, |
| 90 | + }; |
| 91 | + |
| 92 | + let scheduler = core_scheduler(); |
| 93 | + let mut woken = 0; |
| 94 | + while woken != count || count == i32::MAX { |
| 95 | + match queue.get_mut().pop() { |
| 96 | + Some(handle) => scheduler.custom_wakeup(handle), |
| 97 | + None => break, |
| 98 | + } |
| 99 | + woken = woken.saturating_add(1); |
| 100 | + } |
| 101 | + |
| 102 | + if queue.get().is_empty() { |
| 103 | + queue.remove(); |
| 104 | + } |
| 105 | + |
| 106 | + woken |
| 107 | +} |
0 commit comments