radix_engine/kernel/
substate_locks.rs

1use crate::internal_prelude::*;
2
3pub struct SubstateLockError;
4
5#[derive(Debug, Copy, Clone, PartialEq, Eq, Sbor)]
6pub enum SubstateLockState {
7    Read(usize),
8    Write,
9}
10
11impl SubstateLockState {
12    fn no_lock() -> Self {
13        Self::Read(0)
14    }
15
16    fn is_locked(&self) -> bool {
17        !matches!(self, SubstateLockState::Read(0usize))
18    }
19
20    fn try_lock(&mut self, read_only: bool) -> Result<(), SubstateLockError> {
21        match self {
22            SubstateLockState::Read(n) => {
23                if read_only {
24                    *n = *n + 1;
25                } else {
26                    if *n != 0 {
27                        return Err(SubstateLockError);
28                    }
29                    *self = SubstateLockState::Write;
30                }
31            }
32            SubstateLockState::Write => {
33                return Err(SubstateLockError);
34            }
35        }
36
37        Ok(())
38    }
39
40    fn unlock(&mut self) {
41        match self {
42            SubstateLockState::Read(n) => {
43                *n = *n - 1;
44            }
45            SubstateLockState::Write => {
46                *self = SubstateLockState::no_lock();
47            }
48        }
49    }
50}
51
52pub struct SubstateLocks<D> {
53    locks: IndexMap<u32, (NodeId, PartitionNumber, SubstateKey, D)>,
54    substate_lock_states: NonIterMap<(NodeId, PartitionNumber, SubstateKey), SubstateLockState>,
55    node_num_locked: NonIterMap<NodeId, usize>,
56    next_lock_id: u32,
57}
58
59impl<D> SubstateLocks<D> {
60    pub fn new() -> Self {
61        Self {
62            locks: index_map_new(),
63            substate_lock_states: NonIterMap::new(),
64            node_num_locked: NonIterMap::new(),
65            next_lock_id: 0u32,
66        }
67    }
68
69    fn new_lock_handle(
70        &mut self,
71        node_id: &NodeId,
72        partition_num: PartitionNumber,
73        substate_key: &SubstateKey,
74        data: D,
75    ) -> u32 {
76        let new_lock = self.next_lock_id;
77        self.locks.insert(
78            new_lock,
79            (*node_id, partition_num, substate_key.clone(), data),
80        );
81        self.next_lock_id += 1;
82        new_lock
83    }
84
85    pub fn node_is_locked(&self, node_id: &NodeId) -> bool {
86        self.node_num_locked
87            .get(node_id)
88            .map(|e| *e > 0)
89            .unwrap_or(false)
90    }
91
92    pub fn is_locked(
93        &self,
94        node_id: &NodeId,
95        partition_num: PartitionNumber,
96        substate_key: &SubstateKey,
97    ) -> bool {
98        if let Some(state) =
99            self.substate_lock_states
100                .get(&(node_id.clone(), partition_num, substate_key.clone()))
101        {
102            state.is_locked()
103        } else {
104            false
105        }
106    }
107
108    pub fn lock(
109        &mut self,
110        node_id: &NodeId,
111        partition_num: PartitionNumber,
112        substate_key: &SubstateKey,
113        read_only: bool,
114        data: D,
115    ) -> Option<u32> {
116        let lock_state = self
117            .substate_lock_states
118            .entry((node_id.clone(), partition_num, substate_key.clone()))
119            .or_insert(SubstateLockState::no_lock());
120        match lock_state.try_lock(read_only) {
121            Ok(()) => {}
122            Err(_) => {
123                return None;
124            }
125        }
126
127        let count = self.node_num_locked.entry(*node_id).or_insert(0);
128        *count = *count + 1;
129
130        let handle = self.new_lock_handle(node_id, partition_num, substate_key, data);
131        Some(handle)
132    }
133
134    pub fn get(&self, handle: u32) -> &(NodeId, PartitionNumber, SubstateKey, D) {
135        self.locks.get(&handle).unwrap()
136    }
137
138    pub fn get_mut(&mut self, handle: u32) -> &mut (NodeId, PartitionNumber, SubstateKey, D) {
139        self.locks.get_mut(&handle).unwrap()
140    }
141
142    pub fn unlock(&mut self, handle: u32) -> (NodeId, PartitionNumber, SubstateKey, D) {
143        let (node_id, partition_num, substate_key, data) = self.locks.swap_remove(&handle).unwrap();
144        let full_key = (node_id, partition_num, substate_key);
145
146        let lock_state = self.substate_lock_states.get_mut(&full_key).unwrap();
147        lock_state.unlock();
148
149        let count = self.node_num_locked.entry(node_id).or_insert(0);
150        *count = *count - 1;
151
152        (full_key.0, full_key.1, full_key.2, data)
153    }
154}
radix_engine/kernel/substate_locks.rs

radix_engine/kernel/
substate_locks.rs
