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::marker::PhantomData;
7use core::mem;
8use sealed::{DefaultTombstone, DeleteWithTombstone};
9
10#[derive(Debug)]
11/// Errors returned by operations on `ResourceTable`
12pub enum ResourceTableError {
13    /// ResourceTable has no free keys
14    Full,
15    /// Resource not present in table
16    NotPresent,
17    /// Resource present in table, but with a different type
18    WrongType,
19    /// Resource cannot be deleted because child resources exist in the table. Consult wit docs for
20    /// the particular resource to see which methods may return child resources.
21    HasChildren,
22}
23
24impl fmt::Display for ResourceTableError {
25    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
26        match self {
27            Self::Full => write!(f, "resource table has no free keys"),
28            Self::NotPresent => write!(f, "resource not present"),
29            Self::WrongType => write!(f, "resource is of another type"),
30            Self::HasChildren => write!(f, "resource has children"),
31        }
32    }
33}
34
35impl core::error::Error for ResourceTableError {}
36
37/// The `ResourceTable` type maps a `Resource<T>` to its `T`.
38pub struct ResourceTable<D: DeleteWithTombstone = DefaultTombstone> {
39    entries: Vec<Entry>,
40    free_head: Option<usize>,
41    // Number of reusable slots on the free list, excluding debug tombstones.
42    free_count: usize,
43    max_capacity: usize,
44    live: usize,
45    _tombstone: PhantomData<D>,
46}
47
48#[derive(Debug)]
49enum Entry {
50    Free { next: Option<usize> },
51    Occupied { entry: TableEntry },
52}
53
54impl Entry {
55    pub fn occupied(&self) -> Option<&TableEntry> {
56        match self {
57            Self::Occupied { entry } => Some(entry),
58            Self::Free { .. } => None,
59        }
60    }
61
62    pub fn occupied_mut(&mut self) -> Option<&mut TableEntry> {
63        match self {
64            Self::Occupied { entry } => Some(entry),
65            Self::Free { .. } => None,
66        }
67    }
68}
69
70struct Tombstone;
71
72mod sealed {
73    /// Configures whether `ResourceTable` leaves tombstones on delete.
74    ///
75    /// This trait is public but unnameable outside of this module, so it can't
76    /// be implemented elsewhere.
77    pub trait DeleteWithTombstone {
78        const DELETE_WITH_TOMBSTONE: bool;
79    }
80
81    pub struct DefaultTombstone;
82
83    impl DeleteWithTombstone for DefaultTombstone {
84        // Change this to `true` to assist with handle debugging in development if
85        // necessary.
86        const DELETE_WITH_TOMBSTONE: bool = false;
87    }
88}
89
90/// Default setting for `ResourceTable::max_capacity`, chosen to be high
91/// enough that it doesn't need changing all that often but low enough that
92/// exhausting it isn't a massive problem for the host.
93const DEFAULT_MAX_CAPACITY: usize = 1_000_000;
94
95/// This structure tracks parent and child relationships for a given table entry.
96///
97/// Parents and children are referred to by table index. We maintain the
98/// following invariants:
99/// * the parent must exist when adding a child.
100/// * whenever a child is created, its index is added to children.
101/// * whenever a child is deleted, its index is removed from children.
102/// * an entry with children may not be deleted.
103#[derive(Debug)]
104struct TableEntry {
105    /// The entry in the table, as a boxed dynamically-typed object
106    entry: Box<dyn Any + Send>,
107    /// The index of the parent of this entry, if it has one.
108    parent: Option<u32>,
109    /// The indices of any children of this entry.
110    children: BTreeSet<u32>,
111}
112
113impl TableEntry {
114    fn new(entry: Box<dyn Any + Send>, parent: Option<u32>) -> Self {
115        Self {
116            entry,
117            parent,
118            children: BTreeSet::new(),
119        }
120    }
121    fn add_child(&mut self, child: u32) {
122        debug_assert!(!self.children.contains(&child));
123        self.children.insert(child);
124    }
125    fn remove_child(&mut self, child: u32) {
126        let was_removed = self.children.remove(&child);
127        debug_assert!(was_removed);
128    }
129}
130
131impl ResourceTable {
132    /// Create an empty table
133    pub fn new() -> Self {
134        ResourceTable::with_capacity(0)
135    }
136
137    /// Create an empty table with at least the specified capacity.
138    pub fn with_capacity(capacity: usize) -> Self {
139        ResourceTable::empty(capacity)
140    }
141}
142
143impl<D: DeleteWithTombstone> ResourceTable<D> {
144    fn empty(capacity: usize) -> Self {
145        ResourceTable {
146            entries: Vec::with_capacity(capacity),
147            free_head: None,
148            free_count: 0,
149            max_capacity: DEFAULT_MAX_CAPACITY,
150            live: 0,
151            _tombstone: PhantomData,
152        }
153    }
154
155    /// Returns the number of live resources in this table.
156    ///
157    /// Deleted resources, including debug tombstones, are not counted. This is
158    /// an `O(1)` operation.
159    pub fn len(&self) -> usize {
160        self.live
161    }
162
163    /// Returns whether or not this table is empty in `O(1)` time.
164    pub fn is_empty(&self) -> bool {
165        self.len() == 0
166    }
167
168    /// Checks that a batch of pushes has enough free or appendable slots.
169    ///
170    /// Free slots are reused even if the maximum capacity has been lowered below
171    /// the number of backing slots.
172    #[cfg(any(feature = "component-model-async", test))]
173    pub(crate) fn ensure_capacity(&self, additional: usize) -> Result<(), ResourceTableError> {
174        let Some(to_append) = additional.checked_sub(self.free_count) else {
175            return Ok(());
176        };
177        if to_append > self.max_capacity.saturating_sub(self.entries.len()) {
178            Err(ResourceTableError::Full)
179        } else {
180            Ok(())
181        }
182    }
183
184    /// Returns the maximum capacity of this table, in elements, before adding
185    /// any more will be refused.
186    pub fn max_capacity(&self) -> usize {
187        self.max_capacity
188    }
189
190    /// Configures the maximum number of entries that may be present within this
191    /// table.
192    ///
193    /// Note that this does not retroactively shrink the table nor evict
194    /// existing entries should the maximum be smaller than the current size of
195    /// the entry table.
196    pub fn set_max_capacity(&mut self, max: usize) {
197        self.max_capacity = max;
198    }
199
200    /// Inserts a new value `T` into this table, returning a corresponding
201    /// `Resource<T>` which can be used to refer to it after it was inserted.
202    pub fn push<T>(&mut self, entry: T) -> Result<Resource<T>, ResourceTableError>
203    where
204        T: Send + 'static,
205    {
206        let idx = self.push_(TableEntry::new(Box::new(entry), None))?;
207        Ok(Resource::new_own(idx))
208    }
209
210    /// Pop an index off of the free list, if it's not empty.
211    fn pop_free_list(&mut self) -> Option<usize> {
212        if let Some(ix) = self.free_head {
213            // Advance free_head to the next entry if one is available.
214            match &self.entries[ix] {
215                Entry::Free { next } => self.free_head = *next,
216                Entry::Occupied { .. } => unreachable!(),
217            }
218            self.free_count -= 1;
219            Some(ix)
220        } else {
221            None
222        }
223    }
224
225    /// Free an entry in the table, returning its [`TableEntry`]. Add the index to the free list.
226    #[inline]
227    fn free_entry(&mut self, ix: usize) -> TableEntry {
228        if D::DELETE_WITH_TOMBSTONE {
229            // Instead of making this entry available for reuse, we leave a
230            // tombstone in debug mode.  This helps detect use-after-delete and
231            // double-delete bugs.
232            match mem::replace(
233                &mut self.entries[ix],
234                Entry::Occupied {
235                    entry: TableEntry {
236                        entry: Box::new(Tombstone),
237                        parent: None,
238                        children: BTreeSet::new(),
239                    },
240                },
241            ) {
242                Entry::Occupied { entry } => {
243                    if !entry.entry.is::<Tombstone>() {
244                        self.live -= 1;
245                    }
246                    entry
247                }
248                Entry::Free { .. } => unreachable!(),
249            }
250        } else {
251            let entry = match core::mem::replace(
252                &mut self.entries[ix],
253                Entry::Free {
254                    next: self.free_head,
255                },
256            ) {
257                Entry::Occupied { entry } => entry,
258                Entry::Free { .. } => unreachable!(),
259            };
260
261            self.free_head = Some(ix);
262            self.free_count += 1;
263            self.live -= 1;
264
265            entry
266        }
267    }
268
269    /// Push a new entry into the table, returning its handle. This will prefer to use free entries
270    /// if they exist, falling back on pushing new entries onto the end of the table.
271    fn push_(&mut self, e: TableEntry) -> Result<u32, ResourceTableError> {
272        if let Some(free) = self.pop_free_list() {
273            self.entries[free] = Entry::Occupied { entry: e };
274            self.live += 1;
275            Ok(free.try_into().unwrap())
276        } else {
277            if self.entries.len() >= self.max_capacity {
278                return Err(ResourceTableError::Full);
279            }
280            let ix = self
281                .entries
282                .len()
283                .try_into()
284                .map_err(|_| ResourceTableError::Full)?;
285            self.entries.push(Entry::Occupied { entry: e });
286            self.live += 1;
287            Ok(ix)
288        }
289    }
290
291    fn occupied(&self, key: u32) -> Result<&TableEntry, ResourceTableError> {
292        self.entries
293            .get(key as usize)
294            .and_then(Entry::occupied)
295            .ok_or(ResourceTableError::NotPresent)
296    }
297
298    fn occupied_mut(&mut self, key: u32) -> Result<&mut TableEntry, ResourceTableError> {
299        self.entries
300            .get_mut(key as usize)
301            .and_then(Entry::occupied_mut)
302            .ok_or(ResourceTableError::NotPresent)
303    }
304
305    /// Insert a resource at the next available index, and track that it has a
306    /// parent resource.
307    ///
308    /// The parent must exist to create a child. All children resources must
309    /// be destroyed before a parent can be destroyed - otherwise
310    /// [`ResourceTable::delete`] will fail with
311    /// [`ResourceTableError::HasChildren`].
312    ///
313    /// Parent-child relationships are tracked inside the table to ensure that
314    /// a parent resource is not deleted while it has live children. This
315    /// allows child resources to hold "references" to a parent by table
316    /// index, to avoid needing e.g. an `Arc<Mutex<parent>>` and the associated
317    /// locking overhead and design issues, such as child existence extending
318    /// lifetime of parent referent even after parent resource is destroyed,
319    /// possibility for deadlocks.
320    pub fn push_child<T, U>(
321        &mut self,
322        entry: T,
323        parent: &Resource<U>,
324    ) -> Result<Resource<T>, ResourceTableError>
325    where
326        T: Send + 'static,
327        U: 'static,
328    {
329        let parent = parent.rep();
330        self.occupied(parent)?;
331        let child = self.push_(TableEntry::new(Box::new(entry), Some(parent)))?;
332        self.occupied_mut(parent)?.add_child(child);
333        Ok(Resource::new_own(child))
334    }
335
336    /// Add an already-resident child to a resource.
337    pub fn add_child<T: 'static, U: 'static>(
338        &mut self,
339        child: Resource<T>,
340        parent: Resource<U>,
341    ) -> Result<(), ResourceTableError> {
342        let entry = self.occupied_mut(child.rep())?;
343        assert!(entry.parent.is_none());
344        entry.parent = Some(parent.rep());
345        self.occupied_mut(parent.rep())?.add_child(child.rep());
346        Ok(())
347    }
348
349    /// Remove a child to from a resource (but leave it in the table).
350    pub fn remove_child<T: 'static, U: 'static>(
351        &mut self,
352        child: Resource<T>,
353        parent: Resource<U>,
354    ) -> Result<(), ResourceTableError> {
355        let entry = self.occupied_mut(child.rep())?;
356        assert_eq!(entry.parent, Some(parent.rep()));
357        entry.parent = None;
358        self.occupied_mut(parent.rep())?.remove_child(child.rep());
359        Ok(())
360    }
361
362    /// Get an immutable reference to a resource of a given type at a given
363    /// index.
364    ///
365    /// Multiple shared references can be borrowed at any given time.
366    pub fn get<T: Any + Sized>(&self, key: &Resource<T>) -> Result<&T, ResourceTableError> {
367        self.get_(key.rep())?
368            .downcast_ref()
369            .ok_or(ResourceTableError::WrongType)
370    }
371
372    fn get_(&self, key: u32) -> Result<&dyn Any, ResourceTableError> {
373        let r = self.occupied(key)?;
374        Ok(&*r.entry)
375    }
376
377    /// Get an mutable reference to a resource of a given type at a given
378    /// index.
379    pub fn get_mut<T: Any + Sized>(
380        &mut self,
381        key: &Resource<T>,
382    ) -> Result<&mut T, ResourceTableError> {
383        self.get_any_mut(key.rep())?
384            .downcast_mut()
385            .ok_or(ResourceTableError::WrongType)
386    }
387
388    /// Returns the raw `Any` at the `key` index provided.
389    pub fn get_any_mut(&mut self, key: u32) -> Result<&mut dyn Any, ResourceTableError> {
390        let r = self.occupied_mut(key)?;
391        Ok(&mut *r.entry)
392    }
393
394    /// Remove the specified entry from the table.
395    pub fn delete<T>(&mut self, resource: Resource<T>) -> Result<T, ResourceTableError>
396    where
397        T: Any,
398    {
399        debug_assert!(resource.owned());
400        let key = resource.rep();
401        if !self.occupied(key)?.children.is_empty() {
402            return Err(ResourceTableError::HasChildren);
403        }
404        let entry = self.free_entry(key as usize);
405        if let Some(parent) = entry.parent {
406            // Remove deleted resource from parent's child list.
407            // Parent must still be present because it can't be deleted while still having
408            // children:
409            self.occupied_mut(parent)
410                .expect("missing parent")
411                .remove_child(key);
412        }
413        match entry.entry.downcast() {
414            Ok(t) => Ok(*t),
415            Err(_e) => Err(ResourceTableError::WrongType),
416        }
417    }
418
419    /// Zip the values of the map with mutable references to table entries corresponding to each
420    /// key. As the keys in the `BTreeMap` are unique, this iterator can give mutable references
421    /// with the same lifetime as the mutable reference to the [ResourceTable].
422    pub fn iter_entries<'a, T>(
423        &'a mut self,
424        map: BTreeMap<u32, T>,
425    ) -> impl Iterator<Item = (Result<&'a mut dyn Any, ResourceTableError>, T)> {
426        map.into_iter().map(move |(k, v)| {
427            let item = self
428                .occupied_mut(k)
429                .map(|e| Box::as_mut(&mut e.entry))
430                // Safety: extending the lifetime of the mutable reference.
431                .map(|item| unsafe { &mut *(item as *mut dyn Any) });
432            (item, v)
433        })
434    }
435
436    /// Iterate over all children belonging to the provided parent
437    pub fn iter_children<T>(
438        &self,
439        parent: &Resource<T>,
440    ) -> Result<impl Iterator<Item = &(dyn Any + Send)> + use<'_, T, D>, ResourceTableError>
441    where
442        T: 'static,
443    {
444        let parent_entry = self.occupied(parent.rep())?;
445        Ok(parent_entry.children.iter().map(|child_index| {
446            let child = self.occupied(*child_index).expect("missing child");
447            child.entry.as_ref()
448        }))
449    }
450
451    /// Iterate over all the entries in this table.
452    pub fn iter_mut(&mut self) -> impl Iterator<Item = (u32, &mut (dyn Any + Send))> {
453        self.entries
454            .iter_mut()
455            .enumerate()
456            .filter_map(|(rep, entry)| match entry {
457                Entry::Occupied { entry } => Some((u32::try_from(rep).unwrap(), &mut *entry.entry)),
458                Entry::Free { .. } => None,
459            })
460    }
461}
462
463impl Default for ResourceTable {
464    fn default() -> Self {
465        ResourceTable::new()
466    }
467}
468
469impl<D: DeleteWithTombstone> fmt::Debug for ResourceTable<D> {
470    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
471        write!(f, "[")?;
472        let mut wrote = false;
473        for (index, entry) in self.entries.iter().enumerate() {
474            if let Entry::Occupied { entry } = entry {
475                if entry.entry.downcast_ref::<Tombstone>().is_none() {
476                    if wrote {
477                        write!(f, ", ")?;
478                    } else {
479                        wrote = true;
480                    }
481                    write!(f, "{index}")?;
482                }
483            }
484        }
485        write!(f, "]")
486    }
487}
488
489#[test]
490pub fn test_free_list() {
491    let mut table = ResourceTable::new();
492    assert!(table.is_empty());
493
494    let x = table.push(()).unwrap();
495    assert_eq!(x.rep(), 0);
496    assert_eq!(table.len(), 1);
497
498    let y = table.push(()).unwrap();
499    assert_eq!(y.rep(), 1);
500    assert_eq!(table.len(), 2);
501
502    // Deleting x should put it on the free list, so the next entry should have the same rep.
503    table.delete(x).unwrap();
504    assert_eq!(table.len(), 1);
505    let x = table.push(()).unwrap();
506    assert_eq!(x.rep(), 0);
507    assert_eq!(table.len(), 2);
508
509    // Deleting x and then y should yield indices 1 and then 0 for new entries.
510    table.delete(x).unwrap();
511    table.delete(y).unwrap();
512
513    let y = table.push(()).unwrap();
514    assert_eq!(y.rep(), 1);
515
516    let x = table.push(()).unwrap();
517    assert_eq!(x.rep(), 0);
518
519    // As the free list is empty, this entry will have a new id.
520    let x = table.push(()).unwrap();
521    assert_eq!(x.rep(), 2);
522    assert_eq!(table.len(), 3);
523}
524
525#[test]
526fn test_max_capacity() {
527    let mut table = ResourceTable::new();
528    assert_eq!(table.max_capacity(), DEFAULT_MAX_CAPACITY);
529
530    table.set_max_capacity(0);
531    assert_eq!(table.max_capacity(), 0);
532    assert!(table.push(()).is_err());
533
534    table.set_max_capacity(1);
535    assert_eq!(table.max_capacity(), 1);
536    let x = table.push(()).unwrap();
537    assert!(table.push(()).is_err());
538    assert_eq!(table.len(), 1);
539
540    table.set_max_capacity(0);
541    assert!(table.push(()).is_err());
542    table.delete(x).unwrap();
543    let x = table.push(()).unwrap();
544    table.delete(x).unwrap();
545
546    table.set_max_capacity(10);
547
548    let handles = (0..10).map(|_| table.push(()).unwrap()).collect::<Vec<_>>();
549    assert!(table.push(()).is_err());
550    for handle in handles {
551        table.delete(handle).unwrap();
552    }
553    assert!(table.is_empty());
554
555    table.push(()).unwrap();
556}
557
558#[test]
559fn test_len_and_delete_errors() {
560    let mut table = ResourceTable::new();
561    let parent = table.push(()).unwrap();
562    let child = table.push_child(1u32, &parent).unwrap();
563    assert_eq!(table.len(), 2);
564    table.set_max_capacity(2);
565    assert!(matches!(
566        table.push_child(2u32, &parent),
567        Err(ResourceTableError::Full)
568    ));
569    assert_eq!(table.len(), 2);
570    assert!(matches!(
571        table.delete(Resource::<()>::new_own(parent.rep())),
572        Err(ResourceTableError::HasChildren)
573    ));
574    assert_eq!(table.len(), 2);
575    assert!(matches!(
576        table.push_child((), &Resource::<()>::new_own(100)),
577        Err(ResourceTableError::NotPresent)
578    ));
579    assert_eq!(table.len(), 2);
580    assert!(matches!(
581        table.delete(Resource::<()>::new_own(child.rep())),
582        Err(ResourceTableError::WrongType)
583    ));
584    assert_eq!(table.len(), 1);
585    assert!(matches!(
586        table.delete(Resource::<()>::new_own(100)),
587        Err(ResourceTableError::NotPresent)
588    ));
589    assert_eq!(table.len(), 1);
590    table.delete(parent).unwrap();
591    assert!(table.is_empty());
592}
593
594#[test]
595fn test_len_and_tombstones() {
596    struct DebugTombstone;
597    impl DeleteWithTombstone for DebugTombstone {
598        const DELETE_WITH_TOMBSTONE: bool = true;
599    }
600
601    let mut table = ResourceTable::<DebugTombstone>::empty(0);
602    table.set_max_capacity(1);
603    let value = table.push(1u32).unwrap();
604    let rep = value.rep();
605    table.delete(value).unwrap();
606    assert_eq!(table.len(), 0);
607    assert!(table.is_empty());
608    assert_eq!(table.entries.len(), 1);
609    assert!(matches!(
610        table.ensure_capacity(1),
611        Err(ResourceTableError::Full)
612    ));
613    assert!(matches!(
614        table.delete(Resource::<u32>::new_own(rep)),
615        Err(ResourceTableError::WrongType)
616    ));
617    assert_eq!(table.len(), 0);
618    assert!(table.is_empty());
619}
620
621#[test]
622fn test_batch_capacity() {
623    let mut table = ResourceTable::new();
624    table.set_max_capacity(3);
625    table.ensure_capacity(0).unwrap();
626    table.ensure_capacity(3).unwrap();
627    assert!(matches!(
628        table.ensure_capacity(4),
629        Err(ResourceTableError::Full)
630    ));
631    assert_eq!(table.entries.len(), 0);
632    let first = table.push(()).unwrap();
633    table.ensure_capacity(2).unwrap();
634    assert!(matches!(
635        table.ensure_capacity(3),
636        Err(ResourceTableError::Full)
637    ));
638    table.delete(first).unwrap();
639    // A batch can combine free slots and new slots.
640    table.ensure_capacity(3).unwrap();
641    let handles = (0..3).map(|_| table.push(()).unwrap()).collect::<Vec<_>>();
642    for handle in handles {
643        table.delete(handle).unwrap();
644    }
645    table.set_max_capacity(0);
646    // Reusing slots remains valid even after lowering the maximum capacity.
647    table.ensure_capacity(3).unwrap();
648    assert!(matches!(
649        table.ensure_capacity(4),
650        Err(ResourceTableError::Full)
651    ));
652    assert_eq!(table.len(), 0);
653    assert_eq!(table.entries.len(), 3);
654}
wasmtime/runtime/component/resource_table.rs

wasmtime/runtime/component/
resource_table.rs
