wasmtime/runtime/component/
resource_table.rs

1use super::Resource;
2use crate::prelude::*;
3use alloc::collections::{BTreeMap, BTreeSet};
4use core::any::Any;
5use core::fmt;
6use core::mem;
7
8#[derive(Debug)]
9/// Errors returned by operations on `ResourceTable`
10pub enum ResourceTableError {
11    /// ResourceTable has no free keys
12    Full,
13    /// Resource not present in table
14    NotPresent,
15    /// Resource present in table, but with a different type
16    WrongType,
17    /// Resource cannot be deleted because child resources exist in the table. Consult wit docs for
18    /// the particular resource to see which methods may return child resources.
19    HasChildren,
20}
21
22impl fmt::Display for ResourceTableError {
23    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
24        match self {
25            Self::Full => write!(f, "resource table has no free keys"),
26            Self::NotPresent => write!(f, "resource not present"),
27            Self::WrongType => write!(f, "resource is of another type"),
28            Self::HasChildren => write!(f, "resource has children"),
29        }
30    }
31}
32
33impl core::error::Error for ResourceTableError {}
34
35/// The `ResourceTable` type maps a `Resource<T>` to its `T`.
36pub struct ResourceTable {
37    entries: Vec<Entry>,
38    free_head: Option<usize>,
39}
40
41#[derive(Debug)]
42enum Entry {
43    Free { next: Option<usize> },
44    Occupied { entry: TableEntry },
45}
46
47impl Entry {
48    pub fn occupied(&self) -> Option<&TableEntry> {
49        match self {
50            Self::Occupied { entry } => Some(entry),
51            Self::Free { .. } => None,
52        }
53    }
54
55    pub fn occupied_mut(&mut self) -> Option<&mut TableEntry> {
56        match self {
57            Self::Occupied { entry } => Some(entry),
58            Self::Free { .. } => None,
59        }
60    }
61}
62
63struct Tombstone;
64
65// Change this to `true` to assist with handle debugging in development if
66// necessary.
67const DELETE_WITH_TOMBSTONE: bool = false;
68
69/// This structure tracks parent and child relationships for a given table entry.
70///
71/// Parents and children are referred to by table index. We maintain the
72/// following invariants:
73/// * the parent must exist when adding a child.
74/// * whenever a child is created, its index is added to children.
75/// * whenever a child is deleted, its index is removed from children.
76/// * an entry with children may not be deleted.
77#[derive(Debug)]
78struct TableEntry {
79    /// The entry in the table, as a boxed dynamically-typed object
80    entry: Box<dyn Any + Send>,
81    /// The index of the parent of this entry, if it has one.
82    parent: Option<u32>,
83    /// The indices of any children of this entry.
84    children: BTreeSet<u32>,
85}
86
87impl TableEntry {
88    fn new(entry: Box<dyn Any + Send>, parent: Option<u32>) -> Self {
89        Self {
90            entry,
91            parent,
92            children: BTreeSet::new(),
93        }
94    }
95    fn add_child(&mut self, child: u32) {
96        debug_assert!(!self.children.contains(&child));
97        self.children.insert(child);
98    }
99    fn remove_child(&mut self, child: u32) {
100        let was_removed = self.children.remove(&child);
101        debug_assert!(was_removed);
102    }
103}
104
105impl ResourceTable {
106    /// Create an empty table
107    pub fn new() -> Self {
108        ResourceTable {
109            entries: Vec::new(),
110            free_head: None,
111        }
112    }
113
114    /// Returns whether or not this table is empty.
115    ///
116    /// Note that this is an `O(n)` operation, where `n` is the number of
117    /// entries in the backing `Vec`.
118    pub fn is_empty(&self) -> bool {
119        self.entries.iter().all(|entry| match entry {
120            Entry::Free { .. } => true,
121            Entry::Occupied { entry } => entry.entry.downcast_ref::<Tombstone>().is_some(),
122        })
123    }
124
125    /// Create an empty table with at least the specified capacity.
126    pub fn with_capacity(capacity: usize) -> Self {
127        ResourceTable {
128            entries: Vec::with_capacity(capacity),
129            free_head: None,
130        }
131    }
132
133    /// Inserts a new value `T` into this table, returning a corresponding
134    /// `Resource<T>` which can be used to refer to it after it was inserted.
135    pub fn push<T>(&mut self, entry: T) -> Result<Resource<T>, ResourceTableError>
136    where
137        T: Send + 'static,
138    {
139        let idx = self.push_(TableEntry::new(Box::new(entry), None))?;
140        Ok(Resource::new_own(idx))
141    }
142
143    /// Pop an index off of the free list, if it's not empty.
144    fn pop_free_list(&mut self) -> Option<usize> {
145        if let Some(ix) = self.free_head {
146            // Advance free_head to the next entry if one is available.
147            match &self.entries[ix] {
148                Entry::Free { next } => self.free_head = *next,
149                Entry::Occupied { .. } => unreachable!(),
150            }
151            Some(ix)
152        } else {
153            None
154        }
155    }
156
157    /// Free an entry in the table, returning its [`TableEntry`]. Add the index to the free list.
158    fn free_entry(&mut self, ix: usize, debug: bool) -> TableEntry {
159        if debug {
160            // Instead of making this entry available for reuse, we leave a
161            // tombstone in debug mode.  This helps detect use-after-delete and
162            // double-delete bugs.
163            match mem::replace(
164                &mut self.entries[ix],
165                Entry::Occupied {
166                    entry: TableEntry {
167                        entry: Box::new(Tombstone),
168                        parent: None,
169                        children: BTreeSet::new(),
170                    },
171                },
172            ) {
173                Entry::Occupied { entry } => entry,
174                Entry::Free { .. } => unreachable!(),
175            }
176        } else {
177            let entry = match core::mem::replace(
178                &mut self.entries[ix],
179                Entry::Free {
180                    next: self.free_head,
181                },
182            ) {
183                Entry::Occupied { entry } => entry,
184                Entry::Free { .. } => unreachable!(),
185            };
186
187            self.free_head = Some(ix);
188
189            entry
190        }
191    }
192
193    /// Push a new entry into the table, returning its handle. This will prefer to use free entries
194    /// if they exist, falling back on pushing new entries onto the end of the table.
195    fn push_(&mut self, e: TableEntry) -> Result<u32, ResourceTableError> {
196        if let Some(free) = self.pop_free_list() {
197            self.entries[free] = Entry::Occupied { entry: e };
198            Ok(free.try_into().unwrap())
199        } else {
200            let ix = self
201                .entries
202                .len()
203                .try_into()
204                .map_err(|_| ResourceTableError::Full)?;
205            self.entries.push(Entry::Occupied { entry: e });
206            Ok(ix)
207        }
208    }
209
210    fn occupied(&self, key: u32) -> Result<&TableEntry, ResourceTableError> {
211        self.entries
212            .get(key as usize)
213            .and_then(Entry::occupied)
214            .ok_or(ResourceTableError::NotPresent)
215    }
216
217    fn occupied_mut(&mut self, key: u32) -> Result<&mut TableEntry, ResourceTableError> {
218        self.entries
219            .get_mut(key as usize)
220            .and_then(Entry::occupied_mut)
221            .ok_or(ResourceTableError::NotPresent)
222    }
223
224    /// Insert a resource at the next available index, and track that it has a
225    /// parent resource.
226    ///
227    /// The parent must exist to create a child. All children resources must
228    /// be destroyed before a parent can be destroyed - otherwise
229    /// [`ResourceTable::delete`] will fail with
230    /// [`ResourceTableError::HasChildren`].
231    ///
232    /// Parent-child relationships are tracked inside the table to ensure that
233    /// a parent resource is not deleted while it has live children. This
234    /// allows child resources to hold "references" to a parent by table
235    /// index, to avoid needing e.g. an `Arc<Mutex<parent>>` and the associated
236    /// locking overhead and design issues, such as child existence extending
237    /// lifetime of parent referent even after parent resource is destroyed,
238    /// possibility for deadlocks.
239    pub fn push_child<T, U>(
240        &mut self,
241        entry: T,
242        parent: &Resource<U>,
243    ) -> Result<Resource<T>, ResourceTableError>
244    where
245        T: Send + 'static,
246        U: 'static,
247    {
248        let parent = parent.rep();
249        self.occupied(parent)?;
250        let child = self.push_(TableEntry::new(Box::new(entry), Some(parent)))?;
251        self.occupied_mut(parent)?.add_child(child);
252        Ok(Resource::new_own(child))
253    }
254
255    /// Add an already-resident child to a resource.
256    pub fn add_child<T: 'static, U: 'static>(
257        &mut self,
258        child: Resource<T>,
259        parent: Resource<U>,
260    ) -> Result<(), ResourceTableError> {
261        let entry = self.occupied_mut(child.rep())?;
262        assert!(entry.parent.is_none());
263        entry.parent = Some(parent.rep());
264        self.occupied_mut(parent.rep())?.add_child(child.rep());
265        Ok(())
266    }
267
268    /// Remove a child to from a resource (but leave it in the table).
269    pub fn remove_child<T: 'static, U: 'static>(
270        &mut self,
271        child: Resource<T>,
272        parent: Resource<U>,
273    ) -> Result<(), ResourceTableError> {
274        let entry = self.occupied_mut(child.rep())?;
275        assert_eq!(entry.parent, Some(parent.rep()));
276        entry.parent = None;
277        self.occupied_mut(parent.rep())?.remove_child(child.rep());
278        Ok(())
279    }
280
281    /// Get an immutable reference to a resource of a given type at a given
282    /// index.
283    ///
284    /// Multiple shared references can be borrowed at any given time.
285    pub fn get<T: Any + Sized>(&self, key: &Resource<T>) -> Result<&T, ResourceTableError> {
286        self.get_(key.rep())?
287            .downcast_ref()
288            .ok_or(ResourceTableError::WrongType)
289    }
290
291    fn get_(&self, key: u32) -> Result<&dyn Any, ResourceTableError> {
292        let r = self.occupied(key)?;
293        Ok(&*r.entry)
294    }
295
296    /// Get an mutable reference to a resource of a given type at a given
297    /// index.
298    pub fn get_mut<T: Any + Sized>(
299        &mut self,
300        key: &Resource<T>,
301    ) -> Result<&mut T, ResourceTableError> {
302        self.get_any_mut(key.rep())?
303            .downcast_mut()
304            .ok_or(ResourceTableError::WrongType)
305    }
306
307    /// Returns the raw `Any` at the `key` index provided.
308    pub fn get_any_mut(&mut self, key: u32) -> Result<&mut dyn Any, ResourceTableError> {
309        let r = self.occupied_mut(key)?;
310        Ok(&mut *r.entry)
311    }
312
313    /// Remove the specified entry from the table.
314    pub fn delete<T>(&mut self, resource: Resource<T>) -> Result<T, ResourceTableError>
315    where
316        T: Any,
317    {
318        self.delete_maybe_debug(resource, DELETE_WITH_TOMBSTONE)
319    }
320
321    fn delete_maybe_debug<T>(
322        &mut self,
323        resource: Resource<T>,
324        debug: bool,
325    ) -> Result<T, ResourceTableError>
326    where
327        T: Any,
328    {
329        debug_assert!(resource.owned());
330        let entry = self.delete_entry(resource.rep(), debug)?;
331        match entry.entry.downcast() {
332            Ok(t) => Ok(*t),
333            Err(_e) => Err(ResourceTableError::WrongType),
334        }
335    }
336
337    fn delete_entry(&mut self, key: u32, debug: bool) -> Result<TableEntry, ResourceTableError> {
338        if !self.occupied(key)?.children.is_empty() {
339            return Err(ResourceTableError::HasChildren);
340        }
341        let e = self.free_entry(key as usize, debug);
342        if let Some(parent) = e.parent {
343            // Remove deleted resource from parent's child list.
344            // Parent must still be present because it can't be deleted while still having
345            // children:
346            self.occupied_mut(parent)
347                .expect("missing parent")
348                .remove_child(key);
349        }
350        Ok(e)
351    }
352
353    /// Zip the values of the map with mutable references to table entries corresponding to each
354    /// key. As the keys in the `BTreeMap` are unique, this iterator can give mutable references
355    /// with the same lifetime as the mutable reference to the [ResourceTable].
356    pub fn iter_entries<'a, T>(
357        &'a mut self,
358        map: BTreeMap<u32, T>,
359    ) -> impl Iterator<Item = (Result<&'a mut dyn Any, ResourceTableError>, T)> {
360        map.into_iter().map(move |(k, v)| {
361            let item = self
362                .occupied_mut(k)
363                .map(|e| Box::as_mut(&mut e.entry))
364                // Safety: extending the lifetime of the mutable reference.
365                .map(|item| unsafe { &mut *(item as *mut dyn Any) });
366            (item, v)
367        })
368    }
369
370    /// Iterate over all children belonging to the provided parent
371    pub fn iter_children<T>(
372        &self,
373        parent: &Resource<T>,
374    ) -> Result<impl Iterator<Item = &(dyn Any + Send)> + use<'_, T>, ResourceTableError>
375    where
376        T: 'static,
377    {
378        let parent_entry = self.occupied(parent.rep())?;
379        Ok(parent_entry.children.iter().map(|child_index| {
380            let child = self.occupied(*child_index).expect("missing child");
381            child.entry.as_ref()
382        }))
383    }
384
385    /// Iterate over all the entries in this table.
386    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut (dyn Any + Send)> {
387        self.entries.iter_mut().filter_map(|entry| match entry {
388            Entry::Occupied { entry } => Some(&mut *entry.entry),
389            Entry::Free { .. } => None,
390        })
391    }
392}
393
394impl Default for ResourceTable {
395    fn default() -> Self {
396        ResourceTable::new()
397    }
398}
399
400impl fmt::Debug for ResourceTable {
401    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
402        write!(f, "[")?;
403        let mut wrote = false;
404        for (index, entry) in self.entries.iter().enumerate() {
405            if let Entry::Occupied { entry } = entry {
406                if entry.entry.downcast_ref::<Tombstone>().is_none() {
407                    if wrote {
408                        write!(f, ", ")?;
409                    } else {
410                        wrote = true;
411                    }
412                    write!(f, "{index}")?;
413                }
414            }
415        }
416        write!(f, "]")
417    }
418}
419
420#[test]
421pub fn test_free_list() {
422    let mut table = ResourceTable::new();
423
424    let x = table.push(()).unwrap();
425    assert_eq!(x.rep(), 0);
426
427    let y = table.push(()).unwrap();
428    assert_eq!(y.rep(), 1);
429
430    // Deleting x should put it on the free list, so the next entry should have the same rep.
431    table.delete_maybe_debug(x, false).unwrap();
432    let x = table.push(()).unwrap();
433    assert_eq!(x.rep(), 0);
434
435    // Deleting x and then y should yield indices 1 and then 0 for new entries.
436    table.delete_maybe_debug(x, false).unwrap();
437    table.delete_maybe_debug(y, false).unwrap();
438
439    let y = table.push(()).unwrap();
440    assert_eq!(y.rep(), 1);
441
442    let x = table.push(()).unwrap();
443    assert_eq!(x.rep(), 0);
444
445    // As the free list is empty, this entry will have a new id.
446    let x = table.push(()).unwrap();
447    assert_eq!(x.rep(), 2);
448}
wasmtime/runtime/component/resource_table.rs

wasmtime/runtime/component/
resource_table.rs
