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

wasmtime/runtime/component/
resource_table.rs
