ax_task/sync/
local_lock.rs

1//! CPU-local lock acquisition with migration exclusion before CPU selection.
2
3use alloc::vec::Vec;
4use core::ops::{Deref, DerefMut};
5
6use super::{MigrationGuard, SpinLock, SpinLockGuard};
7use crate::{
8    sched::{CpuId, cpu_topology_len},
9    thread::TaskError,
10};
11
12/// One RT lock and protected value per possible CPU.
13///
14/// Acquisition first pins migration, then selects that CPU's lock. Tasks on
15/// the same CPU may preempt one another and synchronize through PI contention.
16pub struct LocalLock<T> {
17    cpus: Vec<SpinLock<T>>,
18}
19
20/// Local lock ownership; releases the lock before the outer migration pin.
21#[must_use]
22pub struct LocalLockGuard<'a, T> {
23    owner: SpinLockGuard<'a, T>,
24    migration: MigrationGuard,
25}
26
27impl<T> LocalLock<T> {
28    /// Initializes one protected value for every configured CPU.
29    pub fn new(mut init: impl FnMut(CpuId) -> T) -> Result<Self, TaskError> {
30        let cpus = (0..cpu_topology_len()?)
31            .map(|cpu| SpinLock::new(init(CpuId::new(cpu as u32))))
32            .collect();
33        Ok(Self { cpus })
34    }
35
36    /// Acquires the calling task's CPU-local value in preemptible task context.
37    pub fn lock(&self) -> LocalLockGuard<'_, T> {
38        let migration = MigrationGuard::new().expect("local lock requires task context");
39        let owner = self.cpus[migration.cpu().as_usize()].lock();
40        LocalLockGuard { owner, migration }
41    }
42
43    /// Attempts acquisition without sleeping.
44    pub fn try_lock(&self) -> Option<LocalLockGuard<'_, T>> {
45        let migration = MigrationGuard::new().ok()?;
46        let owner = self.cpus[migration.cpu().as_usize()].try_lock()?;
47        Some(LocalLockGuard { owner, migration })
48    }
49
50    /// RT IRQ-save spelling; hardware IRQ state is unchanged.
51    pub fn lock_irqsave(&self) -> LocalLockGuard<'_, T> {
52        self.lock()
53    }
54    /// RT IRQ-save spelling; hardware IRQ state is unchanged.
55    pub fn try_lock_irqsave(&self) -> Option<LocalLockGuard<'_, T>> {
56        self.try_lock()
57    }
58}
59
60impl<T> LocalLockGuard<'_, T> {
61    /// The CPU whose value remains protected across preemption.
62    pub fn cpu(&self) -> CpuId {
63        self.migration.cpu()
64    }
65}
66impl<T> Deref for LocalLockGuard<'_, T> {
67    type Target = T;
68    fn deref(&self) -> &T {
69        &self.owner
70    }
71}
72impl<T> DerefMut for LocalLockGuard<'_, T> {
73    fn deref_mut(&mut self) -> &mut T {
74        &mut self.owner
75    }
76}
ax_task/sync/local_lock.rs

ax_task/sync/
local_lock.rs
