strat9_kernel/memory/
ownership.rs

1//! Central ownership tracking for physical blocks.
2
3use alloc::collections::BTreeMap;
4
5use smallvec::{smallvec, SmallVec};
6
7use crate::{capability::CapId, sync::SpinLock};
8
9use super::{
10    block::{BlockHandle, BuddyReserved, Exclusive, PhysBlock, Released},
11    block_meta::get_block_meta,
12};
13
14/// Runtime ownership state of a block.
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16#[repr(u8)]
17pub enum BlockState {
18    /// Block reserved by the buddy and not yet claimed.
19    BuddyReserved = 0,
20    /// Block exclusively owned by a single capability.
21    Exclusive = 1,
22    /// Block shared by multiple capabilities.
23    Shared = 2,
24    /// Block released to the buddy-facing state.
25    Free = 3,
26}
27
28/// Ownership entry associated with a block.
29#[derive(Debug, Clone)]
30pub struct OwnerEntry {
31    /// Current runtime state of the block.
32    pub state: BlockState,
33    /// Number of capabilities referencing the block.
34    pub refcount: u32,
35    /// Capabilities that currently reference the block.
36    pub caps: SmallVec<[CapId; 4]>,
37    /// Temporary non-capability pins held while publishing or revoking a mapping.
38    pub transient_refs: u32,
39}
40
41/// Errors returned by the ownership layer.
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub enum OwnerError {
44    /// The block has no ownership entry.
45    NotFound,
46    /// The block has already been claimed.
47    DoubleClaim,
48    /// The capability is already recorded for the block.
49    CapAlreadyPresent,
50    /// The capability is not recorded for the block.
51    CapNotFound,
52    /// The block still has live references after a release attempt.
53    StillReferenced,
54    /// The reference count cannot be represented in `u32`.
55    CountOverflow,
56}
57
58/// Result of removing one capability reference from a block.
59#[derive(Debug)]
60pub enum RemoveRefResult {
61    /// The block has no remaining owners and can return to the buddy.
62    Freed(PhysBlock<Released>),
63    /// Exactly one owner remains.
64    NowExclusive {
65        /// The last remaining capability.
66        remaining_cap: CapId,
67    },
68    /// No capability remains, but transient pins still keep the block alive.
69    StillPinned {
70        /// Current transient reference count.
71        transient_refs: u32,
72    },
73    /// Multiple owners still reference the block.
74    StillShared {
75        /// Current shared reference count.
76        refcount: u32,
77    },
78}
79
80/// Source of truth for block ownership.
81pub struct OwnershipTable {
82    entries: SpinLock<BTreeMap<BlockHandle, OwnerEntry>>,
83}
84
85impl OwnershipTable {
86    /// Creates an empty ownership table.
87    pub fn new() -> Self {
88        Self {
89            entries: SpinLock::new(BTreeMap::new()),
90        }
91    }
92
93    fn sync_meta(handle: BlockHandle, refcount: u32) {
94        let meta = get_block_meta(handle.base);
95        meta.set_order(handle.order);
96        meta.set_refcount(refcount);
97    }
98
99    fn classify_refcount(refcount: u32) -> BlockState {
100        if refcount <= 1 {
101            BlockState::Exclusive
102        } else {
103            BlockState::Shared
104        }
105    }
106
107    fn total_refs(entry: &OwnerEntry) -> Result<u32, OwnerError> {
108        let cap_refs = u32::try_from(entry.caps.len()).map_err(|_| OwnerError::CountOverflow)?;
109        cap_refs
110            .checked_add(entry.transient_refs)
111            .ok_or(OwnerError::CountOverflow)
112    }
113
114    /// Claims a reserved block for the provided capability.
115    pub fn claim(
116        &self,
117        block: PhysBlock<BuddyReserved>,
118        cap_id: CapId,
119    ) -> Result<PhysBlock<Exclusive>, OwnerError> {
120        let handle = block.handle();
121        let mut entries = self.entries.lock();
122        if entries.contains_key(&handle) {
123            return Err(OwnerError::DoubleClaim);
124        }
125        entries.insert(
126            handle,
127            OwnerEntry {
128                state: BlockState::Exclusive,
129                refcount: 1,
130                caps: smallvec![cap_id],
131                transient_refs: 0,
132            },
133        );
134        Self::sync_meta(handle, 1);
135        Ok(block.into_exclusive())
136    }
137
138    /// Ensures `cap_id` is recorded as a live reference on `handle`.
139    pub fn ensure_ref(&self, handle: BlockHandle, cap_id: CapId) -> Result<BlockState, OwnerError> {
140        let mut entries = self.entries.lock();
141        if let Some(entry) = entries.get_mut(&handle) {
142            if entry.caps.iter().any(|existing| *existing == cap_id) {
143                return Err(OwnerError::CapAlreadyPresent);
144            }
145            entry.caps.push(cap_id);
146            entry.refcount = Self::total_refs(entry)?;
147            entry.state = Self::classify_refcount(entry.refcount);
148            Self::sync_meta(handle, entry.refcount);
149            Ok(entry.state)
150        } else {
151            entries.insert(
152                handle,
153                OwnerEntry {
154                    state: BlockState::Exclusive,
155                    refcount: 1,
156                    caps: smallvec![cap_id],
157                    transient_refs: 0,
158                },
159            );
160            Self::sync_meta(handle, 1);
161            Ok(BlockState::Exclusive)
162        }
163    }
164
165    /// Adds a capability reference to an existing owned block.
166    pub fn add_ref(&self, handle: BlockHandle, cap_id: CapId) -> Result<BlockState, OwnerError> {
167        let mut entries = self.entries.lock();
168        let entry = entries.get_mut(&handle).ok_or(OwnerError::NotFound)?;
169        if entry.caps.iter().any(|existing| *existing == cap_id) {
170            return Err(OwnerError::CapAlreadyPresent);
171        }
172
173        entry.caps.push(cap_id);
174        entry.refcount = Self::total_refs(entry)?;
175        entry.state = Self::classify_refcount(entry.refcount);
176        Self::sync_meta(handle, entry.refcount);
177
178        Ok(entry.state)
179    }
180
181    /// Adds a temporary pin to keep the block alive while publishing a mapping.
182    pub fn pin(&self, handle: BlockHandle) -> Result<u32, OwnerError> {
183        let mut entries = self.entries.lock();
184        let entry = entries.get_mut(&handle).ok_or(OwnerError::NotFound)?;
185        entry.transient_refs = entry
186            .transient_refs
187            .checked_add(1)
188            .ok_or(OwnerError::CountOverflow)?;
189        entry.refcount = Self::total_refs(entry)?;
190        entry.state = Self::classify_refcount(entry.refcount);
191        Self::sync_meta(handle, entry.refcount);
192        Ok(entry.refcount)
193    }
194
195    /// Removes a capability reference from a block.
196    pub fn remove_ref(
197        &self,
198        handle: BlockHandle,
199        cap_id: CapId,
200    ) -> Result<RemoveRefResult, OwnerError> {
201        let mut entries = self.entries.lock();
202        let entry = entries.get_mut(&handle).ok_or(OwnerError::NotFound)?;
203        let position = entry
204            .caps
205            .iter()
206            .position(|existing| *existing == cap_id)
207            .ok_or(OwnerError::CapNotFound)?;
208
209        entry.caps.remove(position);
210        entry.refcount = Self::total_refs(entry)?;
211
212        match entry.refcount {
213            0 => {
214                entries.remove(&handle);
215                Self::sync_meta(handle, 0);
216                Ok(RemoveRefResult::Freed(PhysBlock::from_handle(handle)))
217            }
218            1 => {
219                entry.state = BlockState::Exclusive;
220                Self::sync_meta(handle, entry.refcount);
221                if let Some(&remaining_cap) = entry.caps.first() {
222                    Ok(RemoveRefResult::NowExclusive { remaining_cap })
223                } else {
224                    Ok(RemoveRefResult::StillPinned {
225                        transient_refs: entry.transient_refs,
226                    })
227                }
228            }
229            refcount => {
230                entry.state = BlockState::Shared;
231                Self::sync_meta(handle, entry.refcount);
232                Ok(RemoveRefResult::StillShared { refcount })
233            }
234        }
235    }
236
237    /// Removes one temporary pin and releases the block if this was the last live reference.
238    pub fn unpin(&self, handle: BlockHandle) -> Result<RemoveRefResult, OwnerError> {
239        let mut entries = self.entries.lock();
240        let entry = entries.get_mut(&handle).ok_or(OwnerError::NotFound)?;
241        if entry.transient_refs == 0 {
242            return Err(OwnerError::CapNotFound);
243        }
244
245        entry.transient_refs -= 1;
246        entry.refcount = Self::total_refs(entry)?;
247
248        match entry.refcount {
249            0 => {
250                entries.remove(&handle);
251                Self::sync_meta(handle, 0);
252                Ok(RemoveRefResult::Freed(PhysBlock::from_handle(handle)))
253            }
254            1 => {
255                entry.state = BlockState::Exclusive;
256                Self::sync_meta(handle, entry.refcount);
257                if let Some(&remaining_cap) = entry.caps.first() {
258                    Ok(RemoveRefResult::NowExclusive { remaining_cap })
259                } else {
260                    Ok(RemoveRefResult::StillPinned {
261                        transient_refs: entry.transient_refs,
262                    })
263                }
264            }
265            refcount => {
266                entry.state = BlockState::Shared;
267                Self::sync_meta(handle, entry.refcount);
268                Ok(RemoveRefResult::StillShared { refcount })
269            }
270        }
271    }
272
273    /// Releases an exclusive block and returns a buddy-facing handle if it becomes free.
274    pub fn release(
275        &self,
276        block: PhysBlock<Exclusive>,
277        cap_id: CapId,
278    ) -> Result<PhysBlock<Released>, OwnerError> {
279        match self.remove_ref(block.handle(), cap_id)? {
280            RemoveRefResult::Freed(released) => Ok(released),
281            RemoveRefResult::NowExclusive { .. }
282            | RemoveRefResult::StillPinned { .. }
283            | RemoveRefResult::StillShared { .. } => Err(OwnerError::StillReferenced),
284        }
285    }
286
287    /// Returns a snapshot of the ownership entry for the given block.
288    pub fn get(&self, handle: BlockHandle) -> Option<OwnerEntry> {
289        self.entries.lock().get(&handle).cloned()
290    }
291
292    /// Returns the live handle whose base physical address matches `base`, if any.
293    pub fn handle_for_base(&self, base: x86_64::PhysAddr) -> Option<BlockHandle> {
294        self.entries
295            .lock()
296            .keys()
297            .find(|handle| handle.base == base)
298            .copied()
299    }
300
301    /// Returns the live handle whose block currently contains `phys`, if any.
302    pub fn handle_containing(&self, phys: x86_64::PhysAddr) -> Option<BlockHandle> {
303        self.entries.lock().keys().find_map(|handle| {
304            let start = handle.base.as_u64();
305            let end = start.checked_add(handle.size_bytes())?;
306            let addr = phys.as_u64();
307            (addr >= start && addr < end).then_some(*handle)
308        })
309    }
310
311    /// Returns the current reference count for the given block.
312    pub fn refcount(&self, handle: BlockHandle) -> Option<u32> {
313        self.entries.lock().get(&handle).map(|entry| entry.refcount)
314    }
315}
316
317impl Default for OwnershipTable {
318    /// Creates an empty ownership table.
319    fn default() -> Self {
320        Self::new()
321    }
322}
strat9_kernel/memory/ownership.rs

strat9_kernel/memory/
ownership.rs
