strat9_kernel/syscall/

futex.rs

//! Futex (Fast Userspace Mutex) syscall handlers
use alloc::sync::Arc;
use core::sync::atomic::Ordering;

use super::error::SyscallError;
use crate::{
    memory::userslice::{UserSliceRead, UserSliceReadWrite},
    process::{block_current_task, current_task_id, wake_task},
    sync::{FixedQueue, SpinLock, SpinLockGuard},
};

/// Number of independent hash buckets.  Must be a power of two.
///
/// 256 buckets × 32 slots = 8 192 distinct futex addresses system-wide.
/// Increasing this reduces hot-bucket probability at the cost of 256 ×
/// sizeof(SpinLock<FutexBucket>) of static memory (~16 KiB on x86-64).
const FUTEX_QUEUE_BUCKETS: usize = 256;

/// Max distinct futex addresses per bucket.
///
/// With the Fibonacci hash below, collisions are rare; 32 slots handle
/// workloads with hundreds of concurrent futex addresses without spilling.
const FUTEX_BUCKET_CAPACITY: usize = 32;

/// Max concurrent waiters on a single futex address.
const FUTEX_WAITERS_CAPACITY: usize = 256;

// Fibonacci hashing requires FUTEX_QUEUE_BUCKETS to be a power of two so
// the bit-shift index calculation is correct.
const _: () = assert!(
    FUTEX_QUEUE_BUCKETS.is_power_of_two(),
    "FUTEX_QUEUE_BUCKETS must be a power of two"
);

struct FutexQueue {
    waiters: SpinLock<FixedQueue<crate::process::TaskId, FUTEX_WAITERS_CAPACITY>>,
}

impl FutexQueue {
    /// Creates a new instance.
    const fn new() -> Self {
        FutexQueue {
            waiters: SpinLock::new(FixedQueue::new()),
        }
    }

    /// Performs the pop waiter operation.
    fn pop_waiter(&self) -> Option<crate::process::TaskId> {
        let mut waiters = self.waiters.lock();
        waiters.pop_front()
    }

    /// Performs the remove waiter operation.
    fn remove_waiter(&self, id: crate::process::TaskId) {
        let mut waiters = self.waiters.lock();
        let _ = waiters.remove_first_where(|queued| *queued == id);
    }

    /// Returns whether empty.
    fn is_empty(&self) -> bool {
        self.waiters.lock().is_empty()
    }
}

struct FutexBucketEntry {
    addr: u64,
    queue: Arc<FutexQueue>,
}

struct FutexBucket {
    entries: [Option<FutexBucketEntry>; FUTEX_BUCKET_CAPACITY],
}

impl FutexBucket {
    const fn new() -> Self {
        Self {
            entries: [const { None }; FUTEX_BUCKET_CAPACITY],
        }
    }

    fn lookup(&self, addr: u64) -> Option<Arc<FutexQueue>> {
        self.entries.iter().find_map(|slot| match slot {
            Some(entry) if entry.addr == addr => Some(entry.queue.clone()),
            _ => None,
        })
    }

    fn get_or_insert(&mut self, addr: u64) -> Result<Arc<FutexQueue>, SyscallError> {
        let mut empty_slot = None;

        for (index, slot) in self.entries.iter_mut().enumerate() {
            match slot {
                Some(entry) if entry.addr == addr => return Ok(entry.queue.clone()),
                None if empty_slot.is_none() => empty_slot = Some(index),
                _ => {}
            }
        }

        let Some(index) = empty_slot else {
            return Err(SyscallError::QueueFull);
        };

        let queue = Arc::new(FutexQueue::new());
        self.entries[index] = Some(FutexBucketEntry {
            addr,
            queue: queue.clone(),
        });
        Ok(queue)
    }

