wasmtime/runtime/vm/component/
handle_table.rs

1use super::{TypedResource, TypedResourceIndex};
2use alloc::vec::Vec;
3use anyhow::{Result, bail};
4use core::mem;
5use wasmtime_environ::component::{TypeFutureTableIndex, TypeStreamTableIndex};
6
7/// The maximum handle value is specified in
8/// <https://github.com/WebAssembly/component-model/blob/main/design/mvp/CanonicalABI.md>
9/// currently and keeps the upper bits free for use in the component and ABI.
10const MAX_HANDLE: u32 = 1 << 28;
11
12/// Represents the state of a stream or future handle from the perspective of a
13/// given component instance.
14#[derive(Debug, Eq, PartialEq)]
15pub enum TransmitLocalState {
16    /// The write end of the stream or future.
17    Write {
18        /// Whether the component instance has been notified that the stream or
19        /// future is "done" (i.e. the other end has dropped, or, in the case of
20        /// a future, a value has been transmitted).
21        done: bool,
22    },
23    /// The read end of the stream or future.
24    Read {
25        /// Whether the component instance has been notified that the stream or
26        /// future is "done" (i.e. the other end has dropped, or, in the case of
27        /// a future, a value has been transmitted).
28        done: bool,
29    },
30    /// A read or write is in progress.
31    Busy,
32}
33
34/// Return value from [`HandleTable::remove_resource`].
35pub enum RemovedResource {
36    /// An `own` resource was removed with the specified `rep`
37    Own { rep: u32 },
38    /// A `borrow` resource was removed originally created within `scope`.
39    Borrow { scope: usize },
40}
41
42/// Different kinds of waitables returned by [`HandleTable::waitable_rep`].
43pub enum Waitable {
44    Subtask { is_host: bool },
45    Future,
46    Stream,
47}
48
49enum Slot {
50    Free {
51        next: u32,
52    },
53
54    /// Represents an owned resource handle with the listed representation.
55    ///
56    /// The `lend_count` tracks how many times this has been lent out as a
57    /// `borrow` and if nonzero this can't be removed.
58    ResourceOwn {
59        resource: TypedResource,
60        lend_count: u32,
61    },
62
63    /// Represents a borrowed resource handle connected to the `scope`
64    /// provided.
65    ///
66    /// The `rep` is listed and dropping this borrow will decrement the borrow
67    /// count of the `scope`.
68    ResourceBorrow {
69        resource: TypedResource,
70        scope: usize,
71    },
72
73    /// Represents a host task handle.
74    HostTask {
75        rep: u32,
76    },
77
78    /// Represents a guest task handle.
79    GuestTask {
80        rep: u32,
81    },
82
83    /// Represents a stream handle.
84    Stream {
85        ty: TypeStreamTableIndex,
86        rep: u32,
87        state: TransmitLocalState,
88    },
89
90    /// Represents a future handle.
91    Future {
92        ty: TypeFutureTableIndex,
93        rep: u32,
94        state: TransmitLocalState,
95    },
96
97    /// Represents a waitable-set handle.
98    WaitableSet {
99        rep: u32,
100    },
101
102    /// Represents an error-context handle.
103    ErrorContext {
104        rep: u32,
105    },
106}
107
108pub struct HandleTable {
109    next: u32,
110    slots: Vec<Slot>,
111}
112
113impl Default for HandleTable {
114    fn default() -> Self {
115        Self {
116            next: 0,
117            slots: Vec::new(),
118        }
119    }
120}
121
122impl HandleTable {
123    /// Returns whether or not this table is empty.
124    pub fn is_empty(&self) -> bool {
125        self.slots
126            .iter()
127            .all(|slot| matches!(slot, Slot::Free { .. }))
128    }
129
130    fn insert(&mut self, slot: Slot) -> Result<u32> {
131        let next = self.next as usize;
132        if next == self.slots.len() {
133            self.slots.push(Slot::Free {
134                next: self.next.checked_add(1).unwrap(),
135            });
136        }
137        let ret = self.next;
138        self.next = match mem::replace(&mut self.slots[next], slot) {
139            Slot::Free { next } => next,
140            _ => unreachable!(),
141        };
142        // The component model reserves index 0 as never allocatable so add one
143        // to the table index to start the numbering at 1 instead. Also note
144        // that the component model places an upper-limit per-table on the
145        // maximum allowed index.
146        let ret = ret + 1;
147        if ret >= MAX_HANDLE {
148            bail!("cannot allocate another handle: index overflow");
149        }
150
151        Ok(ret)
152    }
153
154    fn remove(&mut self, idx: u32) -> Result<()> {
155        let to_fill = Slot::Free { next: self.next };
156        let slot = self.get_mut(idx)?;
157        *slot = to_fill;
158        self.next = idx - 1;
159        Ok(())
160    }
161
162    fn handle_index_to_table_index(&self, idx: u32) -> Option<usize> {
163        // NB: `idx` is decremented by one to account for the `+1` above during
164        // allocation.
165        let idx = idx.checked_sub(1)?;
166        usize::try_from(idx).ok()
167    }
168
169    fn get_mut(&mut self, idx: u32) -> Result<&mut Slot> {
170        let slot = self
171            .handle_index_to_table_index(idx)
172            .and_then(|i| self.slots.get_mut(i));
173        match slot {
174            None | Some(Slot::Free { .. }) => bail!("unknown handle index {idx}"),
175            Some(slot) => Ok(slot),
176        }
177    }
178
179    /// Inserts a new `own` resource into this table whose type/rep are
180    /// specified by `resource`.
181    pub fn resource_own_insert(&mut self, resource: TypedResource) -> Result<u32> {
182        self.insert(Slot::ResourceOwn {
183            resource,
184            lend_count: 0,
185        })
186    }
187
188    /// Inserts a new `borrow` resource into this table whose type/rep are
189    /// specified by `resource`. The `scope` specified is used by
190    /// `CallContexts` to manage lending information.
191    pub fn resource_borrow_insert(&mut self, resource: TypedResource, scope: usize) -> Result<u32> {
192        self.insert(Slot::ResourceBorrow { resource, scope })
193    }
194
195    /// Returns the internal "rep" of the resource specified by `idx`.
196    ///
197    /// Returns an error if `idx` is out-of-bounds or doesn't point to a
198    /// resource of the appropriate type.
199    pub fn resource_rep(&mut self, idx: TypedResourceIndex) -> Result<u32> {
200        match self.get_mut(idx.raw_index())? {
201            Slot::ResourceOwn { resource, .. } | Slot::ResourceBorrow { resource, .. } => {
202                resource.rep(&idx)
203            }
204            _ => bail!("index is not a resource"),
205        }
206    }
207
208    /// Accesses the "rep" of the resource pointed to by `idx` as part of a
209    /// lending operation.
210    ///
211    /// This will increase `lend_count` for owned resources and must be paired
212    /// with a `resource_undo_lend` below later on (managed by `CallContexts`).
213    ///
214    /// Upon success returns the "rep" plus whether the borrow came from an
215    /// `own` handle.
216    pub fn resource_lend(&mut self, idx: TypedResourceIndex) -> Result<(u32, bool)> {
217        match self.get_mut(idx.raw_index())? {
218            Slot::ResourceOwn {
219                resource,
220                lend_count,
221            } => {
222                let rep = resource.rep(&idx)?;
223                *lend_count = lend_count.checked_add(1).unwrap();
224                Ok((rep, true))
225            }
226            Slot::ResourceBorrow { resource, .. } => Ok((resource.rep(&idx)?, false)),
227            _ => bail!("index {} is not a resource", idx.raw_index()),
228        }
229    }
230
231    /// For `own` resources that were borrowed in `resource_lend`, undoes the
232    /// lending operation.
233    pub fn resource_undo_lend(&mut self, idx: TypedResourceIndex) -> Result<()> {
234        match self.get_mut(idx.raw_index())? {
235            Slot::ResourceOwn { lend_count, .. } => {
236                *lend_count -= 1;
237                Ok(())
238            }
239            _ => bail!("index {} is not an own resource", idx.raw_index()),
240        }
241    }
242
243    /// Removes the resource specified by `idx` from the table.
244    ///
245    /// This can fail if `idx` doesn't point to a resource, points to a
246    /// borrowed resource, or points to a resource of the wrong type.
247    pub fn remove_resource(&mut self, idx: TypedResourceIndex) -> Result<RemovedResource> {
248        let ret = match self.get_mut(idx.raw_index())? {
249            Slot::ResourceOwn {
250                resource,
251                lend_count,
252            } => {
253                if *lend_count != 0 {
254                    bail!("cannot remove owned resource while borrowed")
255                }
256                RemovedResource::Own {
257                    rep: resource.rep(&idx)?,
258                }
259            }
260            Slot::ResourceBorrow { resource, scope } => {
261                // Ensure the drop is done with the right type
262                resource.rep(&idx)?;
263                RemovedResource::Borrow { scope: *scope }
264            }
265            _ => bail!("index {} is not a resource", idx.raw_index()),
266        };
267        self.remove(idx.raw_index())?;
268        Ok(ret)
269    }
270
271    /// Inserts a readable-end stream of type `ty` and with the specified `rep`
272    /// into this table.
273    ///
274    /// Returns the table-local index of the stream.
275    pub fn stream_insert_read(&mut self, ty: TypeStreamTableIndex, rep: u32) -> Result<u32> {
276        self.insert(Slot::Stream {
277            rep,
278            ty,
279            state: TransmitLocalState::Read { done: false },
280        })
281    }
282
283    /// Inserts a writable-end stream of type `ty` and with the specified `rep`
284    /// into this table.
285    ///
286    /// Returns the table-local index of the stream.
287    pub fn stream_insert_write(&mut self, ty: TypeStreamTableIndex, rep: u32) -> Result<u32> {
288        self.insert(Slot::Stream {
289            rep,
290            ty,
291            state: TransmitLocalState::Write { done: false },
292        })
293    }
294
295    /// Returns the `rep` and `state` associated with the stream pointed to by
296    /// `idx`.
297    ///
298    /// Returns an error if `idx` is out-of-bounds or doesn't point to a stream
299    /// of type `ty`.
300    pub fn stream_rep(
301        &mut self,
302        expected_ty: TypeStreamTableIndex,
303        idx: u32,
304    ) -> Result<(u32, &mut TransmitLocalState)> {
305        match self.get_mut(idx)? {
306            Slot::Stream { rep, ty, state } => {
307                if *ty != expected_ty {
308                    bail!("handle is a stream of a different type");
309                }
310                Ok((*rep, state))
311            }
312            _ => bail!("handle is not a stream"),
313        }
314    }
315
316    /// Removes the stream handle from `idx`, returning its `rep`.
317    ///
318    /// The stream must have the type `ty` and additionally be in a state
319    /// suitable for removal.
320    ///
321    /// Returns the `rep` for the stream along with whether the stream was
322    /// "done" or the writable end was witnessed as being done.
323    pub fn stream_remove_readable(
324        &mut self,
325        expected_ty: TypeStreamTableIndex,
326        idx: u32,
327    ) -> Result<(u32, bool)> {
328        let ret = match self.get_mut(idx)? {
329            Slot::Stream { rep, ty, state } => {
330                if *ty != expected_ty {
331                    bail!("handle is a stream of a different type");
332                }
333                let is_done = match state {
334                    TransmitLocalState::Read { done } => *done,
335                    TransmitLocalState::Write { .. } => {
336                        bail!("handle is not a readable end of a stream")
337                    }
338                    TransmitLocalState::Busy => bail!("cannot remove busy stream"),
339                };
340                (*rep, is_done)
341            }
342            _ => bail!("handle is not a stream"),
343        };
344        self.remove(idx)?;
345        Ok(ret)
346    }
347
348    /// Removes the writable stream handle from `idx`, returning its `rep`.
349    pub fn stream_remove_writable(
350        &mut self,
351        expected_ty: TypeStreamTableIndex,
352        idx: u32,
353    ) -> Result<u32> {
354        let ret = match self.get_mut(idx)? {
355            Slot::Stream { rep, ty, state } => {
356                if *ty != expected_ty {
357                    bail!("handle is a stream of a different type");
358                }
359                match state {
360                    TransmitLocalState::Write { .. } => {}
361                    TransmitLocalState::Read { .. } => {
362                        bail!("passed read end to `stream.drop-writable`")
363                    }
364                    TransmitLocalState::Busy => bail!("cannot drop busy stream"),
365                }
366                *rep
367            }
368            _ => bail!("handle is not a stream"),
369        };
370        self.remove(idx)?;
371        Ok(ret)
372    }
373
374    /// Inserts a readable-end future of type `ty` and with the specified `rep`
375    /// into this table.
376    ///
377    /// Returns the table-local index of the future.
378    pub fn future_insert_read(&mut self, ty: TypeFutureTableIndex, rep: u32) -> Result<u32> {
379        self.insert(Slot::Future {
380            rep,
381            ty,
382            state: TransmitLocalState::Read { done: false },
383        })
384    }
385
386    /// Inserts a writable-end future of type `ty` and with the specified `rep`
387    /// into this table.
388    ///
389    /// Returns the table-local index of the future.
390    pub fn future_insert_write(&mut self, ty: TypeFutureTableIndex, rep: u32) -> Result<u32> {
391        self.insert(Slot::Future {
392            rep,
393            ty,
394            state: TransmitLocalState::Write { done: false },
395        })
396    }
397
398    /// Returns the `rep` and `state` associated with the future pointed to by
399    /// `idx`.
400    ///
401    /// Returns an error if `idx` is out-of-bounds or doesn't point to a future
402    /// of type `ty`.
403    pub fn future_rep(
404        &mut self,
405        expected_ty: TypeFutureTableIndex,
406        idx: u32,
407    ) -> Result<(u32, &mut TransmitLocalState)> {
408        match self.get_mut(idx)? {
409            Slot::Future { rep, ty, state } => {
410                if *ty != expected_ty {
411                    bail!("handle is a future of a different type");
412                }
413                Ok((*rep, state))
414            }
415            _ => bail!("handle is not a future"),
416        }
417    }
418
419    /// Removes the future handle from `idx`, returning its `rep`.
420    ///
421    /// The future must have the type `ty` and additionally be in a state
422    /// suitable for removal.
423    ///
424    /// Returns the `rep` for the future along with whether the future was
425    /// "done" or the writable end was witnessed as being done.
426    pub fn future_remove_readable(
427        &mut self,
428        expected_ty: TypeFutureTableIndex,
429        idx: u32,
430    ) -> Result<(u32, bool)> {
431        let ret = match self.get_mut(idx)? {
432            Slot::Future { rep, ty, state } => {
433                if *ty != expected_ty {
434                    bail!("handle is a future of a different type");
435                }
436                let is_done = match state {
437                    TransmitLocalState::Read { done } => *done,
438                    TransmitLocalState::Write { .. } => {
439                        bail!("handle is not a readable end of a future")
440                    }
441                    TransmitLocalState::Busy => bail!("cannot remove busy future"),
442                };
443                (*rep, is_done)
444            }
445            _ => bail!("handle is not a future"),
446        };
447        self.remove(idx)?;
448        Ok(ret)
449    }
450
451    /// Removes the writable future handle from `idx`, returning its `rep`.
452    pub fn future_remove_writable(
453        &mut self,
454        expected_ty: TypeFutureTableIndex,
455        idx: u32,
456    ) -> Result<u32> {
457        let ret = match self.get_mut(idx)? {
458            Slot::Future { rep, ty, state } => {
459                if *ty != expected_ty {
460                    bail!("handle is a future of a different type");
461                }
462                match state {
463                    TransmitLocalState::Write { .. } => {}
464                    TransmitLocalState::Read { .. } => {
465                        bail!("passed read end to `future.drop-writable`")
466                    }
467                    TransmitLocalState::Busy => bail!("cannot drop busy future"),
468                }
469                *rep
470            }
471            _ => bail!("handle is not a future"),
472        };
473        self.remove(idx)?;
474        Ok(ret)
475    }
476
477    /// Inserts the error-context `rep` into this table, returning the index it
478    /// now resides at.
479    pub fn error_context_insert(&mut self, rep: u32) -> Result<u32> {
480        self.insert(Slot::ErrorContext { rep })
481    }
482
483    /// Returns the `rep` of an error-context pointed to by `idx`.
484    pub fn error_context_rep(&mut self, idx: u32) -> Result<u32> {
485        match self.get_mut(idx)? {
486            Slot::ErrorContext { rep } => Ok(*rep),
487            _ => bail!("handle is not an error-context"),
488        }
489    }
490
491    /// Drops the error-context pointed to by `idx`.
492    ///
493    /// Returns the internal `rep`.
494    pub fn error_context_drop(&mut self, idx: u32) -> Result<u32> {
495        let rep = match self.get_mut(idx)? {
496            Slot::ErrorContext { rep } => *rep,
497            _ => bail!("handle is not an error-context"),
498        };
499        self.remove(idx)?;
500        Ok(rep)
501    }
502
503    /// Inserts `rep` as a guest subtask into this table.
504    pub fn subtask_insert_guest(&mut self, rep: u32) -> Result<u32> {
505        self.insert(Slot::GuestTask { rep })
506    }
507
508    /// Inserts `rep` as a host subtask into this table.
509    pub fn subtask_insert_host(&mut self, rep: u32) -> Result<u32> {
510        self.insert(Slot::HostTask { rep })
511    }
512
513    /// Returns the `rep` of the subtask at `idx` as well as if it's a host
514    /// task or not.
515    pub fn subtask_rep(&mut self, idx: u32) -> Result<(u32, bool)> {
516        match self.get_mut(idx)? {
517            Slot::GuestTask { rep } => Ok((*rep, false)),
518            Slot::HostTask { rep } => Ok((*rep, true)),
519            _ => bail!("handle is not a subtask"),
520        }
521    }
522
523    /// Removes the subtask set at `idx`, returning its `rep`.
524    pub fn subtask_remove(&mut self, idx: u32) -> Result<(u32, bool)> {
525        let ret = match self.get_mut(idx)? {
526            Slot::GuestTask { rep } => (*rep, false),
527            Slot::HostTask { rep } => (*rep, true),
528            _ => bail!("handle is not a subtask"),
529        };
530        self.remove(idx)?;
531        Ok(ret)
532    }
533
534    /// Inserts `rep` as a waitable set into this table.
535    pub fn waitable_set_insert(&mut self, rep: u32) -> Result<u32> {
536        self.insert(Slot::WaitableSet { rep })
537    }
538
539    /// Returns the `rep` of an waitable-set pointed to by `idx`.
540    pub fn waitable_set_rep(&mut self, idx: u32) -> Result<u32> {
541        match self.get_mut(idx)? {
542            Slot::WaitableSet { rep, .. } => Ok(*rep),
543            _ => bail!("handle is not an waitable-set"),
544        }
545    }
546
547    /// Removes the waitable set at `idx`, returning its `rep`.
548    pub fn waitable_set_remove(&mut self, idx: u32) -> Result<u32> {
549        let ret = match self.get_mut(idx)? {
550            Slot::WaitableSet { rep } => *rep,
551            _ => bail!("handle is not a waitable-set"),
552        };
553        self.remove(idx)?;
554        Ok(ret)
555    }
556
557    /// Returns the `rep` for the waitable specified by `idx` along with what
558    /// kind of waitable it is.
559    pub fn waitable_rep(&mut self, idx: u32) -> Result<(u32, Waitable)> {
560        match self.get_mut(idx)? {
561            Slot::GuestTask { rep } => Ok((*rep, Waitable::Subtask { is_host: false })),
562            Slot::HostTask { rep } => Ok((*rep, Waitable::Subtask { is_host: true })),
563            Slot::Future { rep, .. } => Ok((*rep, Waitable::Future)),
564            Slot::Stream { rep, .. } => Ok((*rep, Waitable::Stream)),
565            _ => bail!("handle is not a waitable"),
566        }
567    }
568}