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

wasmtime/runtime/component/
resource_table.rs
