Implement a More Ergonomic & Efficient Mutex for the Embedded

[Congyuwang]

Mutexes are mostly used for static variables, and also a lot of times the variable needs to be initialized later.

The current Mutex implementation doesn't feel very ergonomic:

1. to allow mutable access + late initialization to a variable, we need to use Mutex<RefCell<Option<T>>> with var.borrow(cs).borrow_mut().replace(value) for initialization and var.borrow(cs).borrow_mut().as_mut().unwrap() for mutable access. These are very long expressions for simple purposes.
2. RefCell comes with an extra isize 32bit space with Option adding another u8 4bit overhead. Each time accessing a mutable reference to a variable basically involves two unwraps, once for checking uniqueness and once for checking initialization, which is unnecessary.

I am thinking of an easier to use + more efficient implementation:

pub struct Mutex<T>(UnsafeCell<MutexInner<T>>);

struct MutexInner<T> {
    state: MutexInnerState,
    value: MaybeUninit<T>,
}

pub struct LockGuard<'cs, T>(&'cs mut MutexInner<T>);

We can use a single enum to keep track of the state of the mutex cell. So the value can be

• uninitialized (None),
• locked (borrowed),
• unlock (free to borrow).

enum MutexInnerState {
    Locked,
    Uinit,
    Unlock,
}

The mutex can be either initialized with a value or not.

impl<T> Mutex<T> {
    /// Creates a new mutex.
    pub const fn new(value: T) -> Self {
        Self(UnsafeCell::new(MutexInner {
            state: MutexInnerState::Unlock,
            value: MaybeUninit::new(value),
        }))
    }

    /// Creates a new unit mutex.
    pub const fn new_uinit() -> Self {
        Self(UnsafeCell::new(MutexInner {
            state: MutexInnerState::Uinit,
            value: MaybeUninit::uninit(),
        }))
    }
}

Value can be initialized once if it was not initialized (otherwise panic).

impl<T> Mutex<T> {
    /// Value initialization.
    ///
    /// panic if already initialized.
    pub fn init<'cs>(&'cs self, _cs: &'cs CriticalSection, value: T) {
        let inner = unsafe { &mut *self.0.get() };
        if let MutexInnerState::Uinit = inner.state {
            inner.state = MutexInnerState::Unlock;
            inner.value = MaybeUninit::new(value);
        } else {
            panic!()
        }
    }
}

Locking the mutex returns None if try_lock fails or if the value is uninitialized.

impl<T> Mutex<T> {
    /// Try to lock the mutex.
    pub fn try_lock<'cs>(&'cs self, _cs: &'cs CriticalSection) -> Option<LockGuard<'cs, T>> {
        let inner = unsafe { &mut *self.0.get() };
        match inner.state {
            MutexInnerState::Uinit | MutexInnerState::Locked => None,
            MutexInnerState::Unlock => {
                inner.state = MutexInnerState::Locked;
                Some(LockGuard(inner))
            }
        }
    }
}

The LockGuard restore lock state back to Unlock on drop.

impl<'cs, T> Drop for LockGuard<'cs, T> {
    fn drop(&mut self) {
        self.0.state = MutexInnerState::Unlock;
    }
}

Miscellaneous.

impl<T> Drop for Mutex<T> {
    fn drop(&mut self) {
        let inner = unsafe { &mut *self.0.get() };
        if let MutexInnerState::Unlock | MutexInnerState::Locked = inner.state {
            unsafe { inner.value.assume_init_drop() }
        }
    }
}

impl<'cs, T> Deref for LockGuard<'cs, T> {
    type Target = T;

    #[inline]
    fn deref(&self) -> &Self::Target {
        unsafe { self.0.value.assume_init_ref() }
    }
}

impl<'cs, T> DerefMut for LockGuard<'cs, T> {
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { self.0.value.assume_init_mut() }
    }
}

unsafe impl<T> Sync for Mutex<T> where T: Send {}

Usage

Now we can write like

// compared to `Mutex <RefCell<Option<Rtc<RTC0>>>>`
static RTC: Mutex<Rtc<RTC0>> = Mutex::new_uinit();

// initialization
cortex_m::interrupt::free(|cs| RTC.init(cs, rtc0));

