Skip to main content

embassy_supervisor/data_deps/
leased.rs

1use core::ops::Deref;
2use core::sync::atomic::Ordering;
3
4use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
5use embassy_sync::signal::Signal;
6use portable_atomic::AtomicU32;
7
8use crate::{Sig, TaskNode};
9
10const CLOSED: u32 = 1 << 31;
11const COUNT: u32 = !CLOSED;
12
13#[repr(C)]
14/// A signal-like value whose producer can wait until all outstanding readers finish.
15///
16/// `Leased` counts live claims; [`drain`](Self::drain) blocks until the count
17/// reaches zero, which is used to hold a producer's stop until consumers are done.
18pub struct Leased<T> {
19    inner: T,
20    state: AtomicU32,
21    idle: Signal<CriticalSectionRawMutex, ()>,
22}
23
24impl<T> Leased<T> {
25    /// Wrap `inner` as a leased value.
26    pub const fn new(inner: T) -> Self {
27        Self {
28            inner,
29            state: AtomicU32::new(0),
30            idle: Signal::new(),
31        }
32    }
33
34    /// Return the number of live leases.
35    pub fn leases(&self) -> u32 {
36        self.state.load(Ordering::Acquire) & COUNT
37    }
38
39    /// Return `true` if the value has been drained.
40    pub fn is_drained(&self) -> bool {
41        self.state.load(Ordering::Acquire) & CLOSED != 0
42    }
43
44    /// Close the value to new leases and wait until all current leases drop.
45    pub async fn drain(&self) {
46        self.state.fetch_or(CLOSED, Ordering::AcqRel);
47        while self.state.load(Ordering::Acquire) & COUNT != 0 {
48            self.idle.wait().await;
49        }
50    }
51
52    /// Reopen the value to new leases after a previous drain.
53    pub fn reopen(&self) {
54        self.idle.reset();
55        self.state.fetch_and(!CLOSED, Ordering::AcqRel);
56    }
57
58    /// Acquire a live lease, or `None` if the value is drained.
59    pub fn lease(&'static self) -> Option<Lease<T>> {
60        let mut cur = self.state.load(Ordering::Acquire);
61        loop {
62            if cur & CLOSED != 0 {
63                return None;
64            }
65            match self.state.compare_exchange_weak(
66                cur,
67                cur + 1,
68                Ordering::AcqRel,
69                Ordering::Acquire,
70            ) {
71                Ok(_) => return Some(Lease { target: self }),
72                Err(seen) => cur = seen,
73            }
74        }
75    }
76}
77
78/// The wrapper is transparent to everything that is not a lease.
79impl<T> Deref for Leased<T> {
80    type Target = T;
81    fn deref(&self) -> &T {
82        &self.inner
83    }
84}
85
86/// Polling a wrapped signal is polling what it wraps, so an `observed` entry
87/// keeps working when a signal gains a lease count.
88#[cfg(feature = "coupling-observe")]
89impl<T: crate::Observable> crate::Observable for Leased<T> {
90    fn change_token(&self) -> u32 {
91        self.inner.change_token()
92    }
93}
94
95/// A live claim on a [`Leased`] signal: the producer's `drain` does not return
96/// until every `Lease` has been dropped.
97pub struct Lease<T: 'static> {
98    target: &'static Leased<T>,
99}
100
101impl<T> Deref for Lease<T> {
102    type Target = T;
103    fn deref(&self) -> &T {
104        &self.target.inner
105    }
106}
107
108impl<T> Drop for Lease<T> {
109    fn drop(&mut self) {
110        if self.target.state.fetch_sub(1, Ordering::AcqRel) & COUNT == 1 {
111            self.target.idle.signal(());
112        }
113    }
114}
115
116impl TaskNode {
117    /// Acquire a lease on a `Sig<Leased<T>>` signal.
118    pub fn lease<T: Sync>(&self, s: Sig<Leased<T>>) -> Option<Lease<T>> {
119        s.target.lease()
120    }
121}