    /// Removes the slot for `addr`/`queue` if the queue is still empty.
    ///
    /// # Lock ordering
    ///
    /// This method is called while the **bucket lock** is held.  It
    /// re-acquires `queue.waiters` internally to re-verify emptiness.
    /// Therefore the lock order on this path is:
    ///
    ///   `bucket_lock` → `queue.waiters`
    ///
    /// **Invariant:** no code path may hold `queue.waiters` while trying to
    /// acquire any `bucket_lock`.  Violating this would cause an ABBA
    /// deadlock with the ordering established here.  All callers of
    /// `queue.waiters.lock()` in the syscall handlers release the bucket
    /// lock before acquiring waiters (bucket is released inside `get_queue`/
    /// `lookup_queue` before the returned Arc is used).
    fn remove_if_empty(&mut self, addr: u64, queue: &Arc<FutexQueue>) {
        for slot in &mut self.entries {
            let should_remove = match slot {
                Some(entry)
                    if entry.addr == addr
                        && Arc::ptr_eq(&entry.queue, queue)
                        && entry.queue.is_empty() =>
                {
                    true
                }
                _ => false,
            };

            if should_remove {
                *slot = None;
                return;
            }
        }
    }
}

static FUTEX_QUEUES: [SpinLock<FutexBucket>; FUTEX_QUEUE_BUCKETS] =
    [const { SpinLock::new(FutexBucket::new()) }; FUTEX_QUEUE_BUCKETS];

#[inline]
fn futex_bucket_index(addr: u64) -> usize {
    // Fibonacci hashing: multiply by the 64-bit golden-ratio constant then
    // take the top byte.  This spreads all address bits across the index,
    // avoiding the hot-bucket problem of the old `(addr >> 2) & 63` which
    // only used 6 bits and caused all page-offset-0 addresses to collide.
    // FUTEX_QUEUE_BUCKETS must remain a power of two for the mask to work.
    let h = addr.wrapping_mul(0x9e37_79b9_7f4a_7c15_u64);
    (h >> (64 - FUTEX_QUEUE_BUCKETS.trailing_zeros())) as usize
}

#[inline]
fn futex_bucket(addr: u64) -> &'static SpinLock<FutexBucket> {
    &FUTEX_QUEUES[futex_bucket_index(addr)]
}

fn lookup_queue(addr: u64) -> Option<Arc<FutexQueue>> {
    let bucket = futex_bucket(addr).lock();
    bucket.lookup(addr)
}

/// Returns queue.
fn get_queue(addr: u64) -> Result<Arc<FutexQueue>, SyscallError> {
    let mut bucket = futex_bucket(addr).lock();
    bucket.get_or_insert(addr)
}

/// Reads u32.
fn read_u32(addr: u64) -> Result<u32, SyscallError> {
    let slice =
        UserSliceRead::new(addr, core::mem::size_of::<u32>()).map_err(|_| SyscallError::Fault)?;
    slice.read_val::<u32>().map_err(|_| SyscallError::Fault)
}

/// Performs the atomic cmpxchg u32 operation.
#[inline]
fn atomic_cmpxchg_u32(ptr: *mut u32, expected: u32, desired: u32) -> u32 {
    use core::sync::atomic::{AtomicU32, Ordering};

    // SAFETY: ptr validated by UserSliceReadWrite in caller. Creating a
    // shared &AtomicU32 from user memory is sound as long as the mapping
    // stays live (guaranteed by the UserSlice guard held above).
    let atomic = unsafe { &*ptr.cast::<AtomicU32>() };
    atomic
        .compare_exchange_weak(expected, desired, Ordering::AcqRel, Ordering::Acquire)
        .unwrap_or_else(|x| x)
}

/// Performs the atomic fetch update u32 operation.
fn atomic_fetch_update_u32<F>(addr: u64, update: F) -> Result<u32, SyscallError>
where
    F: Fn(u32) -> u32,
{
    if (addr & 0x3) != 0 {
        return Err(SyscallError::InvalidArgument);
    }
    let _slice = UserSliceReadWrite::new(addr, core::mem::size_of::<u32>())
        .map_err(|_| SyscallError::Fault)?;
    let ptr = addr as *mut u32;
    let mut cur = unsafe { core::ptr::read_volatile(ptr) };
    loop {
        let new = update(cur);
        let observed = atomic_cmpxchg_u32(ptr, cur, new);
        if observed == cur {
            return Ok(cur);
        }
        cur = observed;
    }
}