// access
#[interrupt]
fn RTC0() {
    cortex_m::interrupt::free(|cs| {
        if let (Some(mut rtc), Some(mut other)) = (RTC.try_lock(cs), OTHER.try_lock(cs)) {
            other.handle_interrupt(&mut rtc);
        }
    });
}

Of course we should not use the name Mutex directly since we need to have backward compatibility. I am thinking of naming it MutexCell or perhaps MutexOption. If anyone finds it useful maybe we can PR it into the library.

[Congyuwang]

Related: rust-embedded/cortex-m#224, rust-embedded/cortex-m#359

[skibon02]

When we want mutable access to the value behind mutex, essentially it boils down to two of the most common scenarios:

1. Simple mutable access from different parts of the program. Most frequently, we want to share data between peripheral interrupt and other thread mode code.
2. Initialize first, and only read later.

Here are some types I've come up with in my code:

1. MutexMut<T>

First case is already implemented in critical_section and covered by Mutex<RefCell> implementation, but we can make it more user-friendly.
This is how I use it in my code:

static VAL_MUTEX: MutexMut<Vec<u8>> = MutexMut::new_refcell(Vec::new());

// ...Somewhere in code
critical_section::with(|cs| {
    VAL_MUTEX.borrow_ref_mut(cs).push(13);
});

Implementation:

/// Re-export simple mutexes
pub use critical_section::{CriticalSection, Mutex};

/// Critical section mutex with dynamic borrowing (refcell inside)
pub struct MutexMut<T>(Mutex<RefCell<T>>);

impl<'cs, T> MutexMut<T> {
    pub const fn new_refcell(v: T) -> Self {
        MutexMut(Mutex::new(RefCell::new(v)))
    }

    #[inline]
    #[track_caller]
    pub fn replace(&'cs self, cs: CriticalSection<'cs>, v: T) -> T {
        self.0.replace(cs, v)
    }

    #[inline]
    #[track_caller]
    pub fn borrow_ref(&'cs self, cs: CriticalSection<'cs>) -> Ref<'cs, T> {
        self.0.borrow_ref(cs)
    }
    #[inline]
    #[track_caller]
    pub fn borrow_ref_mut(&'cs self, cs: CriticalSection<'cs>) -> RefMut<'cs, T> {
        self.0.borrow_ref_mut(cs)
    }
}

// Deref used because we can't add method to foreign type -> use your own transparent wrapper
impl<T> Deref for MutexMut<T> {
    type Target = Mutex<RefCell<T>>;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

RefCell is only an option as we want to be able to acquire &mut T in arbitrary parts of our code. Otherwise, without a runtime check, we can end up with several &mut T references in our code, which is potentially UB and strictly prohibited.

From a global perspective, using RefCell comes with a cost of rechecking all borrow_ref and borrow_ref_mut to ensure they aren't called while another reference is still alive, but we can't do anything about it.

2. MutexSlot<T>

The second case is covered by using Mutex<OnceCell> and is more interesting as it has a BONUS FEATURE:

static MY_STRING: MutexSlot<String> = MutexSlot::new_slot();

// ...Somewhere in the code
critical_section::with(|cs| {
    MY_STRING.set(cs, String::from("WHOA!"));
});


// ...Somewhere else
println!("{}", MY_STRING.get_nocs());

Implementation:

/// MutexMut, but specialized for the following usage:
/// 1) Set value using set(tok, v). Can set value only once!
/// 2) Get value using get(tok) in your code (if T:Sync, get_nocs is available from any thread!)
///
/// Respect this order! otherwise, you will get a panic
pub struct MutexSlot<T>(Mutex<OnceCell<T>>, AtomicBool);

impl<T> MutexSlot<T> {
    pub const fn new_slot() -> Self {
        MutexSlot(Mutex::new(OnceCell::new()), AtomicBool::new(false))
    }

    /// Must be called only once! Otherwise, you will get a panic
    #[inline]
    #[track_caller]
    pub fn set(&self, cs: CriticalSection, v: T) {
        self.0.borrow(cs).set(v).ok().unwrap();
        self.1.store(true, Ordering::Relaxed);
    }

