1 // Copyright 2025 The Pigweed Authors 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); you may not 4 // use this file except in compliance with the License. You may obtain a copy of 5 // the License at 6 // 7 // https://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, WITHOUT 11 // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the 12 // License for the specific language governing permissions and limitations under 13 // the License. 14 15 use core::{ 16 cell::UnsafeCell, 17 ops::{Deref, DerefMut}, 18 }; 19 20 use crate::scheduler::{Thread, WaitQueueLock}; 21 22 struct MutexState { 23 count: u32, 24 holder_id: usize, 25 } 26 27 pub struct Mutex<T> { 28 // An future optimization can be made by keeping an atomic count outside of 29 // the spinlock. However, not all architectures support atomics so a pure 30 // SchedLock based approach will always be needed. 31 state: WaitQueueLock<MutexState>, 32 33 inner: UnsafeCell<T>, 34 } 35 unsafe impl<T> Sync for Mutex<T> {} 36 unsafe impl<T> Send for Mutex<T> {} 37 38 pub struct MutexGuard<'lock, T> { 39 lock: &'lock Mutex<T>, 40 } 41 42 impl<T> Deref for MutexGuard<'_, T> { 43 type Target = T; 44 deref(&self) -> &T45 fn deref(&self) -> &T { 46 unsafe { &*self.lock.inner.get() } 47 } 48 } 49 50 impl<T> DerefMut for MutexGuard<'_, T> { deref_mut(&mut self) -> &mut T51 fn deref_mut(&mut self) -> &mut T { 52 unsafe { &mut *self.lock.inner.get() } 53 } 54 } 55 56 impl<T> Drop for MutexGuard<'_, T> { drop(&mut self)57 fn drop(&mut self) { 58 self.lock.unlock(); 59 } 60 } 61 62 impl<T> Mutex<T> { new(initial_value: T) -> Self63 pub const fn new(initial_value: T) -> Self { 64 Self { 65 state: WaitQueueLock::new(MutexState { 66 count: 0, 67 holder_id: Thread::null_id(), 68 }), 69 inner: UnsafeCell::new(initial_value), 70 } 71 } 72 lock(&self) -> MutexGuard<'_, T>73 pub fn lock(&self) -> MutexGuard<'_, T> { 74 let mut state = self.state.lock(); 75 state.count += 1; 76 // TODO - konkers: investigate using core::intrinsics::unlikely() or 77 // core::hint::unlikely() 78 let mut state = if state.count > 1 { 79 state.into_wait_queue().wait(); 80 // `wait()` consumes the state lock so re-acquire it before moving on. 81 self.state.lock() 82 } else { 83 state 84 }; 85 86 state.holder_id = state.sched().current_thread_id(); 87 88 // At this point we have exclusive access to `self.inner`. 89 90 MutexGuard { lock: self } 91 } 92 unlock(&self)93 fn unlock(&self) { 94 let mut state = self.state.lock(); 95 96 assert!(state.count > 0); 97 assert_eq!(state.holder_id, state.sched().current_thread_id()); 98 state.holder_id = Thread::null_id(); 99 100 state.count -= 1; 101 102 // TODO - konkers: investigate using core::intrinsics::unlikely() or 103 // core::hint::unlikely() 104 if state.count > 0 { 105 state.into_wait_queue().wake_one(); 106 } 107 } 108 } 109