use rtos semaphore as mutex in esp-radio refcount

esp-radio/src/refcount.rs

@@ -1,83 +1,119 @@
-use core::sync::atomic::Ordering;
+use core::{cell::UnsafeCell, ptr::null_mut, sync::atomic::Ordering};
 
-use portable_atomic::AtomicU32;
+use esp_radio_rtos_driver::semaphore::{SemaphoreHandle, SemaphoreKind, SemaphorePtr};
+use portable_atomic::AtomicPtr;
 
-// Refcount of 1 is special, indicating that the radio is being initialized or deinitialized. If a
-// caller encounters this state, it must spin until the refcount changes.
+/// Resource guard that handles initialization and deinitalization gracefully
+/// using [`esp_radio_rtos_driver`]'s API.
+pub(crate) struct Refcount {
+    counter: UnsafeCell<u32>,
+    sem: AtomicPtr<()>,
+}
 
-/// A resource guard that uses a lock-free reference count to track usage.
-pub(crate) struct Refcount(AtomicU32);
+unsafe impl Sync for Refcount {}
+
+impl Drop for Refcount {
+    fn drop(&mut self) {
+        let sem = self.sem.load(Ordering::Relaxed);
+        if let Some(sem) = SemaphorePtr::new(sem) {
+            drop(unsafe { SemaphoreHandle::from_ptr(sem) });
+        }
+    }
+}
 
 impl Refcount {
     pub const fn new() -> Self {
-        Self(AtomicU32::new(0))
+        Self {
+            counter: UnsafeCell::new(0),
+            sem: AtomicPtr::new(null_mut()),
+        }
     }
 
-    pub fn increment(&self, on_first: impl FnOnce()) {
-        loop {
-            let op = self
-                .0
-                .fetch_update(Ordering::Release, Ordering::Acquire, |old| {
-                    if old == 1 { None } else { Some(old + 1) }
-                });
-
-            match op {
-                Ok(0) => {
-                    on_first();
-                    self.0.store(2, Ordering::Release);
-                    break;
-                }
-                Ok(_) => break,
-                Err(_) => {}
+    fn use_sem_or_init<T>(&self, f: impl FnOnce(&SemaphoreHandle) -> T) -> T {
+        if self.sem.load(Ordering::Relaxed).is_null() {
+            core::hint::cold_path();
+
+            let sem = SemaphoreHandle::new(SemaphoreKind::Mutex).leak();
+
+            if self
+                .sem
+                .compare_exchange(
+                    null_mut(),
+                    sem.as_ptr(),
+                    Ordering::Release,
+                    Ordering::Relaxed,
+                )
+                .is_err()
+            {
+                core::hint::cold_path();
+
+                drop(unsafe { SemaphoreHandle::from_ptr(sem) });
             }
         }
+
+        let sem = unsafe { SemaphorePtr::new_unchecked(self.sem.load(Ordering::Acquire)) };
+        f(unsafe { SemaphoreHandle::ref_from_ptr(&sem) })
+    }
+
+    fn try_use_sem<T>(&self, f: impl FnOnce(&SemaphoreHandle) -> T) -> Option<T> {
+        if self.sem.load(Ordering::Relaxed).is_null() {
+            core::hint::cold_path();
+
+            None
+        } else {
+            let sem = unsafe { SemaphorePtr::new_unchecked(self.sem.load(Ordering::Acquire)) };
+            Some(f(unsafe { SemaphoreHandle::ref_from_ptr(&sem) }))
+        }
+    }
+
+    fn lock<T>(&self, f: impl FnOnce(&mut u32) -> T) -> T {
+        self.use_sem_or_init(|sem| {
+            sem.take(None);
+            let ret = f(unsafe { self.counter.get().as_mut_unchecked() });
+            sem.give();
+            ret
+        })
+    }
+
+    fn try_lock<T>(&self, f: impl FnOnce(&mut u32) -> T) -> Option<T> {
+        self.try_use_sem(|sem| {
+            sem.take(None);
+            let ret = f(unsafe { self.counter.get().as_mut_unchecked() });
+            sem.give();
+            ret
+        })
+    }
+
+    pub fn increment(&self, on_first: impl FnOnce()) {
+        self.lock(|counter| {
+            if *counter == 0 {
+                on_first();
+            }
+            *counter = counter.checked_add(1).expect("refcount overflow");
+        });
     }
 
     #[cfg(feature = "wifi")]
     pub fn try_increment<E>(&self, on_first: impl FnOnce() -> Result<(), E>) -> Result<bool, E> {
-        loop {
-            let op = self
-                .0
-                .fetch_update(Ordering::Release, Ordering::Acquire, |old| {
-                    if old == 1 { None } else { Some(old + 1) }
-                });
-
-            match op {
-                Ok(0) => {
-                    return match on_first() {
-                        Ok(()) => {
-                            self.0.store(2, Ordering::Release);
-                            Ok(true)
-                        }
-                        Err(e) => {
-                            self.0.store(0, Ordering::Release);
-                            Err(e)
-                        }
-                    };
-                }
-                Ok(_) => return Ok(false),
-                Err(_) => {}
+        self.lock(|counter| {
+            let prev = *counter;
+            *counter = counter.checked_add(1).expect("refcount overflow");
+
+            if prev == 0 {
+                on_first().inspect_err(|_| *counter = 0).map(|_| true)
+            } else {
+                Ok(false)
             }
-        }
+        })
     }
 
     pub fn decrement(&self, on_last: impl FnOnce()) {
-        loop {
-            let op = self
-                .0
-                .fetch_update(Ordering::Release, Ordering::Acquire, |old| {
-                    if old == 1 { None } else { Some(old - 1) }
-                });
-
-            match op {
-                Ok(2) => {
-                    on_last();
-                    self.0.store(0, Ordering::Release);
-                    break;
-                }
-                Ok(_) => break,
-                Err(_) => {}
+        self.try_lock(|counter| {
+            if *counter == 0 {
+                on_last();
             }
-        }
+            *counter = counter.checked_sub(1).expect("decrementing count of zero");
+        })
+        .expect("decrementing before any successful increment")
     }
 }