    /// Must be called after `set` call; otherwise, you will get a panic
    #[inline]
    #[track_caller]
    pub fn get<'cs>(&'cs self, cs: CriticalSection<'cs>) -> &'cs T {
        self.0.borrow(cs).get().unwrap()
    }

    /// Must be called after `set` call; otherwise, you will get a panic
    #[inline]
    pub fn try_get<'cs>(&'cs self, cs: CriticalSection<'cs>) -> Option<&'cs T> {
        self.0.borrow(cs).get()
    }
}

// Deref used because we can't add a method to a foreign type -> use your own transparent wrapper
impl<T> Deref for MutexSlot<T> {
    type Target = Mutex<OnceCell<T>>;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T: Sync> MutexSlot<T> {
    /// If your type is Sync, you can get inner value from anywhere
    ///
    /// Must be called after `set` call; otherwise, you will get a panic
    ///
    /// # Panic
    /// Panics if not set before
    #[inline]
    #[track_caller]
    pub fn get_nocs(&self) -> &T {
        // Additional check required to prevent simultaneous execution with `set`
        if !self.1.load(Ordering::Relaxed) {
            defmt::panic!("Trying to get ThreadModeSlot value without earlier call to `set`!");
        }
        // Safety:
        // 1) The inner type is Sync, we give only shared ref, and value is guaranteed to be set.
        // 2) by atomic flag check, we ensure that set was called before
        unsafe {self.0.borrow(CriticalSection::new()).get().unwrap_unchecked()}
    }

    /// If your type is Sync, you can get the inner value from anywhere
    ///
    /// Must be called after the `set` call; otherwise, you will get a panic
    #[inline]
    pub fn try_get_nocs(&self) -> Option<&T> {
        // Additional check required to prevent simultaneous execution with `set`
        if !self.1.load(Ordering::Relaxed) {
            return None;
        }
        // Safety:
        // 1) Inner type is Sync, we give only shared ref, and the value is guaranteed to be set.
        // 2) by atomic flag check, we ensure that set was called before
        unsafe {Some(self.0.borrow(CriticalSection::new()).get().unwrap_unchecked())}
    }
}

Let's dive deeper!

Without the last impl block, we could stay with Mutex<OnceCell<T>> because it already provides us with the most important methods: get and set. (try_get is defined but never used in my code).

Now we observe the magic:
If the inner type is Sync, the only thing we need to check is whether the value was set before. If it's true, we can safely get shared reference from anywhere in the code without unnecessary critical sections, which makes code much cleaner.

One additional practical considerations is that we usually 100% sure when this value was initialized, so we can safely call 'get_nocs. Again, try_get_nocs` is implemented but never used in my code.

3. ThreadModeUsage<T> (as well as ThreadModeMut<T>, ThreadModeSlot<T>

Not exactly a mutex but very practically useful type.

It is also very ergonomic when your data doesn't need to be shared between interrupt handlers.

Most of the code logic usually happens in thread mode, where contention is impossible. Therefore, we can use a runtime check to acquire ThreadModeToken. By leveraging Rust's Send/Sync rust rules and making it !Sync+!Send, we statically guarantee that it never leaves thread mode.

Usage example:

static MY_NUMBER: ThreadModeUsage<u64> = ThreadModeUsage::new(19);

// .. from thread mode
let tok = ThreadModeToken::acquire();
println!("My number is {}", STATE.borrow(&tok));

The idea is the same, and Cell/RefCell/_Anything_Cell can be placed inside.

ThreadModeMut and ThreadModeSlot variations are also possible.

Implementation:

/// ThreadModeUsage allows putting RefCell, Cell and other !Sync types into static variables, if they are used from thread mode only.
/// Runtime check guarantees that it is used from thread mode only and panics if it is not.
///
/// See ThreadModeToken for more information
pub struct ThreadModeUsage<T> {
    /// Safety: Access to inner only if the token is confirmed to live as long as the returning ref to inner
    inner: T,
}

impl<T> ThreadModeUsage<T> {
    pub const fn new(data: T) -> ThreadModeUsage<T> {
        ThreadModeUsage { inner: data }
    }

    /// Acquire a shared reference to inner data. Token can be acquired with ThreadModeToken::acquire()
    /// Recursive locking is possible, do not worry about it
    #[inline]
    pub fn borrow<'a>(&'a self, _tok: &'a ThreadModeToken) -> &'a T {
        &self.inner
    }
}

/// Use `ThreadModeToken::acquire()` to execute runtime check and unlock `ThreadModeUsage` mutexes.
///
/// !Send and !Sync enforce ThreadModeToken to stay bounded by a single thread forever. \
/// Owning a ThreadModeToken or a reference to ThreadModeToken guarantees that shared access to the same resource is safe. \
/// ThreadModeToken can only be acquired after runtime check for the current active exception vector, therefore it's safe.
#[derive(Copy, Clone)]
pub struct ThreadModeToken {
    _private: PhantomData<*mut ()>,
}

impl ThreadModeToken {
    /// Acquire token for thread mode.
    /// Best place to call is at the beginning of a function to explicitly state that this function is thread mode only.
    /// Panicking:
    /// Panics if called from interrupt mode.
    #[inline(always)]
    #[track_caller]
    pub fn acquire() -> ThreadModeToken {
        if !Self::in_thread_mode() {
            defmt::panic!("ThreadModeToken::acquire() called from thread mode!");
        }
        ThreadModeToken {
            _private: PhantomData,
        }
    }

    /// Acquire token for thread mode. Returns None if called from interrupt mode.
    /// The best place to call is at the beginning of function.
    /// Allows to split logic into two paths: thread mode and interrupt mode.
    #[inline(always)]
    pub fn try_acquire() -> Option<ThreadModeToken> {
        if !Self::in_thread_mode() {
            return None;
        }
        Some(ThreadModeToken {
            _private: PhantomData,
        })
    }

    /// For cortex-m cores only!
    #[inline(always)]
    pub fn in_thread_mode() -> bool {
        (unsafe { (0xE000ED04 as *const u32).read_volatile() } & 0x1FF) == 0
    }

    pub unsafe fn new_unsafe() -> ThreadModeToken {
        ThreadModeToken {
            _private: PhantomData,
        }
    }
}

4. MaskedMutex<T, const INTERRUPT_MASK> (as well as MaskedMutexMut<T, INTERRUPT_MASK>, MaskedMutexSlot<T, INTERRUPT_MASK>)

Will not post the implementation here because it is big and specific to microcontroller, but this type is also very useful and efficient.

Idea is to define critical sections that disable only a specific set of interrupts, so delays in the interrupt handlers will be much smaller.
For example if you have the I2C state under mutex and the UART state under another, both I2C and UART interrupts will be blocked while any of states is accessed from thread mode. By masking only the required set of interrupt handlers, others will continue execution without causing delays.

However, implementing such a type across many different platforms and microcontrollers may be difficult.

[tisonkun]

I abstract the usage as:

#[cfg(feature = "critical-section")]
mod critical_section;
#[cfg(feature = "critical-section")]
pub(crate) use critical_section::Mutex;

#[cfg(all(not(feature = "critical-section"), feature = "std"))]
mod std_mutex;
#[cfg(all(not(feature = "critical-section"), feature = "std"))]
pub(crate) use std_mutex::Mutex;

// std_mutext.rs
pub(crate) struct Mutex<T>(std::sync::Mutex<T>);

impl<T> Mutex<T> {
    #[must_use]
    #[inline]
    pub(crate) const fn new(t: T) -> Self {
        Self(std::sync::Mutex::new(t))
    }

    pub(crate) fn with<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&mut T) -> R,
    {
        let mut this = self.0.lock().unwrap_or_else(PoisonError::into_inner);
        f(&mut this)
    }
}

// critical_section.rs
pub(crate) struct Mutex<T>(critical_section::Mutex<RefCell<T>>);

impl<T> Mutex<T> {
    #[must_use]
    #[inline]
    pub(crate) const fn new(t: T) -> Self {
        Self(critical_section::Mutex::new(t))
    }

    pub(crate) fn with<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&mut T) -> R,
    {
        critical_section::with(|cs| {
            let mut this = self
                .0
                .borrow(cs)
                .try_borrow_mut()
                .expect("borrow with token 'cs' exactly once");
            f(&mut this)
        })
    }
}

It seems work well.