driver_interface/
lock.rs

1use alloc::sync::{Arc, Weak};
2use core::{
3    cell::UnsafeCell,
4    ops::{Deref, DerefMut},
5    sync::atomic::{AtomicBool, Ordering},
6};
7
8use crate::custom_type;
9
10custom_type!(PId, usize, "{:?}");
11
12pub enum LockError {
13    UsedByOthers(PId),
14}
15
16pub struct Lock<T> {
17    data: Arc<LockData<T>>,
18}
19
20impl<T> Lock<T> {
21    pub fn new(data: T) -> Self {
22        Lock {
23            data: Arc::new(LockData::new(data)),
24        }
25    }
26
27    pub fn try_borrow(&self, pid: PId) -> Result<LockGuard<T>, LockError> {
28        match self
29            .data
30            .borrowed
31            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
32        {
33            Ok(_) => {
34                unsafe {
35                    (*self.data.used.get()).replace(pid);
36                }
37                Ok(LockGuard {
38                    data: self.data.clone(),
39                })
40            }
41            Err(_) => {
42                let pid = unsafe { *self.data.used.get() };
43                Err(LockError::UsedByOthers(pid.unwrap()))
44            }
45        }
46    }
47
48    pub fn weak(&self) -> LockWeak<T> {
49        LockWeak {
50            data: Arc::downgrade(&self.data),
51        }
52    }
53
54    /// 强制获取设备
55    ///
56    /// # Safety
57    /// 一般用于中断处理中
58    pub unsafe fn force_use(&self) -> *mut T {
59        self.data.data.get()
60    }
61}
62
63impl<T: Sync + Send> Deref for Lock<T> {
64    type Target = T;
65
66    fn deref(&self) -> &Self::Target {
67        unsafe { &*self.data.data.get() }
68    }
69}
70
71pub struct LockWeak<T> {
72    data: Weak<LockData<T>>,
73}
74
75impl<T> LockWeak<T> {
76    pub fn upgrade(&self) -> Option<Lock<T>> {
77        self.data.upgrade().map(|data| Lock { data })
78    }
79}
80
81struct LockData<T> {
82    borrowed: AtomicBool,
83    used: UnsafeCell<Option<PId>>,
84    data: UnsafeCell<T>,
85}
86
87unsafe impl<T: Send> Send for LockData<T> {}
88unsafe impl<T: Send> Sync for LockData<T> {}
89
90impl<T> LockData<T> {
91    fn new(data: T) -> Self {
92        LockData {
93            borrowed: AtomicBool::new(false),
94            used: UnsafeCell::new(None),
95            data: UnsafeCell::new(data),
96        }
97    }
98}
99
100pub struct LockGuard<T> {
101    data: Arc<LockData<T>>,
102}
103
104impl<T> Drop for LockGuard<T> {
105    fn drop(&mut self) {
106        unsafe {
107            (*self.data.used.get()).take();
108        }
109        self.data.borrowed.store(false, Ordering::Release);
110    }
111}
112
113impl<T> Deref for LockGuard<T> {
114    type Target = T;
115
116    fn deref(&self) -> &Self::Target {
117        unsafe { &*self.data.data.get() }
118    }
119}
120
121impl<T> DerefMut for LockGuard<T> {
122    fn deref_mut(&mut self) -> &mut Self::Target {
123        unsafe { &mut *self.data.data.get() }
124    }
125}