1use core::{
4 cell::{RefCell, UnsafeCell},
5 fmt,
6 mem::MaybeUninit,
7 ptr::NonNull,
8 sync::atomic::{AtomicU8, AtomicU64, Ordering},
9};
10
11use super::entry::{FastReleaseAttempt, try_release_current_owner_word};
12use crate::thread::{ParkTicket, ThreadHandle};
13
14static NEXT_PI_MUTEX_GENERATION: AtomicU64 = AtomicU64::new(1);
15const OWNER_HAS_WAITERS: u64 = 1 << 63;
16const OWNER_ID_MASK: u64 = !OWNER_HAS_WAITERS;
17const WAIT_STORAGE_UNINITIALIZED: u8 = 0;
18const WAIT_STORAGE_INITIALIZING: u8 = 1;
19const WAIT_STORAGE_READY: u8 = 2;
20
21#[doc(hidden)]
27pub const PI_MUTEX_WAIT_STORAGE_WORDS: usize = 5;
28
29pub struct PiMutexWaitStorage {
31 state: AtomicU8,
32 words: UnsafeCell<[MaybeUninit<usize>; PI_MUTEX_WAIT_STORAGE_WORDS]>,
33}
34
35impl PiMutexWaitStorage {
36 const fn new() -> Self {
37 Self {
38 state: AtomicU8::new(WAIT_STORAGE_UNINITIALIZED),
39 words: UnsafeCell::new([MaybeUninit::uninit(); PI_MUTEX_WAIT_STORAGE_WORDS]),
40 }
41 }
42}
43
44#[derive(Clone, Copy, Debug)]
49pub struct PiMutexWaitStorageView<'lock> {
50 state: &'lock AtomicU8,
51 words: &'lock UnsafeCell<[MaybeUninit<usize>; PI_MUTEX_WAIT_STORAGE_WORDS]>,
52}
53
54impl<'lock> PiMutexWaitStorageView<'lock> {
55 #[doc(hidden)]
57 const fn from_parts(
58 state: &'lock AtomicU8,
59 words: &'lock UnsafeCell<[MaybeUninit<usize>; PI_MUTEX_WAIT_STORAGE_WORDS]>,
60 ) -> Self {
61 Self { state, words }
62 }
63
64 #[doc(hidden)]
66 pub const fn as_ptr(self) -> *mut () {
67 self.words.get().cast()
68 }
69
70 #[doc(hidden)]
72 pub fn is_initialized(self) -> bool {
73 self.state.load(Ordering::Acquire) == WAIT_STORAGE_READY
74 }
75
76 #[doc(hidden)]
84 pub unsafe fn get_or_init<T>(self, init: impl FnOnce() -> T) -> &'lock T {
85 assert!(
86 core::mem::size_of::<T>()
87 <= PI_MUTEX_WAIT_STORAGE_WORDS * core::mem::size_of::<usize>(),
88 "PI mutex provider waiter state exceeds inline storage"
89 );
90 assert!(
91 core::mem::align_of::<T>() <= core::mem::align_of::<usize>(),
92 "PI mutex provider waiter state exceeds inline alignment"
93 );
94
95 if self
96 .state
97 .compare_exchange(
98 WAIT_STORAGE_UNINITIALIZED,
99 WAIT_STORAGE_INITIALIZING,
100 Ordering::Acquire,
101 Ordering::Acquire,
102 )
103 .is_ok()
104 {
105 unsafe { self.as_ptr().cast::<T>().write(init()) };
108 self.state.store(WAIT_STORAGE_READY, Ordering::Release);
109 } else {
110 while self.state.load(Ordering::Acquire) == WAIT_STORAGE_INITIALIZING {
111 core::hint::spin_loop();
112 }
113 assert_eq!(
114 self.state.load(Ordering::Acquire),
115 WAIT_STORAGE_READY,
116 "PI mutex waiter storage has an invalid lifecycle"
117 );
118 }
119
120 unsafe { &*self.as_ptr().cast::<T>() }
123 }
124
125 #[doc(hidden)]
131 pub unsafe fn get<T>(self) -> Option<&'lock T> {
132 if self.state.load(Ordering::Acquire) != WAIT_STORAGE_READY {
133 return None;
134 }
135 Some(unsafe { &*self.as_ptr().cast::<T>() })
137 }
138}
139
140fn take_initialized_wait_storage(
141 state: &mut u8,
142 words: &mut [MaybeUninit<usize>; PI_MUTEX_WAIT_STORAGE_WORDS],
143) -> Option<*mut ()> {
144 match *state {
145 WAIT_STORAGE_UNINITIALIZED => None,
146 WAIT_STORAGE_READY => {
147 *state = WAIT_STORAGE_UNINITIALIZED;
148 Some(words.as_mut_ptr().cast())
149 }
150 _ => panic!("destroying PI mutex while waiter storage initializes"),
151 }
152}
153
154unsafe impl Sync for PiMutexWaitStorage {}
157
158#[repr(transparent)]
160#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
161pub struct PiTaskId(u64);
162
163impl PiTaskId {
164 pub const fn new(raw: u64) -> Option<Self> {
166 if raw == 0 || raw & OWNER_HAS_WAITERS != 0 {
167 None
168 } else {
169 Some(Self(raw))
170 }
171 }
172
173 pub const fn get(self) -> u64 {
175 self.0
176 }
177}
178
179#[repr(transparent)]
181#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
182pub struct PiMutexId(u64);
183
184impl PiMutexId {
185 pub const fn get(self) -> u64 {
187 self.0
188 }
189}
190
191#[derive(Clone, Copy, Debug, Eq, PartialEq)]
193pub enum PiMutexStateError {
194 WaiterOwnsLock,
196 InvalidState,
198}
199
200impl fmt::Display for PiMutexStateError {
201 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
202 formatter.write_str(match self {
203 Self::WaiterOwnsLock => "PI mutex waiter already owns the lock",
204 Self::InvalidState => "invalid PI mutex state",
205 })
206 }
207}
208
209impl core::error::Error for PiMutexStateError {}
210
211pub struct PiMutexCore {
217 owner: AtomicU64,
218 generation: AtomicU64,
219 wait_storage: PiMutexWaitStorage,
220}
221
222#[derive(Clone, Copy, Debug)]
228pub struct PiMutexCoreView<'lock> {
229 owner: &'lock AtomicU64,
230 generation: &'lock AtomicU64,
231 wait_storage: PiMutexWaitStorageView<'lock>,
232}
233
234impl<'lock> PiMutexCoreView<'lock> {
235 #[doc(hidden)]
237 pub(in crate::sync) const fn from_parts(
238 owner: &'lock AtomicU64,
239 generation: &'lock AtomicU64,
240 wait_state: &'lock AtomicU8,
241 wait_words: &'lock UnsafeCell<[MaybeUninit<usize>; PI_MUTEX_WAIT_STORAGE_WORDS]>,
242 ) -> Self {
243 Self {
244 owner,
245 generation,
246 wait_storage: PiMutexWaitStorageView::from_parts(wait_state, wait_words),
247 }
248 }
249
250 pub fn try_acquire(self, current: PiTaskId) -> Result<PiMutexAcquire, PiMutexStateError> {
252 match self
253 .owner
254 .compare_exchange(0, current.get(), Ordering::Acquire, Ordering::Relaxed)
255 {
256 Ok(_) => Ok(PiMutexAcquire::Acquired),
257 Err(owner) if owner & OWNER_ID_MASK == current.get() => {
258 Err(PiMutexStateError::WaiterOwnsLock)
259 }
260 Err(_) => Ok(PiMutexAcquire::Contended),
261 }
262 }
263
264 #[doc(hidden)]
271 pub unsafe fn try_acquire_for_thread<T>(
272 self,
273 current: T,
274 ) -> Result<PiMutexAcquire, PiMutexStateError>
275 where
276 T: Into<PiTaskId>,
277 {
278 self.try_acquire(current.into())
279 }
280
281 #[doc(hidden)]
288 pub unsafe fn try_release_for_thread<T>(self, current: T) -> Result<bool, PiMutexStateError>
289 where
290 T: Into<PiTaskId>,
291 {
292 let current = current.into();
293 match try_release_current_owner_word(self.owner, current.get(), OWNER_ID_MASK) {
294 FastReleaseAttempt::Released => Ok(true),
295 FastReleaseAttempt::Contended => Ok(false),
296 FastReleaseAttempt::InvalidOwner => Err(PiMutexStateError::InvalidState),
297 }
298 }
299
300 pub unsafe fn try_release_owned(
308 self,
309 current: PiTaskId,
310 ) -> Result<PiMutexOwnedRelease, PiMutexStateError> {
311 match try_release_current_owner_word(self.owner, current.get(), OWNER_ID_MASK) {
312 FastReleaseAttempt::Released => Ok(PiMutexOwnedRelease::Released),
313 FastReleaseAttempt::Contended => Ok(PiMutexOwnedRelease::Contended(current)),
314 FastReleaseAttempt::InvalidOwner => Err(PiMutexStateError::InvalidState),
315 }
316 }
317
318 pub fn is_owned_by(self, current: PiTaskId) -> bool {
320 owner_from_word(self.owner.load(Ordering::Acquire)) == Some(current)
321 }
322
323 pub fn is_locked(self) -> bool {
325 self.owner.load(Ordering::Relaxed) != 0
326 }
327
328 pub fn mutex_ref(self) -> Result<PiMutexRef<'lock>, PiMutexStateError> {
330 let observed = self.generation.load(Ordering::Acquire);
331 if observed != 0 {
332 return Ok(PiMutexRef {
333 core: self,
334 id: PiMutexId(observed),
335 });
336 }
337
338 let allocated = NEXT_PI_MUTEX_GENERATION
339 .try_update(Ordering::AcqRel, Ordering::Acquire, |next| {
340 next.checked_add(1)
341 })
342 .map(PiMutexId)
343 .map_err(|_| PiMutexStateError::InvalidState)?;
344 match self
345 .generation
346 .compare_exchange(0, allocated.0, Ordering::AcqRel, Ordering::Acquire)
347 {
348 Ok(_) => Ok(PiMutexRef {
349 core: self,
350 id: allocated,
351 }),
352 Err(installed) if installed != 0 => Ok(PiMutexRef {
353 core: self,
354 id: PiMutexId(installed),
355 }),
356 Err(_) => Err(PiMutexStateError::InvalidState),
357 }
358 }
359
360 #[doc(hidden)]
362 pub fn owner_snapshot(self) -> PiMutexOwnerSnapshot {
363 let word = self.owner.load(Ordering::Acquire);
364 PiMutexOwnerSnapshot {
365 word,
366 owner: owner_from_word(word),
367 }
368 }
369
370 #[doc(hidden)]
372 pub fn try_acquire_snapshot(self, snapshot: PiMutexOwnerSnapshot, current: PiTaskId) -> bool {
373 debug_assert_eq!(snapshot.word, 0);
374 self.owner
375 .compare_exchange(
376 snapshot.word,
377 current.get(),
378 Ordering::Acquire,
379 Ordering::Relaxed,
380 )
381 .is_ok()
382 }
383
384 #[doc(hidden)]
386 pub fn try_mark_waiters(self, snapshot: PiMutexOwnerSnapshot) -> bool {
387 if snapshot.has_waiters() {
388 return self.owner.load(Ordering::Acquire) == snapshot.word;
389 }
390 self.owner
391 .compare_exchange(
392 snapshot.word,
393 snapshot.word | OWNER_HAS_WAITERS,
394 Ordering::AcqRel,
395 Ordering::Acquire,
396 )
397 .is_ok()
398 }
399
400 #[doc(hidden)]
402 pub fn publish_owner(self, owner: PiTaskId, has_waiters: bool) {
403 self.owner.store(
404 owner.get() | if has_waiters { OWNER_HAS_WAITERS } else { 0 },
405 Ordering::Release,
406 );
407 }
408
409 #[doc(hidden)]
411 pub fn publish_ownerless(self) {
412 self.owner.store(OWNER_HAS_WAITERS, Ordering::Release);
413 }
414
415 #[doc(hidden)]
417 pub fn publish_unlocked(self) {
418 self.owner.store(0, Ordering::Release);
419 }
420
421 #[doc(hidden)]
423 pub fn clear_waiters_bit(self, owner: PiTaskId) {
424 self.owner.store(owner.get(), Ordering::Release);
425 }
426
427 #[doc(hidden)]
429 pub const fn wait_storage(self) -> PiMutexWaitStorageView<'lock> {
430 self.wait_storage
431 }
432}
433
434impl PiMutexCore {
435 pub const fn new() -> Self {
437 Self {
438 owner: AtomicU64::new(0),
439 generation: AtomicU64::new(0),
440 wait_storage: PiMutexWaitStorage::new(),
441 }
442 }
443
444 #[doc(hidden)]
446 pub const fn view(&self) -> PiMutexCoreView<'_> {
447 PiMutexCoreView::from_parts(
448 &self.owner,
449 &self.generation,
450 &self.wait_storage.state,
451 &self.wait_storage.words,
452 )
453 }
454}
455
456impl fmt::Debug for PiMutexCore {
457 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
458 formatter
459 .debug_struct("PiMutexCore")
460 .field(
461 "owner",
462 &owner_from_word(self.owner.load(Ordering::Relaxed)),
463 )
464 .field("generation", &self.generation.load(Ordering::Relaxed))
465 .finish_non_exhaustive()
466 }
467}
468
469impl Default for PiMutexCore {
470 fn default() -> Self {
471 Self::new()
472 }
473}
474
475impl Drop for PiMutexCore {
476 fn drop(&mut self) {
477 destroy_pi_mutex_storage(
478 &mut self.owner,
479 &mut self.generation,
480 &mut self.wait_storage.state,
481 &mut self.wait_storage.words,
482 );
483 }
484}
485
486pub(in crate::sync) fn destroy_pi_mutex_storage(
487 owner: &mut AtomicU64,
488 generation: &mut AtomicU64,
489 wait_state: &mut AtomicU8,
490 wait_words: &mut UnsafeCell<[MaybeUninit<usize>; PI_MUTEX_WAIT_STORAGE_WORDS]>,
491) {
492 *owner.get_mut() = 0;
493 *generation.get_mut() = 0;
494 if let Some(wait_handle) =
495 take_initialized_wait_storage(wait_state.get_mut(), wait_words.get_mut())
496 {
497 unsafe { crate::thread::drop_pi_mutex_wait_handle(wait_handle) };
501 }
502}
503
504#[derive(Clone, Copy, Debug)]
506pub struct PiMutexRef<'lock> {
507 core: PiMutexCoreView<'lock>,
508 id: PiMutexId,
509}
510
511impl<'lock> PiMutexRef<'lock> {
512 pub const fn id(self) -> PiMutexId {
514 self.id
515 }
516
517 #[doc(hidden)]
519 pub const fn core(self) -> PiMutexCoreView<'lock> {
520 self.core
521 }
522
523 #[doc(hidden)]
525 pub fn raw(self) -> PiMutexRaw {
526 PiMutexRaw {
527 owner: NonNull::from(self.core.owner),
528 generation: NonNull::from(self.core.generation),
529 wait_state: NonNull::from(self.core.wait_storage.state),
530 wait_words: NonNull::from(self.core.wait_storage.words),
531 id: self.id,
532 }
533 }
534}
535
536#[derive(Clone, Copy, Debug, Eq, PartialEq)]
538pub struct PiMutexRaw {
539 owner: NonNull<AtomicU64>,
540 generation: NonNull<AtomicU64>,
541 wait_state: NonNull<AtomicU8>,
542 wait_words: NonNull<UnsafeCell<[MaybeUninit<usize>; PI_MUTEX_WAIT_STORAGE_WORDS]>>,
543 id: PiMutexId,
544}
545
546impl PiMutexRaw {
547 pub const fn id(self) -> PiMutexId {
549 self.id
550 }
551
552 #[doc(hidden)]
559 pub unsafe fn core(self) -> PiMutexCoreView<'static> {
560 PiMutexCoreView {
561 owner: unsafe { self.owner.as_ref() },
564 generation: unsafe { self.generation.as_ref() },
566 wait_storage: PiMutexWaitStorageView {
567 state: unsafe { self.wait_state.as_ref() },
569 words: unsafe { self.wait_words.as_ref() },
571 },
572 }
573 }
574}
575
576unsafe impl Send for PiMutexRaw {}
579unsafe impl Sync for PiMutexRaw {}
580
581#[derive(Clone, Copy, Debug, Eq, PartialEq)]
583pub struct PiMutexOwnerSnapshot {
584 word: u64,
585 owner: Option<PiTaskId>,
586}
587
588impl PiMutexOwnerSnapshot {
589 pub const fn owner(self) -> Option<PiTaskId> {
591 self.owner
592 }
593
594 pub const fn is_unlocked(self) -> bool {
596 self.word == 0
597 }
598
599 pub const fn is_ownerless(self) -> bool {
601 self.word == OWNER_HAS_WAITERS
602 }
603
604 pub const fn has_waiters(self) -> bool {
606 self.word & OWNER_HAS_WAITERS != 0
607 }
608}
609
610#[derive(Clone, Copy, Debug, Eq, PartialEq)]
612pub enum PiMutexAcquire {
613 Acquired,
615 Contended,
617}
618
619#[derive(Clone, Copy, Debug, Eq, PartialEq)]
621pub enum PiMutexOwnedRelease {
622 Released,
624 Contended(PiTaskId),
626}
627
628#[must_use = "a PI wait token must be granted or explicitly cancelled"]
630#[derive(Debug)]
631pub struct PiWaitToken {
632 thread: PiTaskId,
633 initial_owner: Option<ThreadHandle>,
634 generation: u64,
635 lock: PiMutexRaw,
636 provider_waiter: NonNull<()>,
637 prepared_park: RefCell<Option<ParkTicket>>,
638}
639
640impl PiWaitToken {
641 #[doc(hidden)]
649 pub unsafe fn from_registration(
650 lock: PiMutexRaw,
651 thread: PiTaskId,
652 initial_owner: Option<ThreadHandle>,
653 generation: u64,
654 provider_waiter: NonNull<()>,
655 ) -> Self {
656 Self {
657 thread,
658 initial_owner,
659 generation,
660 lock,
661 provider_waiter,
662 prepared_park: RefCell::new(None),
663 }
664 }
665
666 pub(crate) fn install_prepared_park(&self, ticket: ParkTicket) {
667 assert_eq!(
668 ticket.thread().as_u64(),
669 self.thread.get(),
670 "PI waiter park ticket must belong to the registered task"
671 );
672 assert!(
673 self.prepared_park.replace(Some(ticket)).is_none(),
674 "PI waiter may own only one prepared park"
675 );
676 }
677
678 pub(crate) fn take_prepared_park(&self) -> Option<ParkTicket> {
679 self.prepared_park.take()
680 }
681
682 pub const fn thread_id(&self) -> PiTaskId {
684 self.thread
685 }
686
687 pub fn initial_owner(&self) -> Option<PiTaskId> {
689 self.initial_owner
690 .as_ref()
691 .map(|owner| PiTaskId::from(owner.id()))
692 }
693
694 #[doc(hidden)]
696 pub(crate) fn initial_owner_handle(&self) -> Option<&ThreadHandle> {
697 self.initial_owner.as_ref()
698 }
699
700 #[doc(hidden)]
702 pub const fn generation(&self) -> u64 {
703 self.generation
704 }
705
706 #[doc(hidden)]
708 pub const fn lock_raw(&self) -> PiMutexRaw {
709 self.lock
710 }
711
712 #[doc(hidden)]
719 pub const unsafe fn provider_waiter(&self) -> NonNull<()> {
720 self.provider_waiter
721 }
722
723 pub fn is_granted(&self) -> bool {
725 crate::runtime::sync::pi_waiter_is_granted(self)
726 }
727
728 pub fn can_claim(&self) -> bool {
730 self.is_top_waiter() && unsafe { self.lock.core() }.owner_snapshot().is_ownerless()
731 }
732
733 pub fn is_top_waiter(&self) -> bool {
735 crate::runtime::sync::pi_waiter_is_top(self)
736 }
737
738 pub fn initial_owner_is_on_cpu(&self) -> bool {
740 super::task_result(
741 crate::runtime::sync::pi_initial_owner_is_on_cpu(self),
742 "observe PI mutex owner execution state",
743 )
744 }
745}
746
747#[must_use = "a registered PI waiter must be blocked, claimed, or cancelled"]
749#[derive(Debug)]
750pub enum PiMutexLockResult {
751 Acquired,
753 Waiting(PiWaitToken),
755}
756
757#[derive(Clone, Copy, Debug, Eq, PartialEq)]
759pub enum PiMutexClaimOutcome {
760 Claimed,
762 Retry,
764}
765
766#[derive(Clone, Copy, Debug, Eq, PartialEq)]
768pub enum PiWaitCancelOutcome {
769 Cancelled,
771 HandoffPending,
773}
774
775fn owner_from_word(state: u64) -> Option<PiTaskId> {
776 PiTaskId::new(state & OWNER_ID_MASK)
777}
778
779#[cfg(test)]
780mod tests;