ax_task/runtime/sync/
pi.rs

1use alloc::sync::Arc;
2
3use crate::{
4    runtime::{
5        TaskSystem,
6        context::{RuntimeIrqGuard, runtime_current_cpu_mut, runtime_task_system},
7        sync::{
8            PiMutexClaimOutcome, PiMutexLockResult, PiMutexRef, PiWaitCancelOutcome, PiWaitToken,
9        },
10    },
11    thread::{
12        CurrentThreadToken, ParkPrepare, TaskError, ThreadCore, ThreadId,
13        current::{
14            acquire_blocking_permit, current_thread_core_arc,
15            park::{cancel_current_park, commit_current_park},
16        },
17    },
18};
19
20pub(crate) enum PiParkAttempt {
21    Complete,
22    Retry,
23    Prepared(Arc<ThreadCore>, crate::thread::ParkTicket),
24}
25
26/// Enters the scheduler-owned PI mutex slow path.
27pub fn pi_mutex_lock_slow(
28    lock: PiMutexRef<'_>,
29    current: &CurrentThreadToken,
30    sequence: u64,
31) -> Result<PiMutexLockResult, TaskError> {
32    let _permit = acquire_blocking_permit()?;
33    let current_core = current_thread_core_arc()?;
34    if current_core.id() != current.id() {
35        return Err(TaskError::InvalidPiState);
36    }
37    let system = runtime_task_system()?;
38    let mut irq = RuntimeIrqGuard::enter();
39    let mut cpu = runtime_current_cpu_mut(&mut irq)?;
40    system.drain_owner_control(cpu.as_mut())?;
41    let mut park = loop {
42        match system.prepare_current_park(&current_core)? {
43            ParkPrepare::Notified => continue,
44            ParkPrepare::Prepared(ticket) => break ticket,
45        }
46    };
47    match system.pi_mutex_lock_slow(lock, current.id(), sequence) {
48        Ok(PiMutexLockResult::Acquired) => {
49            system.cancel_current_park(cpu.as_mut(), &current_core, &mut park)?;
50            Ok(PiMutexLockResult::Acquired)
51        }
52        Ok(PiMutexLockResult::Waiting(token)) => {
53            token.install_prepared_park(park);
54            Ok(PiMutexLockResult::Waiting(token))
55        }
56        Err(error) => {
57            system.cancel_current_park(cpu.as_mut(), &current_core, &mut park)?;
58            Err(error)
59        }
60    }
61}
62
63/// Performs one scheduler park attempt for a PI waiter.
64///
65/// The caller must recheck ownership, interruption, and timeout after this
66/// function returns. This mirrors Linux `rt_mutex_schedule()`: an unrelated
67/// wake returns control to the rtmutex state loop instead of being consumed by
68/// an inner uninterruptible wait.
69pub fn pi_park_current_once(token: &PiWaitToken) -> Result<(), TaskError> {
70    let system = runtime_task_system()?;
71    let prepared = if let Some(ticket) = token.take_prepared_park() {
72        let current = current_thread_core_arc()?;
73        if current.id() != token.thread_id().into() || ticket.thread() != current.id() {
74            return Err(TaskError::InvalidPiState);
75        }
76        PiParkAttempt::Prepared(current, ticket)
77    } else {
78        if token.can_claim() || token.is_granted() {
79            return Ok(());
80        }
81        prepare_pi_park_attempt(system, token)?
82    };
83    let (current, mut ticket) = match prepared {
84        PiParkAttempt::Complete | PiParkAttempt::Retry => return Ok(()),
85        PiParkAttempt::Prepared(current, ticket) => (current, ticket),
86    };
87    if token.can_claim() || token.is_granted() {
88        cancel_current_park(&current, &mut ticket)?;
89        return Ok(());
90    }
91    commit_current_park(&current, &mut ticket).map(|_| ())
92}
93
94pub(crate) fn cancel_prepared_pi_park(token: &PiWaitToken) -> Result<(), TaskError> {
95    let Some(mut ticket) = token.take_prepared_park() else {
96        return Ok(());
97    };
98    let current = current_thread_core_arc()?;
99    if current.id() != token.thread_id().into() || ticket.thread() != current.id() {
100        return Err(TaskError::InvalidPiState);
101    }
102    cancel_current_park(&current, &mut ticket)
103}
104
105pub(crate) fn prepare_pi_park_attempt(
106    system: &TaskSystem,
107    token: &PiWaitToken,
108) -> Result<PiParkAttempt, TaskError> {
109    let _permit = acquire_blocking_permit()?;
110    let current = current_thread_core_arc()?;
111    if current.id() != token.thread_id().into() {
112        return Err(TaskError::InvalidPiState);
113    }
114    let mut irq = RuntimeIrqGuard::enter();
115    let mut cpu = runtime_current_cpu_mut(&mut irq)?;
116    system.drain_owner_control(cpu.as_mut())?;
117    if token.can_claim() || token.is_granted() {
118        return Ok(PiParkAttempt::Complete);
119    }
120    match system.prepare_current_park(&current)? {
121        ParkPrepare::Notified => Ok(PiParkAttempt::Retry),
122        ParkPrepare::Prepared(ticket) => Ok(PiParkAttempt::Prepared(current, ticket)),
123    }
124}
125
126/// Cancels a PI wait token after a handoff-before-block race.
127pub fn pi_wait_cancel(token: PiWaitToken) -> Result<(), TaskError> {
128    let outcome = runtime_task_system()?.pi_wait_try_cancel(&token)?;
129    cancel_prepared_pi_park(&token)?;
130    match outcome {
131        PiWaitCancelOutcome::Cancelled => Ok(()),
132        PiWaitCancelOutcome::HandoffPending => Err(TaskError::InvalidPiState),
133    }
134}
135
136/// Tries to cancel one PI waiter while preserving an ownerless handoff that
137/// already selected it.
138pub fn pi_wait_try_cancel(token: &PiWaitToken) -> Result<PiWaitCancelOutcome, TaskError> {
139    runtime_task_system()?.pi_wait_try_cancel(token)
140}
141
142/// Publishes a raw-mutex-owner PI handoff and wakes the selected waiter.
143///
144/// # Safety
145///
146/// `old_owner` must be the executing identity passed to
147/// [`PiMutexCoreView::try_release_owned`] on `lock`, and the caller must retain the
148/// higher-level raw-mutex owner authority until this complete release
149/// transaction returns.
150pub unsafe fn pi_mutex_release_owned(
151    lock: PiMutexRef<'_>,
152    old_owner: ThreadId,
153) -> Result<(), TaskError> {
154    runtime_task_system()?.pi_mutex_release(lock, old_owner)
155}
156
157/// Claims the ownerless PI mutex handoff selected for this waiter.
158pub fn pi_mutex_claim(
159    token: &PiWaitToken,
160    current: &CurrentThreadToken,
161) -> Result<PiMutexClaimOutcome, TaskError> {
162    if current.id() != token.thread_id().into() {
163        return Err(TaskError::InvalidPiState);
164    }
165    runtime_task_system()?.pi_mutex_claim(token)
166}
167
168/// Tests whether the task-local waiter capability has received handoff.
169pub fn pi_waiter_is_granted(token: &PiWaitToken) -> bool {
170    let waiter = unsafe {
171        // SAFETY: the task-system registration retains this task-local wait
172        // state until the token is claimed or cancelled.
173        token
174            .provider_waiter()
175            .cast::<crate::thread::PiWaitState>()
176            .as_ref()
177    };
178    waiter.is_granted(token.generation())
179}
180
181/// Tests whether the task-local waiter capability is first in its lock queue.
182pub fn pi_waiter_is_top(token: &PiWaitToken) -> bool {
183    let waiter = unsafe {
184        // SAFETY: identical provider capability contract to
185        // `pi_waiter_is_granted` above.
186        token
187            .provider_waiter()
188            .cast::<crate::thread::PiWaitState>()
189            .as_ref()
190    };
191    waiter.is_top(token.generation())
192}
193
194/// Tests whether the waiter token's initial owner is still executing.
195pub fn pi_initial_owner_is_on_cpu(token: &PiWaitToken) -> Result<bool, TaskError> {
196    runtime_task_system()?.pi_initial_owner_is_on_cpu(token)
197}
198
199/// Drops the scheduler-owned waiter handle transferred by a physical PI mutex.
200///
201/// # Safety
202///
203/// `wait_handle` must be the unique initialized inline handle transferred by
204/// `PiMutexCore` after every safe lock reference and waiter became unreachable.
205pub unsafe fn pi_drop_wait_handle(wait_handle: *mut ()) {
206    unsafe { crate::thread::drop_pi_mutex_wait_handle(wait_handle) };
207}
ax_task/runtime/sync/pi.rs

ax_task/runtime/sync/
pi.rs
