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

wasmtime/runtime/component/
resource_table.rs