/// Performs the lock two queues operation.
fn lock_two_queues<'a>(
    addr1: u64,
    q1: &'a FutexQueue,
    addr2: u64,
    q2: &'a FutexQueue,
) -> (
    SpinLockGuard<'a, FixedQueue<crate::process::TaskId, FUTEX_WAITERS_CAPACITY>>,
    SpinLockGuard<'a, FixedQueue<crate::process::TaskId, FUTEX_WAITERS_CAPACITY>>,
) {
    debug_assert_ne!(addr1, addr2, "lock_two_queues: same address would deadlock");
    if addr1 < addr2 {
        let g1 = q1.waiters.lock();
        let g2 = q2.waiters.lock();
        (g1, g2)
    } else {
        let g2 = q2.waiters.lock();
        let g1 = q1.waiters.lock();
        (g1, g2)
    }
}

/// Performs the wake from waiters operation.
fn wake_from_waiters(
    waiters: &mut FixedQueue<crate::process::TaskId, FUTEX_WAITERS_CAPACITY>,
    max_wake: u32,
) -> u64 {
    let mut woke = 0u64;
    while woke < max_wake as u64 {
        if let Some(id) = waiters.pop_front() {
            if wake_task(id) {
                woke += 1;
            }
        } else {
            break;
        }
    }
    woke
}

/// Performs the do requeue operation.
fn do_requeue(
    addr1: u64,
    max_wake: u32,
    max_requeue: u32,
    addr2: u64,
) -> Result<u64, SyscallError> {
    if addr1 == addr2 {
        return sys_futex_wake(addr1, max_wake);
    }

    let queue1 = lookup_queue(addr1);
    let Some(queue1) = queue1 else {
        return Ok(0);
    };

    let queue2 = get_queue(addr2)?;

    let mut woke = 0u64;
    let mut requeued = 0u64;

    {
        let (mut w1, mut w2) = lock_two_queues(addr1, &queue1, addr2, &queue2);

        while woke < max_wake as u64 {
            if let Some(id) = w1.pop_front() {
                if wake_task(id) {
                    woke += 1;
                }
            } else {
                break;
            }
        }

        while requeued < max_requeue as u64 {
            if let Some(id) = w1.pop_front() {
                if w2.push_back(id).is_err() {
                    let _ = w1.push_back(id);
                    break;
                }
                requeued += 1;
            } else {
                break;
            }
        }
    }

    try_gc_queue(addr1, &queue1);
    try_gc_queue(addr2, &queue2);

    Ok(woke + requeued)
}

/// Removes the queue entry for `addr` from its bucket if the queue is empty.
///
/// Two-phase to avoid holding `bucket_lock` while checking emptiness (which
/// would require acquiring `waiters` under `bucket`, violating the intended
/// ordering of acquiring bucket before waiters):
///
///   Phase 1 : check under `waiters`: fast exit if non-empty.
///   Phase 2 : remove under `bucket_lock`: `remove_if_empty` re-verifies
///              emptiness while already holding the bucket lock.  A new
///              waiter that enqueued between the two phases will be visible
///              in that re-check and will block the removal.
///
/// TOCTOU note: a task may enqueue between phase 1 and phase 2.  This is
/// safe: `remove_if_empty` re-checks `is_empty()` under the bucket lock,
/// so a queue with live waiters is never removed.
fn try_gc_queue(addr: u64, queue: &Arc<FutexQueue>) {
    // Phase 1: early exit if still has waiters.
    {
        let waiters = queue.waiters.lock();
        if !waiters.is_empty() {
            return;
        }
    }
    // Phase 2: remove under bucket lock, with re-verification inside.
    let mut bucket = futex_bucket(addr).lock();
    bucket.remove_if_empty(addr, queue);
}

struct FutexWakeOpEncode {
    op: u32,
    is_oparg_shift: bool,
    cmp: u32,
    oparg: u32,
    cmparg: u32,
}

impl FutexWakeOpEncode {
    /// Performs the decode operation.
    fn decode(bits: u32) -> Result<Self, SyscallError> {
        let is_oparg_shift = ((bits >> 31) & 1) == 1;
        let op = (bits >> 28) & 0x7;
        let cmp = (bits >> 24) & 0xF;
        let oparg = (bits >> 12) & 0xFFF;
        let cmparg = bits & 0xFFF;

        if op > 4 || cmp > 5 {
            return Err(SyscallError::InvalidArgument);
        }

        Ok(Self {
            op,
            is_oparg_shift,
            cmp,
            oparg,
            cmparg,
        })
    }

    /// Performs the effective oparg operation.
    fn effective_oparg(&self) -> u32 {
        if self.is_oparg_shift {
            1u32 << (self.oparg & 31)
        } else {
            self.oparg
        }
    }

    /// Performs the calculate new val operation.
    fn calculate_new_val(&self, old_val: u32) -> u32 {
        let oparg = self.effective_oparg();
        match self.op {
            0 => oparg,
            1 => old_val.wrapping_add(oparg),
            2 => old_val | oparg,
            3 => old_val & !oparg,
            4 => old_val ^ oparg,
            _ => old_val,
        }
    }

    /// Performs the should wake operation.
    fn should_wake(&self, old_val: u32) -> bool {
        match self.cmp {
            0 => old_val == self.cmparg,
            1 => old_val != self.cmparg,
            2 => old_val < self.cmparg,
            3 => old_val <= self.cmparg,
            4 => old_val > self.cmparg,
            5 => old_val >= self.cmparg,
            _ => false,
        }
    }
}

/// SYS_FUTEX_WAIT: Wait on a futex
pub fn sys_futex_wait(_addr: u64, _val: u32, _timeout_ns: u64) -> Result<u64, SyscallError> {
    let addr = _addr;
    let val = _val;
    let timeout_ns = _timeout_ns;

    if (addr & 0x3) != 0 {
        return Err(SyscallError::InvalidArgument);
    }

    let id = current_task_id().ok_or(SyscallError::PermissionDenied)?;
    let queue = get_queue(addr)?;

    let saved_deadline = if timeout_ns != 0 {
        let deadline = crate::syscall::time::current_time_ns().saturating_add(timeout_ns);
        if let Some(task) = crate::process::get_task_by_id(id) {
            task.wake_deadline_ns.store(deadline, Ordering::Relaxed);
        }
        deadline
    } else {
        0
    };

    // Hold the futex queue lock while validating the futex word and enqueuing.
    {
        let mut waiters = queue.waiters.lock();
        let cur = read_u32(addr)?;
        if cur != val {
            if let Some(task) = crate::process::get_task_by_id(id) {
                task.wake_deadline_ns.store(0, Ordering::Relaxed);
            }
            return Err(SyscallError::Again);
        }
        waiters.push_back(id).map_err(|_| SyscallError::QueueFull)?;
    }

    block_current_task();

    queue.remove_waiter(id);
    try_gc_queue(addr, &queue);

    if let Some(task) = crate::process::get_task_by_id(id) {
        task.wake_deadline_ns.store(0, Ordering::Relaxed);
    }
    if timeout_ns != 0 && saved_deadline != 0 {
        let now = crate::syscall::time::current_time_ns();
        if now >= saved_deadline {
            return Err(SyscallError::TimedOut);
        }
    }

    Ok(0)
}

/// SYS_FUTEX_WAKE: Wake waiters on a futex
pub fn sys_futex_wake(_addr: u64, _max_wake: u32) -> Result<u64, SyscallError> {
    let addr = _addr;
    let max_wake = _max_wake;
    let queue = lookup_queue(addr);

    let Some(queue) = queue else {
        return Ok(0);
    };

    let mut woke = 0u64;
    while woke < max_wake as u64 {
        if let Some(id) = queue.pop_waiter() {
            if wake_task(id) {
                woke += 1;
            }
        } else {
            break;
        }
    }

    try_gc_queue(addr, &queue);

    Ok(woke)
}

/// SYS_FUTEX_REQUEUE: Requeue waiters from one futex to another
pub fn sys_futex_requeue(
    _addr1: u64,
    _max_wake: u32,
    _max_requeue: u32,
    _addr2: u64,
) -> Result<u64, SyscallError> {
    do_requeue(_addr1, _max_wake, _max_requeue, _addr2)
}

/// SYS_FUTEX_CMP_REQUEUE: Conditional requeue
pub fn sys_futex_cmp_requeue(
    _addr1: u64,
    _max_wake: u32,
    _max_requeue: u32,
    _addr2: u64,
    _expected_val: u32,
) -> Result<u64, SyscallError> {
    if _addr1 == _addr2 {
        let cur = read_u32(_addr1)?;
        if cur != _expected_val {
            return Err(SyscallError::Again);
        }
        return sys_futex_wake(_addr1, _max_wake);
    }

    let queue1 = lookup_queue(_addr1);
    let Some(queue1) = queue1 else {
        let cur = read_u32(_addr1)?;
        if cur != _expected_val {
            return Err(SyscallError::Again);
        }
        return Ok(0);
    };

    // Preliminary val check before allocating a bucket slot for addr2.
    // This avoids consuming a slot on the common EAGAIN path (e.g. lock
    // already taken by another thread).  The authoritative check is
    // re-done under the waiter locks below to close the race window.
    let cur_prelim = read_u32(_addr1)?;
    if cur_prelim != _expected_val {
        return Err(SyscallError::Again);
    }

    let queue2 = get_queue(_addr2)?;
    let mut woke = 0u64;
    let mut requeued = 0u64;

    {
        let (mut w1, mut w2) = lock_two_queues(_addr1, &queue1, _addr2, &queue2);
        // Authoritative check: re-read under both waiter locks so no waker
        // can race between the val read and the requeue.
        let cur = read_u32(_addr1)?;
        if cur != _expected_val {
            return Err(SyscallError::Again);
        }
        while woke < _max_wake as u64 {
            if let Some(id) = w1.pop_front() {
                if wake_task(id) {
                    woke += 1;
                }
            } else {
                break;
            }
        }
        while requeued < _max_requeue as u64 {
            if let Some(id) = w1.pop_front() {
                if w2.push_back(id).is_err() {
                    let _ = w1.push_back(id);
                    break;
                }
                requeued += 1;
            } else {
                break;
            }
        }
    }

    try_gc_queue(_addr1, &queue1);
    try_gc_queue(_addr2, &queue2);

    Ok(woke + requeued)
}

/// SYS_FUTEX_WAKE_OP: Wake with atomic operation
pub fn sys_futex_wake_op(
    _addr1: u64,
    _max_wake1: u32,
    _max_wake2: u32,
    _addr2: u64,
    _op: u32,
) -> Result<u64, SyscallError> {
    let addr1 = _addr1;
    let addr2 = _addr2;
    let max_wake1 = _max_wake1;
    let max_wake2 = _max_wake2;
    let wake_op = FutexWakeOpEncode::decode(_op)?;

    // Materialize queues so wake and wait operations serialize on queue locks.
    let q1 = get_queue(addr1)?;
    let q2 = get_queue(addr2)?;

    let woke = if addr1 == addr2 {
        let mut waiters = q1.waiters.lock();
        let old = atomic_fetch_update_u32(addr2, |v| wake_op.calculate_new_val(v))?;
        let mut woke = wake_from_waiters(&mut waiters, max_wake1);
        if wake_op.should_wake(old) {
            woke += wake_from_waiters(&mut waiters, max_wake2);
        }
        woke
    } else {
        let (mut w1, mut w2) = lock_two_queues(addr1, &q1, addr2, &q2);
        let old = atomic_fetch_update_u32(addr2, |v| wake_op.calculate_new_val(v))?;
        let mut woke = wake_from_waiters(&mut w1, max_wake1);
        if wake_op.should_wake(old) {
            woke += wake_from_waiters(&mut w2, max_wake2);
        }
        woke
    };
    try_gc_queue(addr1, &q1);
    try_gc_queue(addr2, &q2);
    Ok(woke)
}