use super::table::{TableDebug, TableId}; use super::{Event, GlobalErrorContextRefCount, Waitable, WaitableCommon}; use crate::component::concurrent::{ConcurrentState, QualifiedThreadId, WorkItem, tls}; use crate::component::func::{self, LiftContext, LowerContext}; use crate::component::matching::InstanceType; use crate::component::types; use crate::component::values::ErrorContextAny; use crate::component::{ AsAccessor, ComponentInstanceId, ComponentType, FutureAny, Instance, Lift, Lower, StreamAny, Val, WasmList, }; use crate::store::{StoreOpaque, StoreToken}; use crate::vm::component::{ComponentInstance, HandleTable, TransmitLocalState}; use crate::vm::{AlwaysMut, VMStore}; use crate::{AsContext, AsContextMut, StoreContextMut, ValRaw}; use crate::{ Error, Result, Trap, bail, bail_bug, ensure, error::{Context as _, format_err}, }; use buffers::{Extender, SliceBuffer, UntypedWriteBuffer}; use core::fmt; use core::future; use core::iter; use core::marker::PhantomData; use core::mem::{self, ManuallyDrop, MaybeUninit}; use core::ops::{Deref, DerefMut}; use core::pin::Pin; use core::task::{Context, Poll, Waker, ready}; use futures::channel::oneshot; use futures::{FutureExt as _, stream}; use std::any::{Any, TypeId}; use std::boxed::Box; use std::io::Cursor; use std::string::String; use std::sync::{Arc, Mutex, MutexGuard}; use std::vec::Vec; use wasmtime_environ::component::{ CanonicalAbiInfo, ComponentTypes, InterfaceType, OptionsIndex, RuntimeComponentInstanceIndex, TypeComponentGlobalErrorContextTableIndex, TypeComponentLocalErrorContextTableIndex, TypeFutureTableIndex, TypeStreamTableIndex, }; pub use buffers::{ReadBuffer, VecBuffer, WriteBuffer}; mod buffers; /// Enum for distinguishing between a stream or future in functions that handle /// both. #[derive(Copy, Clone, Debug)] pub enum TransmitKind { Stream, Future, } /// Represents `{stream,future}.{read,write}` results. #[derive(Copy, Clone, Debug, PartialEq)] pub enum ReturnCode { Blocked, Completed(u32), Dropped(u32), Cancelled(u32), } impl ReturnCode { /// Pack `self` into a single 32-bit integer that may be returned to the /// guest. /// /// This corresponds to `pack_copy_result` in the Component Model spec. pub fn encode(&self) -> u32 { const BLOCKED: u32 = 0xffff_ffff; const COMPLETED: u32 = 0x0; const DROPPED: u32 = 0x1; const CANCELLED: u32 = 0x2; match self { ReturnCode::Blocked => BLOCKED, ReturnCode::Completed(n) => { debug_assert!(*n < (1 << 28)); (n << 4) | COMPLETED } ReturnCode::Dropped(n) => { debug_assert!(*n < (1 << 28)); (n << 4) | DROPPED } ReturnCode::Cancelled(n) => { debug_assert!(*n < (1 << 28)); (n << 4) | CANCELLED } } } /// Returns `Self::Completed` with the specified count (or zero if /// `matches!(kind, TransmitKind::Future)`) fn completed(kind: TransmitKind, count: u32) -> Self { Self::Completed(if let TransmitKind::Future = kind { 0 } else { count }) } } /// Represents a stream or future type index. /// /// This is useful as a parameter type for functions which operate on either a /// future or a stream. #[derive(Copy, Clone, Debug)] pub enum TransmitIndex { Stream(TypeStreamTableIndex), Future(TypeFutureTableIndex), } impl TransmitIndex { pub fn kind(&self) -> TransmitKind { match self { TransmitIndex::Stream(_) => TransmitKind::Stream, TransmitIndex::Future(_) => TransmitKind::Future, } } /// Retrieve the payload type of the specified stream or future, or `None` /// if it has no payload type. fn payload<'a>(&self, types: &'a ComponentTypes) -> Option<&'a InterfaceType> { match self { TransmitIndex::Stream(i) => { let ty = types[*i].ty; types[ty].payload.as_ref() } TransmitIndex::Future(i) => { let ty = types[*i].ty; types[ty].payload.as_ref() } } } } /// Retrieve the host rep and state for the specified guest-visible waitable /// handle. fn get_mut_by_index_from( handle_table: &mut HandleTable, ty: TransmitIndex, index: u32, ) -> Result<(u32, &mut TransmitLocalState)> { match ty { TransmitIndex::Stream(ty) => handle_table.stream_rep(ty, index), TransmitIndex::Future(ty) => handle_table.future_rep(ty, index), } } fn lower, U: 'static>( mut store: StoreContextMut, instance: Instance, caller_thread: QualifiedThreadId, options: OptionsIndex, ty: TransmitIndex, address: usize, count: usize, buffer: &mut B, ) -> Result<()> { let count = buffer.remaining().len().min(count); // If lowering may call realloc in the guest, then the guest may need // to access its thread context, so we need to set the current thread before lowering // and restore the old one afterward. let (lower, old_thread) = if T::MAY_REQUIRE_REALLOC { let old_thread = store.0.set_thread(caller_thread)?; ( &mut LowerContext::new(store.as_context_mut(), options, instance), Some(old_thread), ) } else { ( &mut LowerContext::new_without_realloc(store.as_context_mut(), options, instance), None, ) }; if address % usize::try_from(T::ALIGN32)? != 0 { bail!("read pointer not aligned"); } lower .as_slice_mut() .get_mut(address..) .and_then(|b| b.get_mut(..T::SIZE32 * count)) .ok_or_else(|| crate::format_err!("read pointer out of bounds of memory"))?; if let Some(ty) = ty.payload(lower.types) { T::linear_store_list_to_memory(lower, *ty, address, &buffer.remaining()[..count])?; } if let Some(old_thread) = old_thread { store.0.set_thread(old_thread)?; } buffer.skip(count); Ok(()) } fn lift>( lift: &mut LiftContext<'_>, ty: Option, buffer: &mut B, address: usize, count: usize, ) -> Result<()> { let count = count.min(buffer.remaining_capacity()); if T::IS_RUST_UNIT_TYPE { // SAFETY: `T::IS_RUST_UNIT_TYPE` is only true for `()`, a // zero-sized type, so `MaybeUninit::uninit().assume_init()` // is a valid way to populate the zero-sized buffer. buffer.extend( iter::repeat_with(|| unsafe { MaybeUninit::uninit().assume_init() }).take(count), ) } else { let ty = match ty { Some(ty) => ty, None => bail_bug!("type required for non-unit lift"), }; if address % usize::try_from(T::ALIGN32)? != 0 { bail!("write pointer not aligned"); } lift.memory() .get(address..) .and_then(|b| b.get(..T::SIZE32 * count)) .ok_or_else(|| crate::format_err!("write pointer out of bounds of memory"))?; let list = &WasmList::new(address, count, lift, ty)?; T::linear_lift_into_from_memory(lift, list, &mut Extender(buffer))? } Ok(()) } /// Represents the state associated with an error context #[derive(Debug, PartialEq, Eq, PartialOrd)] pub(super) struct ErrorContextState { /// Debug message associated with the error context pub(crate) debug_msg: String, } /// Represents the size and alignment for a "flat" Component Model type, /// i.e. one containing no pointers or handles. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(super) struct FlatAbi { pub(super) size: u32, pub(super) align: u32, } /// Represents the buffer for a host- or guest-initiated stream read. pub struct Destination<'a, T, B> { id: TableId, buffer: &'a mut B, host_buffer: Option<&'a mut Cursor>>, _phantom: PhantomData T>, } impl<'a, T, B> Destination<'a, T, B> { /// Reborrow `self` so it can be used again later. pub fn reborrow(&mut self) -> Destination<'_, T, B> { Destination { id: self.id, buffer: &mut *self.buffer, host_buffer: self.host_buffer.as_deref_mut(), _phantom: PhantomData, } } /// Take the buffer out of `self`, leaving a default-initialized one in its /// place. /// /// This can be useful for reusing the previously-stored buffer's capacity /// instead of allocating a fresh one. pub fn take_buffer(&mut self) -> B where B: Default, { mem::take(self.buffer) } /// Store the specified buffer in `self`. /// /// Any items contained in the buffer will be delivered to the reader after /// the `StreamProducer::poll_produce` call to which this `Destination` was /// passed returns (unless overwritten by another call to `set_buffer`). /// /// If items are stored via this buffer _and_ written via a /// `DirectDestination` view of `self`, then the items in the buffer will be /// delivered after the ones written using `DirectDestination`. pub fn set_buffer(&mut self, buffer: B) { *self.buffer = buffer; } /// Return the remaining number of items the current read has capacity to /// accept, if known. /// /// This will return `Some(_)` if the reader is a guest; it will return /// `None` if the reader is the host. /// /// Note that this can return `Some(0)`. This means that the guest is /// attempting to perform a zero-length read which typically means that it's /// trying to wait for this stream to be ready-to-read but is not actually /// ready to receive the items yet. The host in this case is allowed to /// either block waiting for readiness or immediately complete the /// operation. The guest is expected to handle both cases. Some more /// discussion about this case can be found in the discussion of ["Stream /// Readiness" in the component-model repo][docs]. /// /// [docs]: https://github.com/WebAssembly/component-model/blob/main/design/mvp/Concurrency.md#stream-readiness pub fn remaining(&self, mut store: impl AsContextMut) -> Option { // Note that this unwrap should only trigger for bugs in Wasmtime, and // this is modeled here to centralize the `.unwrap()` for this method in // one location. self.remaining_(store.as_context_mut().0).unwrap() } fn remaining_(&self, store: &mut StoreOpaque) -> Result> { let transmit = store.concurrent_state_mut().get_mut(self.id)?; if let &ReadState::GuestReady { count, .. } = &transmit.read { let &WriteState::HostReady { guest_offset, .. } = &transmit.write else { bail_bug!("expected WriteState::HostReady") }; Ok(Some(count - guest_offset)) } else { Ok(None) } } } impl<'a, B> Destination<'a, u8, B> { /// Return a `DirectDestination` view of `self`. /// /// If the reader is a guest, this will provide direct access to the guest's /// read buffer. If the reader is a host, this will provide access to a /// buffer which will be delivered to the host before any items stored using /// `Destination::set_buffer`. /// /// `capacity` will only be used if the reader is a host, in which case it /// will update the length of the buffer, possibly zero-initializing the new /// elements if the new length is larger than the old length. pub fn as_direct( mut self, store: StoreContextMut<'a, D>, capacity: usize, ) -> DirectDestination<'a, D> { if let Some(buffer) = self.host_buffer.as_deref_mut() { buffer.set_position(0); if buffer.get_mut().is_empty() { buffer.get_mut().resize(capacity, 0); } } DirectDestination { id: self.id, host_buffer: self.host_buffer, store, } } } /// Represents a read from a `stream`, providing direct access to the /// writer's buffer. pub struct DirectDestination<'a, D: 'static> { id: TableId, host_buffer: Option<&'a mut Cursor>>, store: StoreContextMut<'a, D>, } impl std::io::Write for DirectDestination<'_, D> { fn write(&mut self, buf: &[u8]) -> std::io::Result { let rem = self.remaining(); let n = rem.len().min(buf.len()); rem[..n].copy_from_slice(&buf[..n]); self.mark_written(n); Ok(n) } fn flush(&mut self) -> std::io::Result<()> { Ok(()) } } impl DirectDestination<'_, D> { /// Provide direct access to the writer's buffer. pub fn remaining(&mut self) -> &mut [u8] { // Note that this unwrap should only trigger for bugs in Wasmtime, and // this is modeled here to centralize the `.unwrap()` for this method in // one location. self.remaining_().unwrap() } fn remaining_(&mut self) -> Result<&mut [u8]> { if let Some(buffer) = self.host_buffer.as_deref_mut() { return Ok(buffer.get_mut()); } let transmit = self .store .as_context_mut() .0 .concurrent_state_mut() .get_mut(self.id)?; let &ReadState::GuestReady { address, count, options, instance, .. } = &transmit.read else { bail_bug!("expected ReadState::GuestReady") }; let &WriteState::HostReady { guest_offset, .. } = &transmit.write else { bail_bug!("expected WriteState::HostReady") }; let memory = instance .options_memory_mut(self.store.0, options) .get_mut((address + guest_offset)..) .and_then(|b| b.get_mut(..(count - guest_offset))); match memory { Some(memory) => Ok(memory), None => bail_bug!("guest buffer unexpectedly out of bounds"), } } /// Mark the specified number of bytes as written to the writer's buffer. /// /// # Panics /// /// This will panic if the count is larger than the size of the /// buffer returned by `Self::remaining`. pub fn mark_written(&mut self, count: usize) { // Note that this unwrap should only trigger for bugs in Wasmtime, and // this is modeled here to centralize the `.unwrap()` for this method in // one location. self.mark_written_(count).unwrap() } fn mark_written_(&mut self, count: usize) -> Result<()> { if let Some(buffer) = self.host_buffer.as_deref_mut() { buffer.set_position( // Note that these `.unwrap`s are documented panic conditions of // `mark_written`. buffer .position() .checked_add(u64::try_from(count).unwrap()) .unwrap(), ); } else { let transmit = self .store .as_context_mut() .0 .concurrent_state_mut() .get_mut(self.id)?; let ReadState::GuestReady { count: read_count, .. } = &transmit.read else { bail_bug!("expected ReadState::GuestReady") }; let WriteState::HostReady { guest_offset, .. } = &mut transmit.write else { bail_bug!("expected WriteState::HostReady"); }; if *guest_offset + count > *read_count { // Note that this `panic` is a documented panic condition of // `mark_written`. panic!( "write count ({count}) must be less than or equal to read count ({read_count})" ) } else { *guest_offset += count; } } Ok(()) } } /// Represents the state of a `Stream{Producer,Consumer}`. #[derive(Copy, Clone, Debug)] pub enum StreamResult { /// The operation completed normally, and the producer or consumer may be /// able to produce or consume more items, respectively. Completed, /// The operation was interrupted (i.e. it wrapped up early after receiving /// a `finish` parameter value of true in a call to `poll_produce` or /// `poll_consume`), and the producer or consumer may be able to produce or /// consume more items, respectively. Cancelled, /// The operation completed normally, but the producer or consumer will /// _not_ able to produce or consume more items, respectively. Dropped, } /// Represents the host-owned write end of a stream. pub trait StreamProducer: Send + 'static { /// The payload type of this stream. type Item; /// The `WriteBuffer` type to use when delivering items. type Buffer: WriteBuffer + Default; /// Handle a host- or guest-initiated read by delivering zero or more items /// to the specified destination. /// /// This will be called whenever the reader starts a read. /// /// # Arguments /// /// * `self` - a `Pin`'d version of self to perform Rust-level /// future-related operations on. /// * `cx` - a Rust-related [`Context`] which is passed to other /// future-related operations or used to acquire a waker. /// * `store` - the Wasmtime store that this operation is happening within. /// Used, for example, to consult the state `D` associated with the store. /// * `destination` - the location that items are to be written to. /// * `finish` - a flag indicating whether the host should strive to /// immediately complete/cancel any pending operation. See below for more /// details. /// /// # Behavior /// /// If the implementation is able to produce one or more items immediately, /// it should write them to `destination` and return either /// `Poll::Ready(Ok(StreamResult::Completed))` if it expects to produce more /// items, or `Poll::Ready(Ok(StreamResult::Dropped))` if it cannot produce /// any more items. /// /// If the implementation is unable to produce any items immediately, but /// expects to do so later, and `finish` is _false_, it should store the /// waker from `cx` for later and return `Poll::Pending` without writing /// anything to `destination`. Later, it should alert the waker when either /// the items arrive, the stream has ended, or an error occurs. /// /// If more items are written to `destination` than the reader has immediate /// capacity to accept, they will be retained in memory by the caller and /// used to satisfy future reads, in which case `poll_produce` will only be /// called again once all those items have been delivered. /// /// # Zero-length reads /// /// This function may be called with a zero-length capacity buffer /// (i.e. `Destination::remaining` returns `Some(0)`). This indicates that /// the guest wants to wait to see if an item is ready without actually /// reading the item. For example think of a UNIX `poll` function run on a /// TCP stream, seeing if it's readable without actually reading it. /// /// In this situation the host is allowed to either return immediately or /// wait for readiness. Note that waiting for readiness is not always /// possible. For example it's impossible to test if a Rust-native `Future` /// is ready without actually reading the item. Stream-specific /// optimizations, such as testing if a TCP stream is readable, may be /// possible however. /// /// For a zero-length read, the host is allowed to: /// /// - Return `Poll::Ready(Ok(StreamResult::Completed))` without writing /// anything if it expects to be able to produce items immediately (i.e. /// without first returning `Poll::Pending`) the next time `poll_produce` /// is called with non-zero capacity. This is the best-case scenario of /// fulfilling the guest's desire -- items aren't read/buffered but the /// host is saying it's ready when the guest is. /// /// - Return `Poll::Ready(Ok(StreamResult::Completed))` without actually /// testing for readiness. The guest doesn't know this yet, but the guest /// will realize that zero-length reads won't work on this stream when a /// subsequent nonzero read attempt is made which returns `Poll::Pending` /// here. /// /// - Return `Poll::Pending` if the host has performed necessary async work /// to wait for this stream to be readable without actually reading /// anything. This is also a best-case scenario where the host is letting /// the guest know that nothing is ready yet. Later the zero-length read /// will complete and then the guest will attempt a nonzero-length read to /// actually read some bytes. /// /// - Return `Poll::Ready(Ok(StreamResult::Completed))` after calling /// `Destination::set_buffer` with one more more items. Note, however, /// that this creates the hazard that the items will never be received by /// the guest if it decides not to do another non-zero-length read before /// closing the stream. Moreover, if `Self::Item` is e.g. a /// `Resource<_>`, they may end up leaking in that scenario. It is not /// recommended to do this and it's better to return /// `StreamResult::Completed` without buffering anything instead. /// /// For more discussion on zero-length reads see the [documentation in the /// component-model repo itself][docs]. /// /// [docs]: https://github.com/WebAssembly/component-model/blob/main/design/mvp/Concurrency.md#stream-readiness /// /// # Return /// /// This function can return a number of possible cases from this function: /// /// * `Poll::Pending` - this operation cannot complete at this time. The /// Rust-level `Future::poll` contract applies here where a waker should /// be stored from the `cx` argument and be arranged to receive a /// notification when this implementation can make progress. For example /// if you call `Future::poll` on a sub-future, that's enough. If items /// were written to `destination` then a trap in the guest will be raised. /// /// Note that implementations should strive to avoid this return value /// when `finish` is `true`. In such a situation the guest is attempting /// to, for example, cancel a previous operation. By returning /// `Poll::Pending` the guest will be blocked during the cancellation /// request. If `finish` is `true` then `StreamResult::Cancelled` is /// favored to indicate that no items were read. If a short read happened, /// however, it's ok to return `StreamResult::Completed` indicating some /// items were read. /// /// * `Poll::Ok(StreamResult::Completed)` - items, if applicable, were /// written to the `destination`. /// /// * `Poll::Ok(StreamResult::Cancelled)` - used when `finish` is `true` and /// the implementation was able to successfully cancel any async work that /// a previous read kicked off, if any. The host should not buffer values /// received after returning `Cancelled` because the guest will not be /// aware of these values and the guest could close the stream after /// cancelling a read. Hosts should only return `Cancelled` when there are /// no more async operations in flight for a previous read. /// /// If items were written to `destination` then a trap in the guest will /// be raised. If `finish` is `false` then this return value will raise a /// trap in the guest. /// /// * `Poll::Ok(StreamResult::Dropped)` - end-of-stream marker, indicating /// that this producer should not be polled again. Note that items may /// still be written to `destination`. /// /// # Errors /// /// The implementation may alternatively choose to return `Err(_)` to /// indicate an unrecoverable error. This will cause the guest (if any) to /// trap and render the component instance (if any) unusable. The /// implementation should report errors that _are_ recoverable by other /// means (e.g. by writing to a `future`) and return /// `Poll::Ready(Ok(StreamResult::Dropped))`. fn poll_produce<'a>( self: Pin<&mut Self>, cx: &mut Context<'_>, store: StoreContextMut<'a, D>, destination: Destination<'a, Self::Item, Self::Buffer>, finish: bool, ) -> Poll>; /// Attempt to convert the specified object into a `Box` which may /// be downcast to the specified type. /// /// The implementation must ensure that, if it returns `Ok(_)`, a downcast /// to the specified type is guaranteed to succeed. fn try_into(me: Pin>, _ty: TypeId) -> Result, Pin>> { Err(me) } } impl StreamProducer for iter::Empty where T: Send + Sync + 'static, { type Item = T; type Buffer = Option; fn poll_produce<'a>( self: Pin<&mut Self>, _: &mut Context<'_>, _: StoreContextMut<'a, D>, _: Destination<'a, Self::Item, Self::Buffer>, _: bool, ) -> Poll> { Poll::Ready(Ok(StreamResult::Dropped)) } } impl StreamProducer for stream::Empty where T: Send + Sync + 'static, { type Item = T; type Buffer = Option; fn poll_produce<'a>( self: Pin<&mut Self>, _: &mut Context<'_>, _: StoreContextMut<'a, D>, _: Destination<'a, Self::Item, Self::Buffer>, _: bool, ) -> Poll> { Poll::Ready(Ok(StreamResult::Dropped)) } } impl StreamProducer for Vec where T: Unpin + Send + Sync + 'static, { type Item = T; type Buffer = VecBuffer; fn poll_produce<'a>( self: Pin<&mut Self>, _: &mut Context<'_>, _: StoreContextMut<'a, D>, mut dst: Destination<'a, Self::Item, Self::Buffer>, _: bool, ) -> Poll> { dst.set_buffer(mem::take(self.get_mut()).into()); Poll::Ready(Ok(StreamResult::Dropped)) } } impl StreamProducer for Box<[T]> where T: Unpin + Send + Sync + 'static, { type Item = T; type Buffer = VecBuffer; fn poll_produce<'a>( self: Pin<&mut Self>, _: &mut Context<'_>, _: StoreContextMut<'a, D>, mut dst: Destination<'a, Self::Item, Self::Buffer>, _: bool, ) -> Poll> { dst.set_buffer(mem::take(self.get_mut()).into_vec().into()); Poll::Ready(Ok(StreamResult::Dropped)) } } #[cfg(feature = "component-model-async-bytes")] impl StreamProducer for bytes::Bytes { type Item = u8; type Buffer = Cursor; fn poll_produce<'a>( self: Pin<&mut Self>, _: &mut Context<'_>, _store: StoreContextMut<'a, D>, mut dst: Destination<'a, Self::Item, Self::Buffer>, _: bool, ) -> Poll> { dst.set_buffer(Cursor::new(mem::take(self.get_mut()))); Poll::Ready(Ok(StreamResult::Dropped)) } } #[cfg(feature = "component-model-async-bytes")] impl StreamProducer for bytes::BytesMut { type Item = u8; type Buffer = Cursor; fn poll_produce<'a>( self: Pin<&mut Self>, _: &mut Context<'_>, _store: StoreContextMut<'a, D>, mut dst: Destination<'a, Self::Item, Self::Buffer>, _: bool, ) -> Poll> { dst.set_buffer(Cursor::new(mem::take(self.get_mut()))); Poll::Ready(Ok(StreamResult::Dropped)) } } /// Represents the buffer for a host- or guest-initiated stream write. pub struct Source<'a, T> { id: TableId, host_buffer: Option<&'a mut dyn WriteBuffer>, } impl<'a, T> Source<'a, T> { /// Reborrow `self` so it can be used again later. pub fn reborrow(&mut self) -> Source<'_, T> { Source { id: self.id, host_buffer: self.host_buffer.as_deref_mut(), } } /// Accept zero or more items from the writer. pub fn read(&mut self, mut store: S, buffer: &mut B) -> Result<()> where T: func::Lift + 'static, B: ReadBuffer, { if let Some(input) = &mut self.host_buffer { let count = input.remaining().len().min(buffer.remaining_capacity()); buffer.move_from(*input, count); } else { let store = store.as_context_mut(); let transmit = store.0.concurrent_state_mut().get_mut(self.id)?; let &ReadState::HostReady { guest_offset, .. } = &transmit.read else { bail_bug!("expected ReadState::HostReady"); }; let &WriteState::GuestReady { ty, address, count, options, instance, .. } = &transmit.write else { bail_bug!("expected WriteState::GuestReady"); }; let cx = &mut LiftContext::new(store.0.store_opaque_mut(), options, instance); let ty = ty.payload(cx.types); let old_remaining = buffer.remaining_capacity(); lift::( cx, ty.copied(), buffer, address + (T::SIZE32 * guest_offset), count - guest_offset, )?; let transmit = store.0.concurrent_state_mut().get_mut(self.id)?; let ReadState::HostReady { guest_offset, .. } = &mut transmit.read else { bail_bug!("expected ReadState::HostReady"); }; *guest_offset += old_remaining - buffer.remaining_capacity(); } Ok(()) } /// Return the number of items remaining to be read from the current write /// operation. pub fn remaining(&self, mut store: impl AsContextMut) -> usize where T: 'static, { // Note that this unwrap should only trigger for bugs in Wasmtime, and // this is modeled here to centralize the `.unwrap()` for this method in // one location. self.remaining_(store.as_context_mut().0).unwrap() } fn remaining_(&self, store: &mut StoreOpaque) -> Result where T: 'static, { let transmit = store.concurrent_state_mut().get_mut(self.id)?; if let &WriteState::GuestReady { count, .. } = &transmit.write { let &ReadState::HostReady { guest_offset, .. } = &transmit.read else { bail_bug!("expected ReadState::HostReady") }; Ok(count - guest_offset) } else if let Some(host_buffer) = &self.host_buffer { Ok(host_buffer.remaining().len()) } else { bail_bug!("expected either WriteState::GuestReady or host buffer") } } } impl<'a> Source<'a, u8> { /// Return a `DirectSource` view of `self`. pub fn as_direct(self, store: StoreContextMut<'a, D>) -> DirectSource<'a, D> { DirectSource { id: self.id, host_buffer: self.host_buffer, store, } } } /// Represents a write to a `stream`, providing direct access to the /// writer's buffer. pub struct DirectSource<'a, D: 'static> { id: TableId, host_buffer: Option<&'a mut dyn WriteBuffer>, store: StoreContextMut<'a, D>, } impl std::io::Read for DirectSource<'_, D> { fn read(&mut self, buf: &mut [u8]) -> std::io::Result { let rem = self.remaining(); let n = rem.len().min(buf.len()); buf[..n].copy_from_slice(&rem[..n]); self.mark_read(n); Ok(n) } } impl DirectSource<'_, D> { /// Provide direct access to the writer's buffer. pub fn remaining(&mut self) -> &[u8] { // Note that this unwrap should only trigger for bugs in Wasmtime, and // this is modeled here to centralize the `.unwrap()` for this method in // one location. self.remaining_().unwrap() } fn remaining_(&mut self) -> Result<&[u8]> { if let Some(buffer) = self.host_buffer.as_deref_mut() { return Ok(buffer.remaining()); } let transmit = self .store .as_context_mut() .0 .concurrent_state_mut() .get_mut(self.id)?; let &WriteState::GuestReady { address, count, options, instance, .. } = &transmit.write else { bail_bug!("expected WriteState::GuestReady") }; let &ReadState::HostReady { guest_offset, .. } = &transmit.read else { bail_bug!("expected ReadState::HostReady") }; let memory = instance .options_memory(self.store.0, options) .get((address + guest_offset)..) .and_then(|b| b.get(..(count - guest_offset))); match memory { Some(memory) => Ok(memory), None => bail_bug!("guest buffer unexpectedly out of bounds"), } } /// Mark the specified number of bytes as read from the writer's buffer. /// /// # Panics /// /// This will panic if the count is larger than the size of the buffer /// returned by `Self::remaining`. pub fn mark_read(&mut self, count: usize) { // Note that this unwrap should only trigger for bugs in Wasmtime, and // this is modeled here to centralize the `.unwrap()` for this method in // one location. self.mark_read_(count).unwrap() } fn mark_read_(&mut self, count: usize) -> Result<()> { if let Some(buffer) = self.host_buffer.as_deref_mut() { buffer.skip(count); return Ok(()); } let transmit = self .store .as_context_mut() .0 .concurrent_state_mut() .get_mut(self.id)?; let WriteState::GuestReady { count: write_count, .. } = &transmit.write else { bail_bug!("expected WriteState::GuestReady"); }; let ReadState::HostReady { guest_offset, .. } = &mut transmit.read else { bail_bug!("expected ReadState::HostReady"); }; if *guest_offset + count > *write_count { // Note that this is a documented panic condition of `mark_read`. panic!("read count ({count}) must be less than or equal to write count ({write_count})") } else { *guest_offset += count; } Ok(()) } } /// Represents the host-owned read end of a stream. pub trait StreamConsumer: Send + 'static { /// The payload type of this stream. type Item; /// Handle a host- or guest-initiated write by accepting zero or more items /// from the specified source. /// /// This will be called whenever the writer starts a write. /// /// If the implementation is able to consume one or more items immediately, /// it should take them from `source` and return either /// `Poll::Ready(Ok(StreamResult::Completed))` if it expects to be able to consume /// more items, or `Poll::Ready(Ok(StreamResult::Dropped))` if it cannot /// accept any more items. Alternatively, it may return `Poll::Pending` to /// indicate that the caller should delay sending a `COMPLETED` event to the /// writer until a later call to this function returns `Poll::Ready(_)`. /// For more about that, see the `Backpressure` section below. /// /// If the implementation cannot consume any items immediately and `finish` /// is _false_, it should store the waker from `cx` for later and return /// `Poll::Pending` without writing anything to `destination`. Later, it /// should alert the waker when either (1) the items arrive, (2) the stream /// has ended, or (3) an error occurs. /// /// If the implementation cannot consume any items immediately and `finish` /// is _true_, it should, if possible, return /// `Poll::Ready(Ok(StreamResult::Cancelled))` immediately without taking /// anything from `source`. However, that might not be possible if an /// earlier call to `poll_consume` kicked off an asynchronous operation /// which needs to be completed (and possibly interrupted) gracefully, in /// which case the implementation may return `Poll::Pending` and later alert /// the waker as described above. In other words, when `finish` is true, /// the implementation should prioritize returning a result to the reader /// (even if no items can be consumed) rather than wait indefinitely for at /// capacity to free up. /// /// In all of the above cases, the implementation may alternatively choose /// to return `Err(_)` to indicate an unrecoverable error. This will cause /// the guest (if any) to trap and render the component instance (if any) /// unusable. The implementation should report errors that _are_ /// recoverable by other means (e.g. by writing to a `future`) and return /// `Poll::Ready(Ok(StreamResult::Dropped))`. /// /// Note that the implementation should only return /// `Poll::Ready(Ok(StreamResult::Cancelled))` without having taken any /// items from `source` if called with `finish` set to true. If it does so /// when `finish` is false, the caller will trap. Additionally, it should /// only return `Poll::Ready(Ok(StreamResult::Completed))` after taking at /// least one item from `source` if there is an item available; otherwise, /// the caller will trap. If `poll_consume` is called with no items in /// `source`, it should only return `Poll::Ready(_)` once it is able to /// accept at least one item during the next call to `poll_consume`. /// /// Note that any items which the implementation of this trait takes from /// `source` become the responsibility of that implementation. For that /// reason, an implementation which forwards items to an upstream sink /// should reserve capacity in that sink before taking items out of /// `source`, if possible. Alternatively, it might buffer items which can't /// be forwarded immediately and send them once capacity is freed up. /// /// ## Backpressure /// /// As mentioned above, an implementation might choose to return /// `Poll::Pending` after taking items from `source`, which tells the caller /// to delay sending a `COMPLETED` event to the writer. This can be used as /// a form of backpressure when the items are forwarded to an upstream sink /// asynchronously. Note, however, that it's not possible to "put back" /// items into `source` once they've been taken out, so if the upstream sink /// is unable to accept all the items, that cannot be communicated to the /// writer at this level of abstraction. Just as with application-specific, /// recoverable errors, information about which items could be forwarded and /// which could not must be communicated out-of-band, e.g. by writing to an /// application-specific `future`. /// /// Similarly, if the writer cancels the write after items have been taken /// from `source` but before the items have all been forwarded to an /// upstream sink, `poll_consume` will be called with `finish` set to true, /// and the implementation may either: /// /// - Interrupt the forwarding process gracefully. This may be preferable /// if there is an out-of-band channel for communicating to the writer how /// many items were forwarded before being interrupted. /// /// - Allow the forwarding to complete without interrupting it. This is /// usually preferable if there's no out-of-band channel for reporting back /// to the writer how many items were forwarded. fn poll_consume( self: Pin<&mut Self>, cx: &mut Context<'_>, store: StoreContextMut, source: Source<'_, Self::Item>, finish: bool, ) -> Poll>; } /// Represents a host-owned write end of a future. pub trait FutureProducer: Send + 'static { /// The payload type of this future. type Item; /// Handle a host- or guest-initiated read by producing a value. /// /// This is equivalent to `StreamProducer::poll_produce`, but with a /// simplified interface for futures. /// /// If `finish` is true, the implementation may return /// `Poll::Ready(Ok(None))` to indicate the operation was canceled before it /// could produce a value. Otherwise, it must either return /// `Poll::Ready(Ok(Some(_)))`, `Poll::Ready(Err(_))`, or `Poll::Pending`. fn poll_produce( self: Pin<&mut Self>, cx: &mut Context<'_>, store: StoreContextMut, finish: bool, ) -> Poll>>; } impl FutureProducer for Fut where E: Into, Fut: Future> + ?Sized + Send + 'static, { type Item = T; fn poll_produce<'a>( self: Pin<&mut Self>, cx: &mut Context<'_>, _: StoreContextMut<'a, D>, finish: bool, ) -> Poll>> { match self.poll(cx) { Poll::Ready(Ok(v)) => Poll::Ready(Ok(Some(v))), Poll::Ready(Err(err)) => Poll::Ready(Err(err.into())), Poll::Pending if finish => Poll::Ready(Ok(None)), Poll::Pending => Poll::Pending, } } } /// Represents a host-owned read end of a future. pub trait FutureConsumer: Send + 'static { /// The payload type of this future. type Item; /// Handle a host- or guest-initiated write by consuming a value. /// /// This is equivalent to `StreamProducer::poll_produce`, but with a /// simplified interface for futures. /// /// If `finish` is true, the implementation may return `Poll::Ready(Ok(()))` /// without taking the item from `source`, which indicates the operation was /// canceled before it could consume the value. Otherwise, it must either /// take the item from `source` and return `Poll::Ready(Ok(()))`, or else /// return `Poll::Ready(Err(_))` or `Poll::Pending` (with or without taking /// the item). fn poll_consume( self: Pin<&mut Self>, cx: &mut Context<'_>, store: StoreContextMut, source: Source<'_, Self::Item>, finish: bool, ) -> Poll>; } /// Represents the readable end of a Component Model `future`. /// /// Note that `FutureReader` instances must be disposed of using either `pipe` /// or `close`; otherwise the in-store representation will leak and the writer /// end will hang indefinitely. Consider using [`GuardedFutureReader`] to /// ensure that disposal happens automatically. pub struct FutureReader { id: TableId, _phantom: PhantomData, } impl FutureReader { /// Create a new future with the specified producer. /// /// # Errors /// /// Returns an error if the resource table for this store is full or if /// [`Config::concurrency_support`] is not enabled. /// /// [`Config::concurrency_support`]: crate::Config::concurrency_support pub fn new( mut store: S, producer: impl FutureProducer, ) -> Result where T: func::Lower + func::Lift + Send + Sync + 'static, { ensure!( store.as_context().0.concurrency_support(), "concurrency support is not enabled" ); struct Producer

(P); impl> StreamProducer for Producer

{ type Item = P::Item; type Buffer = Option; fn poll_produce<'a>( self: Pin<&mut Self>, cx: &mut Context<'_>, store: StoreContextMut, mut destination: Destination<'a, Self::Item, Self::Buffer>, finish: bool, ) -> Poll> { // SAFETY: This is a standard pin-projection, and we never move // out of `self`. let producer = unsafe { self.map_unchecked_mut(|v| &mut v.0) }; Poll::Ready(Ok( if let Some(value) = ready!(producer.poll_produce(cx, store, finish))? { destination.set_buffer(Some(value)); // Here we return `StreamResult::Completed` even though // we've produced the last item we'll ever produce. // That's because the ABI expects // `ReturnCode::Completed(1)` rather than // `ReturnCode::Dropped(1)`. In any case, we won't be // called again since the future will have resolved. StreamResult::Completed } else { StreamResult::Cancelled }, )) } } Ok(Self::new_( store .as_context_mut() .new_transmit(TransmitKind::Future, Producer(producer))?, )) } pub(super) fn new_(id: TableId) -> Self { Self { id, _phantom: PhantomData, } } pub(super) fn id(&self) -> TableId { self.id } /// Set the consumer that accepts the result of this future. /// /// # Errors /// /// Returns an error if this future has already been closed. /// /// # Panics /// /// Panics if this future does not belong to `store`. pub fn pipe( self, mut store: S, consumer: impl FutureConsumer + Unpin, ) -> Result<()> where T: func::Lift + 'static, { struct Consumer(C); impl> StreamConsumer for Consumer { type Item = T; fn poll_consume( self: Pin<&mut Self>, cx: &mut Context<'_>, mut store: StoreContextMut, mut source: Source, finish: bool, ) -> Poll> { // SAFETY: This is a standard pin-projection, and we never move // out of `self`. let consumer = unsafe { self.map_unchecked_mut(|v| &mut v.0) }; ready!(consumer.poll_consume( cx, store.as_context_mut(), source.reborrow(), finish ))?; Poll::Ready(Ok(if source.remaining(store) == 0 { // Here we return `StreamResult::Completed` even though // we've consumed the last item we'll ever consume. That's // because the ABI expects `ReturnCode::Completed(1)` rather // than `ReturnCode::Dropped(1)`. In any case, we won't be // called again since the future will have resolved. StreamResult::Completed } else { StreamResult::Cancelled })) } } store .as_context_mut() .set_consumer(self.id, TransmitKind::Future, Consumer(consumer)) } /// Transfer ownership of the read end of a future from a guest to the host. fn lift_from_index(cx: &mut LiftContext<'_>, ty: InterfaceType, index: u32) -> Result { let id = lift_index_to_future(cx, ty, index)?; Ok(Self::new_(id)) } /// Close this `FutureReader`. /// /// This will close this half of the future which will signal to a pending /// write, if any, that the reader side is dropped. If the writer half has /// not yet written a value then when it attempts to write a value it will /// see that this end is closed. /// /// # Errors /// /// Returns an error if this future has already been closed. /// /// # Panics /// /// Panics if the store that the [`Accessor`] is derived from does not own /// this future. /// /// [`Accessor`]: crate::component::Accessor pub fn close(&mut self, mut store: impl AsContextMut) -> Result<()> { future_close(store.as_context_mut().0, &mut self.id) } /// Convenience method around [`Self::close`]. pub fn close_with(&mut self, accessor: impl AsAccessor) -> Result<()> { accessor.as_accessor().with(|access| self.close(access)) } /// Returns a [`GuardedFutureReader`] which will auto-close this future on /// drop and clean it up from the store. /// /// Note that the `accessor` provided must own this future and is /// additionally transferred to the `GuardedFutureReader` return value. pub fn guard(self, accessor: A) -> GuardedFutureReader where A: AsAccessor, { GuardedFutureReader::new(accessor, self) } /// Attempts to convert this [`FutureReader`] to a [`FutureAny`]. /// /// # Errors /// /// This function will return an error if `self` does not belong to /// `store`. pub fn try_into_future_any(self, store: impl AsContextMut) -> Result where T: ComponentType + 'static, { FutureAny::try_from_future_reader(store, self) } /// Attempts to convert a [`FutureAny`] into a [`FutureReader`]. /// /// # Errors /// /// This function will fail if `T` doesn't match the type of the value that /// `future` is sending. pub fn try_from_future_any(future: FutureAny) -> Result where T: ComponentType + 'static, { future.try_into_future_reader() } } impl fmt::Debug for FutureReader { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("FutureReader") .field("id", &self.id) .finish() } } pub(super) fn future_close( store: &mut StoreOpaque, id: &mut TableId, ) -> Result<()> { let id = mem::replace(id, TableId::new(u32::MAX)); store.host_drop_reader(id, TransmitKind::Future) } /// Transfer ownership of the read end of a future from the host to a guest. pub(super) fn lift_index_to_future( cx: &mut LiftContext<'_>, ty: InterfaceType, index: u32, ) -> Result> { match ty { InterfaceType::Future(src) => { let handle_table = cx .instance_mut() .table_for_transmit(TransmitIndex::Future(src)); let (rep, is_done) = handle_table.future_remove_readable(src, index)?; if is_done { bail!("cannot lift future after being notified that the writable end dropped"); } let id = TableId::::new(rep); let concurrent_state = cx.concurrent_state_mut(); let future = concurrent_state.get_mut(id)?; future.common.handle = None; let state = future.state; if concurrent_state.get_mut(state)?.done { bail!("cannot lift future after previous read succeeded"); } Ok(id) } _ => func::bad_type_info(), } } /// Transfer ownership of the read end of a future from the host to a guest. pub(super) fn lower_future_to_index( id: TableId, cx: &mut LowerContext<'_, U>, ty: InterfaceType, ) -> Result { match ty { InterfaceType::Future(dst) => { let concurrent_state = cx.store.0.concurrent_state_mut(); let state = concurrent_state.get_mut(id)?.state; let rep = concurrent_state.get_mut(state)?.read_handle.rep(); let handle = cx .instance_mut() .table_for_transmit(TransmitIndex::Future(dst)) .future_insert_read(dst, rep)?; cx.store.0.concurrent_state_mut().get_mut(id)?.common.handle = Some(handle); Ok(handle) } _ => func::bad_type_info(), } } // SAFETY: This relies on the `ComponentType` implementation for `u32` being // safe and correct since we lift and lower future handles as `u32`s. unsafe impl ComponentType for FutureReader { const ABI: CanonicalAbiInfo = CanonicalAbiInfo::SCALAR4; type Lower = ::Lower; fn typecheck(ty: &InterfaceType, types: &InstanceType<'_>) -> Result<()> { match ty { InterfaceType::Future(ty) => { let ty = types.types[*ty].ty; types::typecheck_payload::(types.types[ty].payload.as_ref(), types) } other => bail!("expected `future`, found `{}`", func::desc(other)), } } } // SAFETY: See the comment on the `ComponentType` `impl` for this type. unsafe impl func::Lower for FutureReader { fn linear_lower_to_flat( &self, cx: &mut LowerContext<'_, U>, ty: InterfaceType, dst: &mut MaybeUninit, ) -> Result<()> { lower_future_to_index(self.id, cx, ty)?.linear_lower_to_flat(cx, InterfaceType::U32, dst) } fn linear_lower_to_memory( &self, cx: &mut LowerContext<'_, U>, ty: InterfaceType, offset: usize, ) -> Result<()> { lower_future_to_index(self.id, cx, ty)?.linear_lower_to_memory( cx, InterfaceType::U32, offset, ) } } // SAFETY: See the comment on the `ComponentType` `impl` for this type. unsafe impl func::Lift for FutureReader { fn linear_lift_from_flat( cx: &mut LiftContext<'_>, ty: InterfaceType, src: &Self::Lower, ) -> Result { let index = u32::linear_lift_from_flat(cx, InterfaceType::U32, src)?; Self::lift_from_index(cx, ty, index) } fn linear_lift_from_memory( cx: &mut LiftContext<'_>, ty: InterfaceType, bytes: &[u8], ) -> Result { let index = u32::linear_lift_from_memory(cx, InterfaceType::U32, bytes)?; Self::lift_from_index(cx, ty, index) } } /// A [`FutureReader`] paired with an [`Accessor`]. /// /// This is an RAII wrapper around [`FutureReader`] that ensures it is closed /// when dropped. This can be created through [`GuardedFutureReader::new`] or /// [`FutureReader::guard`]. /// /// [`Accessor`]: crate::component::Accessor pub struct GuardedFutureReader where A: AsAccessor, { // This field is `None` to implement the conversion from this guard back to // `FutureReader`. When `None` is seen in the destructor it will cause the // destructor to do nothing. reader: Option>, accessor: A, } impl GuardedFutureReader where A: AsAccessor, { /// Create a new `GuardedFutureReader` with the specified `accessor` and `reader`. /// /// # Panics /// /// Panics if [`Config::concurrency_support`] is not enabled. /// /// [`Config::concurrency_support`]: crate::Config::concurrency_support pub fn new(accessor: A, reader: FutureReader) -> Self { assert!( accessor .as_accessor() .with(|a| a.as_context().0.concurrency_support()) ); Self { reader: Some(reader), accessor, } } /// Extracts the underlying [`FutureReader`] from this guard, returning it /// back. pub fn into_future(self) -> FutureReader { self.into() } } impl From> for FutureReader where A: AsAccessor, { fn from(mut guard: GuardedFutureReader) -> Self { guard.reader.take().unwrap() } } impl Drop for GuardedFutureReader where A: AsAccessor, { fn drop(&mut self) { if let Some(reader) = &mut self.reader { // Currently this can only fail if the future is closed twice, which // this guard prevents, so this error shouldn't happen. let result = reader.close_with(&self.accessor); debug_assert!(result.is_ok()); } } } /// Represents the readable end of a Component Model `stream`. /// /// Note that `StreamReader` instances must be disposed of using `close`; /// otherwise the in-store representation will leak and the writer end will hang /// indefinitely. Consider using [`GuardedStreamReader`] to ensure that /// disposal happens automatically. pub struct StreamReader { id: TableId, _phantom: PhantomData, } impl StreamReader { /// Create a new stream with the specified producer. /// /// # Errors /// /// Returns an error if the resource table for this store is full or if /// [`Config::concurrency_support`] is not enabled. /// /// [`Config::concurrency_support`]: crate::Config::concurrency_support pub fn new( mut store: S, producer: impl StreamProducer, ) -> Result where T: func::Lower + func::Lift + Send + Sync + 'static, { ensure!( store.as_context().0.concurrency_support(), "concurrency support is not enabled", ); Ok(Self::new_( store .as_context_mut() .new_transmit(TransmitKind::Stream, producer)?, )) } pub(super) fn new_(id: TableId) -> Self { Self { id, _phantom: PhantomData, } } pub(super) fn id(&self) -> TableId { self.id } /// Attempt to consume this object by converting it into the specified type. /// /// This can be useful for "short-circuiting" host-to-host streams, /// bypassing the guest entirely. For example, if a guest task returns a /// host-created stream and then exits, this function may be used to /// retrieve the write end, after which the guest instance and store may be /// disposed of if no longer needed. /// /// This will return `Ok(_)` if and only if the following conditions are /// met: /// /// - The stream was created by the host (i.e. not by the guest). /// /// - The `StreamProducer::try_into` function returns `Ok(_)` when given the /// producer provided to `StreamReader::new` when the stream was created, /// along with `TypeId::of::()`. /// /// # Panics /// /// Panics if this stream has already been closed, or if this stream doesn't /// belong to the specified `store`. pub fn try_into(mut self, mut store: impl AsContextMut) -> Result { let store = store.as_context_mut(); let state = store.0.concurrent_state_mut(); let id = state.get_mut(self.id).unwrap().state; if let WriteState::HostReady { try_into, .. } = &state.get_mut(id).unwrap().write { match try_into(TypeId::of::()) { Some(result) => { self.close(store).unwrap(); Ok(*result.downcast::().unwrap()) } None => Err(self), } } else { Err(self) } } /// Set the consumer that accepts the items delivered to this stream. /// /// # Errors /// /// Returns an error if this stream has already been closed. /// /// # Panics /// /// Panics if this stream does not belong to `store`. pub fn pipe( self, mut store: S, consumer: impl StreamConsumer, ) -> Result<()> where T: 'static, { store .as_context_mut() .set_consumer(self.id, TransmitKind::Stream, consumer) } /// Transfer ownership of the read end of a stream from a guest to the host. fn lift_from_index(cx: &mut LiftContext<'_>, ty: InterfaceType, index: u32) -> Result { let id = lift_index_to_stream(cx, ty, index)?; Ok(Self::new_(id)) } /// Close this `StreamReader`. /// /// This will signal that this portion of the stream is closed causing all /// future writes to return immediately with "DROPPED". /// /// # Errors /// /// Returns an error if this stream has already been closed. /// /// # Panics /// /// Panics if the store that the [`Accessor`] is derived from does not own /// this stream. /// /// [`Accessor`]: crate::component::Accessor pub fn close(&mut self, mut store: impl AsContextMut) -> Result<()> { stream_close(store.as_context_mut().0, &mut self.id) } /// Convenience method around [`Self::close`]. pub fn close_with(&mut self, accessor: impl AsAccessor) -> Result<()> { accessor.as_accessor().with(|access| self.close(access)) } /// Returns a [`GuardedStreamReader`] which will auto-close this stream on /// drop and clean it up from the store. /// /// Note that the `accessor` provided must own this future and is /// additionally transferred to the `GuardedStreamReader` return value. pub fn guard(self, accessor: A) -> GuardedStreamReader where A: AsAccessor, { GuardedStreamReader::new(accessor, self) } /// Attempts to convert this [`StreamReader`] to a [`StreamAny`]. /// /// # Errors /// /// This function will return an error if `self` does not belong to /// `store`. pub fn try_into_stream_any(self, store: impl AsContextMut) -> Result where T: ComponentType + 'static, { StreamAny::try_from_stream_reader(store, self) } /// Attempts to convert a [`StreamAny`] into a [`StreamReader`]. /// /// # Errors /// /// This function will fail if `T` doesn't match the type of the value that /// `stream` is sending. pub fn try_from_stream_any(stream: StreamAny) -> Result where T: ComponentType + 'static, { stream.try_into_stream_reader() } } impl fmt::Debug for StreamReader { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("StreamReader") .field("id", &self.id) .finish() } } pub(super) fn stream_close( store: &mut StoreOpaque, id: &mut TableId, ) -> Result<()> { let id = mem::replace(id, TableId::new(u32::MAX)); store.host_drop_reader(id, TransmitKind::Stream) } /// Transfer ownership of the read end of a stream from a guest to the host. pub(super) fn lift_index_to_stream( cx: &mut LiftContext<'_>, ty: InterfaceType, index: u32, ) -> Result> { match ty { InterfaceType::Stream(src) => { let handle_table = cx .instance_mut() .table_for_transmit(TransmitIndex::Stream(src)); let (rep, is_done) = handle_table.stream_remove_readable(src, index)?; if is_done { bail!("cannot lift stream after being notified that the writable end dropped"); } let id = TableId::::new(rep); cx.concurrent_state_mut().get_mut(id)?.common.handle = None; Ok(id) } _ => func::bad_type_info(), } } /// Transfer ownership of the read end of a stream from the host to a guest. pub(super) fn lower_stream_to_index( id: TableId, cx: &mut LowerContext<'_, U>, ty: InterfaceType, ) -> Result { match ty { InterfaceType::Stream(dst) => { let concurrent_state = cx.store.0.concurrent_state_mut(); let state = concurrent_state.get_mut(id)?.state; let rep = concurrent_state.get_mut(state)?.read_handle.rep(); let handle = cx .instance_mut() .table_for_transmit(TransmitIndex::Stream(dst)) .stream_insert_read(dst, rep)?; cx.store.0.concurrent_state_mut().get_mut(id)?.common.handle = Some(handle); Ok(handle) } _ => func::bad_type_info(), } } // SAFETY: This relies on the `ComponentType` implementation for `u32` being // safe and correct since we lift and lower stream handles as `u32`s. unsafe impl ComponentType for StreamReader { const ABI: CanonicalAbiInfo = CanonicalAbiInfo::SCALAR4; type Lower = ::Lower; fn typecheck(ty: &InterfaceType, types: &InstanceType<'_>) -> Result<()> { match ty { InterfaceType::Stream(ty) => { let ty = types.types[*ty].ty; types::typecheck_payload::(types.types[ty].payload.as_ref(), types) } other => bail!("expected `stream`, found `{}`", func::desc(other)), } } } // SAFETY: See the comment on the `ComponentType` `impl` for this type. unsafe impl func::Lower for StreamReader { fn linear_lower_to_flat( &self, cx: &mut LowerContext<'_, U>, ty: InterfaceType, dst: &mut MaybeUninit, ) -> Result<()> { lower_stream_to_index(self.id, cx, ty)?.linear_lower_to_flat(cx, InterfaceType::U32, dst) } fn linear_lower_to_memory( &self, cx: &mut LowerContext<'_, U>, ty: InterfaceType, offset: usize, ) -> Result<()> { lower_stream_to_index(self.id, cx, ty)?.linear_lower_to_memory( cx, InterfaceType::U32, offset, ) } } // SAFETY: See the comment on the `ComponentType` `impl` for this type. unsafe impl func::Lift for StreamReader { fn linear_lift_from_flat( cx: &mut LiftContext<'_>, ty: InterfaceType, src: &Self::Lower, ) -> Result { let index = u32::linear_lift_from_flat(cx, InterfaceType::U32, src)?; Self::lift_from_index(cx, ty, index) } fn linear_lift_from_memory( cx: &mut LiftContext<'_>, ty: InterfaceType, bytes: &[u8], ) -> Result { let index = u32::linear_lift_from_memory(cx, InterfaceType::U32, bytes)?; Self::lift_from_index(cx, ty, index) } } /// A [`StreamReader`] paired with an [`Accessor`]. /// /// This is an RAII wrapper around [`StreamReader`] that ensures it is closed /// when dropped. This can be created through [`GuardedStreamReader::new`] or /// [`StreamReader::guard`]. /// /// [`Accessor`]: crate::component::Accessor pub struct GuardedStreamReader where A: AsAccessor, { // This field is `None` to implement the conversion from this guard back to // `StreamReader`. When `None` is seen in the destructor it will cause the // destructor to do nothing. reader: Option>, accessor: A, } impl GuardedStreamReader where A: AsAccessor, { /// Create a new `GuardedStreamReader` with the specified `accessor` and /// `reader`. /// /// # Panics /// /// Panics if [`Config::concurrency_support`] is not enabled. /// /// [`Config::concurrency_support`]: crate::Config::concurrency_support pub fn new(accessor: A, reader: StreamReader) -> Self { assert!( accessor .as_accessor() .with(|a| a.as_context().0.concurrency_support()) ); Self { reader: Some(reader), accessor, } } /// Extracts the underlying [`StreamReader`] from this guard, returning it /// back. pub fn into_stream(self) -> StreamReader { self.into() } } impl From> for StreamReader where A: AsAccessor, { fn from(mut guard: GuardedStreamReader) -> Self { guard.reader.take().unwrap() } } impl Drop for GuardedStreamReader where A: AsAccessor, { fn drop(&mut self) { if let Some(reader) = &mut self.reader { // Currently this can only fail if the future is closed twice, which // this guard prevents, so this error shouldn't happen. let result = reader.close_with(&self.accessor); debug_assert!(result.is_ok()); } } } /// Represents a Component Model `error-context`. pub struct ErrorContext { rep: u32, } impl ErrorContext { pub(crate) fn new(rep: u32) -> Self { Self { rep } } /// Convert this `ErrorContext` into a [`Val`]. pub fn into_val(self) -> Val { Val::ErrorContext(ErrorContextAny(self.rep)) } /// Attempt to convert the specified [`Val`] to a `ErrorContext`. pub fn from_val(_: impl AsContextMut, value: &Val) -> Result { let Val::ErrorContext(ErrorContextAny(rep)) = value else { bail!("expected `error-context`; got `{}`", value.desc()); }; Ok(Self::new(*rep)) } fn lift_from_index(cx: &mut LiftContext<'_>, ty: InterfaceType, index: u32) -> Result { match ty { InterfaceType::ErrorContext(src) => { let rep = cx .instance_mut() .table_for_error_context(src) .error_context_rep(index)?; Ok(Self { rep }) } _ => func::bad_type_info(), } } } pub(crate) fn lower_error_context_to_index( rep: u32, cx: &mut LowerContext<'_, U>, ty: InterfaceType, ) -> Result { match ty { InterfaceType::ErrorContext(dst) => { let tbl = cx.instance_mut().table_for_error_context(dst); tbl.error_context_insert(rep) } _ => func::bad_type_info(), } } // SAFETY: This relies on the `ComponentType` implementation for `u32` being // safe and correct since we lift and lower future handles as `u32`s. unsafe impl func::ComponentType for ErrorContext { const ABI: CanonicalAbiInfo = CanonicalAbiInfo::SCALAR4; type Lower = ::Lower; fn typecheck(ty: &InterfaceType, _types: &InstanceType<'_>) -> Result<()> { match ty { InterfaceType::ErrorContext(_) => Ok(()), other => bail!("expected `error`, found `{}`", func::desc(other)), } } } // SAFETY: See the comment on the `ComponentType` `impl` for this type. unsafe impl func::Lower for ErrorContext { fn linear_lower_to_flat( &self, cx: &mut LowerContext<'_, T>, ty: InterfaceType, dst: &mut MaybeUninit, ) -> Result<()> { lower_error_context_to_index(self.rep, cx, ty)?.linear_lower_to_flat( cx, InterfaceType::U32, dst, ) } fn linear_lower_to_memory( &self, cx: &mut LowerContext<'_, T>, ty: InterfaceType, offset: usize, ) -> Result<()> { lower_error_context_to_index(self.rep, cx, ty)?.linear_lower_to_memory( cx, InterfaceType::U32, offset, ) } } // SAFETY: See the comment on the `ComponentType` `impl` for this type. unsafe impl func::Lift for ErrorContext { fn linear_lift_from_flat( cx: &mut LiftContext<'_>, ty: InterfaceType, src: &Self::Lower, ) -> Result { let index = u32::linear_lift_from_flat(cx, InterfaceType::U32, src)?; Self::lift_from_index(cx, ty, index) } fn linear_lift_from_memory( cx: &mut LiftContext<'_>, ty: InterfaceType, bytes: &[u8], ) -> Result { let index = u32::linear_lift_from_memory(cx, InterfaceType::U32, bytes)?; Self::lift_from_index(cx, ty, index) } } /// Represents the read or write end of a stream or future. pub(super) struct TransmitHandle { pub(super) common: WaitableCommon, /// See `TransmitState` state: TableId, } impl TransmitHandle { fn new(state: TableId) -> Self { Self { common: WaitableCommon::default(), state, } } } impl TableDebug for TransmitHandle { fn type_name() -> &'static str { "TransmitHandle" } } /// Represents the state of a stream or future. struct TransmitState { /// The write end of the stream or future. write_handle: TableId, /// The read end of the stream or future. read_handle: TableId, /// See `WriteState` write: WriteState, /// See `ReadState` read: ReadState, /// Whether further values may be transmitted via this stream or future. done: bool, /// The original creator of this stream, used for type-checking with /// `{Future,Stream}Any`. pub(super) origin: TransmitOrigin, } #[derive(Copy, Clone)] pub(super) enum TransmitOrigin { Host, GuestFuture(ComponentInstanceId, TypeFutureTableIndex), GuestStream(ComponentInstanceId, TypeStreamTableIndex), } impl TransmitState { fn new(origin: TransmitOrigin) -> Self { Self { write_handle: TableId::new(u32::MAX), read_handle: TableId::new(u32::MAX), read: ReadState::Open, write: WriteState::Open, done: false, origin, } } } impl TableDebug for TransmitState { fn type_name() -> &'static str { "TransmitState" } } impl TransmitOrigin { fn guest(id: ComponentInstanceId, index: TransmitIndex) -> Self { match index { TransmitIndex::Future(ty) => TransmitOrigin::GuestFuture(id, ty), TransmitIndex::Stream(ty) => TransmitOrigin::GuestStream(id, ty), } } } type PollStream = Box< dyn Fn() -> Pin> + Send + 'static>> + Send + Sync, >; type TryInto = Box Option> + Send + Sync>; /// Represents the state of the write end of a stream or future. enum WriteState { /// The write end is open, but no write is pending. Open, /// The write end is owned by a guest task and a write is pending. GuestReady { instance: Instance, caller: RuntimeComponentInstanceIndex, ty: TransmitIndex, flat_abi: Option, options: OptionsIndex, address: usize, count: usize, handle: u32, }, /// The write end is owned by the host, which is ready to produce items. HostReady { produce: PollStream, try_into: TryInto, guest_offset: usize, cancel: bool, cancel_waker: Option, }, /// The write end has been dropped. Dropped, } impl fmt::Debug for WriteState { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self { Self::Open => f.debug_tuple("Open").finish(), Self::GuestReady { .. } => f.debug_tuple("GuestReady").finish(), Self::HostReady { .. } => f.debug_tuple("HostReady").finish(), Self::Dropped => f.debug_tuple("Dropped").finish(), } } } /// Represents the state of the read end of a stream or future. enum ReadState { /// The read end is open, but no read is pending. Open, /// The read end is owned by a guest task and a read is pending. GuestReady { ty: TransmitIndex, caller_instance: RuntimeComponentInstanceIndex, caller_thread: QualifiedThreadId, flat_abi: Option, instance: Instance, options: OptionsIndex, address: usize, count: usize, handle: u32, }, /// The read end is owned by a host task, and it is ready to consume items. HostReady { consume: PollStream, guest_offset: usize, cancel: bool, cancel_waker: Option, }, /// Both the read and write ends are owned by the host. HostToHost { accept: Box< dyn for<'a> Fn( &'a mut UntypedWriteBuffer<'a>, ) -> Pin> + Send + 'a>> + Send + Sync, >, buffer: Vec, limit: usize, }, /// The read end has been dropped. Dropped, } impl fmt::Debug for ReadState { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self { Self::Open => f.debug_tuple("Open").finish(), Self::GuestReady { .. } => f.debug_tuple("GuestReady").finish(), Self::HostReady { .. } => f.debug_tuple("HostReady").finish(), Self::HostToHost { .. } => f.debug_tuple("HostToHost").finish(), Self::Dropped => f.debug_tuple("Dropped").finish(), } } } fn return_code(kind: TransmitKind, state: StreamResult, guest_offset: usize) -> Result { let count = guest_offset.try_into()?; Ok(match state { StreamResult::Dropped => ReturnCode::Dropped(count), StreamResult::Completed => ReturnCode::completed(kind, count), StreamResult::Cancelled => ReturnCode::Cancelled(count), }) } impl StoreOpaque { fn pipe_from_guest( &mut self, kind: TransmitKind, id: TableId, future: Pin> + Send + 'static>>, ) { let future = async move { let stream_state = future.await?; tls::get(|store| { let state = store.concurrent_state_mut(); let transmit = state.get_mut(id)?; let ReadState::HostReady { consume, guest_offset, .. } = mem::replace(&mut transmit.read, ReadState::Open) else { bail_bug!("expected ReadState::HostReady") }; let code = return_code(kind, stream_state, guest_offset)?; transmit.read = match stream_state { StreamResult::Dropped => ReadState::Dropped, StreamResult::Completed | StreamResult::Cancelled => ReadState::HostReady { consume, guest_offset: 0, cancel: false, cancel_waker: None, }, }; let WriteState::GuestReady { ty, handle, .. } = mem::replace(&mut transmit.write, WriteState::Open) else { bail_bug!("expected WriteState::HostReady") }; state.send_write_result(ty, id, handle, code)?; Ok(()) }) }; self.concurrent_state_mut().push_future(future.boxed()); } fn pipe_to_guest( &mut self, kind: TransmitKind, id: TableId, future: Pin> + Send + 'static>>, ) { let future = async move { let stream_state = future.await?; tls::get(|store| { let state = store.concurrent_state_mut(); let transmit = state.get_mut(id)?; let WriteState::HostReady { produce, try_into, guest_offset, .. } = mem::replace(&mut transmit.write, WriteState::Open) else { bail_bug!("expected WriteState::HostReady") }; let code = return_code(kind, stream_state, guest_offset)?; transmit.write = match stream_state { StreamResult::Dropped => WriteState::Dropped, StreamResult::Completed | StreamResult::Cancelled => WriteState::HostReady { produce, try_into, guest_offset: 0, cancel: false, cancel_waker: None, }, }; let ReadState::GuestReady { ty, handle, .. } = mem::replace(&mut transmit.read, ReadState::Open) else { bail_bug!("expected ReadState::HostReady") }; state.send_read_result(ty, id, handle, code)?; Ok(()) }) }; self.concurrent_state_mut().push_future(future.boxed()); } /// Drop the read end of a stream or future read from the host. fn host_drop_reader(&mut self, id: TableId, kind: TransmitKind) -> Result<()> { let state = self.concurrent_state_mut(); let transmit_id = state.get_mut(id)?.state; let transmit = state .get_mut(transmit_id) .with_context(|| format!("error closing reader {transmit_id:?}"))?; log::trace!( "host_drop_reader state {transmit_id:?}; read state {:?} write state {:?}", transmit.read, transmit.write ); transmit.read = ReadState::Dropped; // If the write end is already dropped, it should stay dropped, // otherwise, it should be opened. let new_state = if let WriteState::Dropped = &transmit.write { WriteState::Dropped } else { WriteState::Open }; let write_handle = transmit.write_handle; match mem::replace(&mut transmit.write, new_state) { // If a guest is waiting to write, notify it that the read end has // been dropped. WriteState::GuestReady { ty, handle, .. } => { state.update_event( write_handle.rep(), match ty { TransmitIndex::Future(ty) => Event::FutureWrite { code: ReturnCode::Dropped(0), pending: Some((ty, handle)), }, TransmitIndex::Stream(ty) => Event::StreamWrite { code: ReturnCode::Dropped(0), pending: Some((ty, handle)), }, }, )?; } WriteState::HostReady { .. } => {} WriteState::Open => { state.update_event( write_handle.rep(), match kind { TransmitKind::Future => Event::FutureWrite { code: ReturnCode::Dropped(0), pending: None, }, TransmitKind::Stream => Event::StreamWrite { code: ReturnCode::Dropped(0), pending: None, }, }, )?; } WriteState::Dropped => { log::trace!("host_drop_reader delete {transmit_id:?}"); state.delete_transmit(transmit_id)?; } } Ok(()) } /// Drop the write end of a stream or future read from the host. fn host_drop_writer( &mut self, id: TableId, on_drop_open: Option Result<()>>, ) -> Result<()> { let state = self.concurrent_state_mut(); let transmit_id = state.get_mut(id)?.state; let transmit = state .get_mut(transmit_id) .with_context(|| format!("error closing writer {transmit_id:?}"))?; log::trace!( "host_drop_writer state {transmit_id:?}; write state {:?} read state {:?}", transmit.read, transmit.write ); // Existing queued transmits must be updated with information for the impending writer closure match &mut transmit.write { WriteState::GuestReady { .. } => { bail_bug!("can't call `host_drop_writer` on a guest-owned writer"); } WriteState::HostReady { .. } => {} v @ WriteState::Open => { if let (Some(on_drop_open), false) = ( on_drop_open, transmit.done || matches!(transmit.read, ReadState::Dropped), ) { on_drop_open()?; } else { *v = WriteState::Dropped; } } WriteState::Dropped => bail_bug!("write state is already dropped"), } let transmit = self.concurrent_state_mut().get_mut(transmit_id)?; // If the existing read state is dropped, then there's nothing to read // and we can keep it that way. // // If the read state was any other state, then we must set the new state to open // to indicate that there *is* data to be read let new_state = if let ReadState::Dropped = &transmit.read { ReadState::Dropped } else { ReadState::Open }; let read_handle = transmit.read_handle; // Swap in the new read state match mem::replace(&mut transmit.read, new_state) { // If the guest was ready to read, then we cannot drop the reader (or writer); // we must deliver the event, and update the state associated with the handle to // represent that a read must be performed ReadState::GuestReady { ty, handle, .. } => { // Ensure the final read of the guest is queued, with appropriate closure indicator self.concurrent_state_mut().update_event( read_handle.rep(), match ty { TransmitIndex::Future(ty) => Event::FutureRead { code: ReturnCode::Dropped(0), pending: Some((ty, handle)), }, TransmitIndex::Stream(ty) => Event::StreamRead { code: ReturnCode::Dropped(0), pending: Some((ty, handle)), }, }, )?; } ReadState::HostReady { .. } | ReadState::HostToHost { .. } => {} // If the read state is open, then there are no registered readers of the stream/future ReadState::Open => { self.concurrent_state_mut().update_event( read_handle.rep(), match on_drop_open { Some(_) => Event::FutureRead { code: ReturnCode::Dropped(0), pending: None, }, None => Event::StreamRead { code: ReturnCode::Dropped(0), pending: None, }, }, )?; } // If the read state was already dropped, then we can remove the transmit state completely // (both writer and reader have been dropped) ReadState::Dropped => { log::trace!("host_drop_writer delete {transmit_id:?}"); self.concurrent_state_mut().delete_transmit(transmit_id)?; } } Ok(()) } pub(super) fn transmit_origin( &mut self, id: TableId, ) -> Result { let state = self.concurrent_state_mut(); let state_id = state.get_mut(id)?.state; Ok(state.get_mut(state_id)?.origin) } } impl StoreContextMut<'_, T> { fn new_transmit>( mut self, kind: TransmitKind, producer: P, ) -> Result> where P::Item: func::Lower, { let token = StoreToken::new(self.as_context_mut()); let state = self.0.concurrent_state_mut(); let (_, read) = state.new_transmit(TransmitOrigin::Host)?; let producer = Arc::new(LockedState::new((Box::pin(producer), P::Buffer::default()))); let id = state.get_mut(read)?.state; let mut dropped = false; let produce = Box::new({ let producer = producer.clone(); move || { let producer = producer.clone(); async move { let mut state = producer.take()?; let (mine, buffer) = &mut *state; let (result, cancelled) = if buffer.remaining().is_empty() { future::poll_fn(|cx| { tls::get(|store| { let transmit = store.concurrent_state_mut().get_mut(id)?; let &WriteState::HostReady { cancel, .. } = &transmit.write else { bail_bug!("expected WriteState::HostReady") }; let mut host_buffer = if let ReadState::HostToHost { buffer, .. } = &mut transmit.read { Some(Cursor::new(mem::take(buffer))) } else { None }; let poll = mine.as_mut().poll_produce( cx, token.as_context_mut(store), Destination { id, buffer, host_buffer: host_buffer.as_mut(), _phantom: PhantomData, }, cancel, ); let transmit = store.concurrent_state_mut().get_mut(id)?; let host_offset = if let ( Some(host_buffer), ReadState::HostToHost { buffer, limit, .. }, ) = (host_buffer, &mut transmit.read) { *limit = usize::try_from(host_buffer.position())?; *buffer = host_buffer.into_inner(); *limit } else { 0 }; { let WriteState::HostReady { guest_offset, cancel, cancel_waker, .. } = &mut transmit.write else { bail_bug!("expected WriteState::HostReady") }; if poll.is_pending() { if !buffer.remaining().is_empty() || *guest_offset > 0 || host_offset > 0 { bail!( "StreamProducer::poll_produce returned Poll::Pending \ after producing at least one item" ) } *cancel_waker = Some(cx.waker().clone()); } else { *cancel_waker = None; *cancel = false; } } Ok(poll.map(|v| v.map(|result| (result, cancel)))) })? }) .await? } else { (StreamResult::Completed, false) }; let (guest_offset, host_offset, count) = tls::get(|store| { let transmit = store.concurrent_state_mut().get_mut(id)?; let (count, host_offset) = match &transmit.read { &ReadState::GuestReady { count, .. } => (count, 0), &ReadState::HostToHost { limit, .. } => (1, limit), _ => bail_bug!("invalid read state"), }; let guest_offset = match &transmit.write { &WriteState::HostReady { guest_offset, .. } => guest_offset, _ => bail_bug!("invalid write state"), }; Ok((guest_offset, host_offset, count)) })?; match result { StreamResult::Completed => { if count > 1 && buffer.remaining().is_empty() && guest_offset == 0 && host_offset == 0 { bail!( "StreamProducer::poll_produce returned StreamResult::Completed \ without producing any items" ); } } StreamResult::Cancelled => { if !cancelled { bail!( "StreamProducer::poll_produce returned StreamResult::Cancelled \ without being given a `finish` parameter value of true" ); } } StreamResult::Dropped => { dropped = true; } } let write_buffer = !buffer.remaining().is_empty() || host_offset > 0; drop(state); if write_buffer { write(token, id, producer.clone(), kind).await?; } Ok(if dropped { if producer.with(|p| p.1.remaining().is_empty())? { StreamResult::Dropped } else { StreamResult::Completed } } else { result }) } .boxed() } }); let try_into = Box::new(move |ty| { let (mine, buffer) = producer.try_lock().ok()?.take()?; match P::try_into(mine, ty) { Ok(value) => Some(value), Err(mine) => { *producer.try_lock().ok()? = Some((mine, buffer)); None } } }); state.get_mut(id)?.write = WriteState::HostReady { produce, try_into, guest_offset: 0, cancel: false, cancel_waker: None, }; Ok(read) } fn set_consumer>( mut self, id: TableId, kind: TransmitKind, consumer: C, ) -> Result<()> { let token = StoreToken::new(self.as_context_mut()); let state = self.0.concurrent_state_mut(); let id = state.get_mut(id)?.state; let transmit = state.get_mut(id)?; let consumer = Arc::new(LockedState::new(Box::pin(consumer))); let consume_with_buffer = { let consumer = consumer.clone(); async move |mut host_buffer: Option<&mut dyn WriteBuffer>| { let mut mine = consumer.take()?; let host_buffer_remaining_before = host_buffer.as_deref_mut().map(|v| v.remaining().len()); let (result, cancelled) = future::poll_fn(|cx| { tls::get(|store| { let cancel = match &store.concurrent_state_mut().get_mut(id)?.read { &ReadState::HostReady { cancel, .. } => cancel, ReadState::Open => false, _ => bail_bug!("unexpected read state"), }; let poll = mine.as_mut().poll_consume( cx, token.as_context_mut(store), Source { id, host_buffer: host_buffer.as_deref_mut(), }, cancel, ); if let ReadState::HostReady { cancel_waker, cancel, .. } = &mut store.concurrent_state_mut().get_mut(id)?.read { if poll.is_pending() { *cancel_waker = Some(cx.waker().clone()); } else { *cancel_waker = None; *cancel = false; } } Ok(poll.map(|v| v.map(|result| (result, cancel)))) })? }) .await?; let (guest_offset, count) = tls::get(|store| { let transmit = store.concurrent_state_mut().get_mut(id)?; Ok(( match &transmit.read { &ReadState::HostReady { guest_offset, .. } => guest_offset, ReadState::Open => 0, _ => bail_bug!("invalid read state"), }, match &transmit.write { &WriteState::GuestReady { count, .. } => count, WriteState::HostReady { .. } => match host_buffer_remaining_before { Some(n) => n, None => bail_bug!("host_buffer_remaining_before should be set"), }, _ => bail_bug!("invalid write state"), }, )) })?; match result { StreamResult::Completed => { if count > 0 && guest_offset == 0 && host_buffer_remaining_before .zip(host_buffer.map(|v| v.remaining().len())) .map(|(before, after)| before == after) .unwrap_or(false) { bail!( "StreamConsumer::poll_consume returned StreamResult::Completed \ without consuming any items" ); } if let TransmitKind::Future = kind { tls::get(|store| { store.concurrent_state_mut().get_mut(id)?.done = true; crate::error::Ok(()) })?; } } StreamResult::Cancelled => { if !cancelled { bail!( "StreamConsumer::poll_consume returned StreamResult::Cancelled \ without being given a `finish` parameter value of true" ); } } StreamResult::Dropped => {} } Ok(result) } }; let consume = { let consume = consume_with_buffer.clone(); Box::new(move || { let consume = consume.clone(); async move { consume(None).await }.boxed() }) }; match &transmit.write { WriteState::Open => { transmit.read = ReadState::HostReady { consume, guest_offset: 0, cancel: false, cancel_waker: None, }; } &WriteState::GuestReady { .. } => { let future = consume(); transmit.read = ReadState::HostReady { consume, guest_offset: 0, cancel: false, cancel_waker: None, }; self.0.pipe_from_guest(kind, id, future); } WriteState::HostReady { .. } => { let WriteState::HostReady { produce, .. } = mem::replace( &mut transmit.write, WriteState::HostReady { produce: Box::new(|| { Box::pin(async { bail_bug!("unexpected invocation of `produce`") }) }), try_into: Box::new(|_| None), guest_offset: 0, cancel: false, cancel_waker: None, }, ) else { bail_bug!("expected WriteState::HostReady") }; transmit.read = ReadState::HostToHost { accept: Box::new(move |input| { let consume = consume_with_buffer.clone(); async move { consume(Some(input.get_mut::())).await }.boxed() }), buffer: Vec::new(), limit: 0, }; let future = async move { loop { if tls::get(|store| { crate::error::Ok(matches!( store.concurrent_state_mut().get_mut(id)?.read, ReadState::Dropped )) })? { break Ok(()); } match produce().await? { StreamResult::Completed | StreamResult::Cancelled => {} StreamResult::Dropped => break Ok(()), } if let TransmitKind::Future = kind { break Ok(()); } } } .map(move |result| { tls::get(|store| store.concurrent_state_mut().delete_transmit(id))?; result }); state.push_future(Box::pin(future)); } WriteState::Dropped => { let reader = transmit.read_handle; self.0.host_drop_reader(reader, kind)?; } } Ok(()) } } async fn write>( token: StoreToken, id: TableId, pair: Arc>, kind: TransmitKind, ) -> Result<()> { let (read, guest_offset) = tls::get(|store| { let transmit = store.concurrent_state_mut().get_mut(id)?; let guest_offset = if let &WriteState::HostReady { guest_offset, .. } = &transmit.write { Some(guest_offset) } else { None }; crate::error::Ok(( mem::replace(&mut transmit.read, ReadState::Open), guest_offset, )) })?; match read { ReadState::GuestReady { ty, flat_abi, options, address, count, handle, instance, caller_instance, caller_thread, } => { let guest_offset = match guest_offset { Some(i) => i, None => bail_bug!("guest_offset should be present if ready"), }; if let TransmitKind::Future = kind { tls::get(|store| { store.concurrent_state_mut().get_mut(id)?.done = true; crate::error::Ok(()) })?; } let old_remaining = pair.with(|p| p.1.remaining().len())?; let accept = { let pair = pair.clone(); move |mut store: StoreContextMut| { let mut state = pair.take()?; lower::( store.as_context_mut(), instance, caller_thread, options, ty, address + (T::SIZE32 * guest_offset), count - guest_offset, &mut state.1, )?; crate::error::Ok(()) } }; if guest_offset < count { if T::MAY_REQUIRE_REALLOC { // For payloads which may require a realloc call, use a // oneshot::channel and background task. This is // necessary because calling the guest while there are // host embedder frames on the stack is unsound. let (tx, rx) = oneshot::channel(); tls::get(move |store| { store .concurrent_state_mut() .push_high_priority(WorkItem::WorkerFunction(AlwaysMut::new(Box::new( move |store| { _ = tx.send(accept(token.as_context_mut(store))?); Ok(()) }, )))) }); rx.await? } else { // Optimize flat payloads (i.e. those which do not // require calling the guest's realloc function) by // lowering directly instead of using a oneshot::channel // and background task. tls::get(|store| accept(token.as_context_mut(store)))? }; } tls::get(|store| { let count = old_remaining - pair.with(|p| p.1.remaining().len())?; let transmit = store.concurrent_state_mut().get_mut(id)?; let WriteState::HostReady { guest_offset, .. } = &mut transmit.write else { bail_bug!("expected WriteState::HostReady") }; *guest_offset += count; transmit.read = ReadState::GuestReady { ty, flat_abi, options, address, count, handle, instance, caller_instance, caller_thread, }; crate::error::Ok(()) })?; Ok(()) } ReadState::HostToHost { accept, mut buffer, limit, } => { let mut state = StreamResult::Completed; let mut position = 0; while !matches!(state, StreamResult::Dropped) && position < limit { let mut slice_buffer = SliceBuffer::new(buffer, position, limit); state = accept(&mut UntypedWriteBuffer::new(&mut slice_buffer)).await?; (buffer, position, _) = slice_buffer.into_parts(); } { let mut pair = pair.take()?; let (_, buffer) = &mut *pair; while !(matches!(state, StreamResult::Dropped) || buffer.remaining().is_empty()) { state = accept(&mut UntypedWriteBuffer::new(buffer)).await?; } } tls::get(|store| { store.concurrent_state_mut().get_mut(id)?.read = match state { StreamResult::Dropped => ReadState::Dropped, StreamResult::Completed | StreamResult::Cancelled => ReadState::HostToHost { accept, buffer, limit: 0, }, }; crate::error::Ok(()) })?; Ok(()) } _ => bail_bug!("unexpected read state"), } } impl Instance { /// Handle a host- or guest-initiated write by delivering the item(s) to the /// `StreamConsumer` for the specified stream or future. fn consume( self, store: &mut dyn VMStore, kind: TransmitKind, transmit_id: TableId, consume: PollStream, guest_offset: usize, cancel: bool, ) -> Result { let mut future = consume(); store.concurrent_state_mut().get_mut(transmit_id)?.read = ReadState::HostReady { consume, guest_offset, cancel, cancel_waker: None, }; let poll = tls::set(store, || { future .as_mut() .poll(&mut Context::from_waker(&Waker::noop())) }); Ok(match poll { Poll::Ready(state) => { let transmit = store.concurrent_state_mut().get_mut(transmit_id)?; let ReadState::HostReady { guest_offset, .. } = &mut transmit.read else { bail_bug!("expected ReadState::HostReady") }; let code = return_code(kind, state?, mem::replace(guest_offset, 0))?; transmit.write = WriteState::Open; code } Poll::Pending => { store.pipe_from_guest(kind, transmit_id, future); ReturnCode::Blocked } }) } /// Handle a host- or guest-initiated read by polling the `StreamProducer` /// for the specified stream or future for items. fn produce( self, store: &mut dyn VMStore, kind: TransmitKind, transmit_id: TableId, produce: PollStream, try_into: TryInto, guest_offset: usize, cancel: bool, ) -> Result { let mut future = produce(); store.concurrent_state_mut().get_mut(transmit_id)?.write = WriteState::HostReady { produce, try_into, guest_offset, cancel, cancel_waker: None, }; let poll = tls::set(store, || { future .as_mut() .poll(&mut Context::from_waker(&Waker::noop())) }); Ok(match poll { Poll::Ready(state) => { let transmit = store.concurrent_state_mut().get_mut(transmit_id)?; let WriteState::HostReady { guest_offset, .. } = &mut transmit.write else { bail_bug!("expected WriteState::HostReady") }; let code = return_code(kind, state?, mem::replace(guest_offset, 0))?; transmit.read = ReadState::Open; code } Poll::Pending => { store.pipe_to_guest(kind, transmit_id, future); ReturnCode::Blocked } }) } /// Drop the writable end of the specified stream or future from the guest. pub(super) fn guest_drop_writable( self, store: &mut StoreOpaque, ty: TransmitIndex, writer: u32, ) -> Result<()> { let table = self.id().get_mut(store).table_for_transmit(ty); let transmit_rep = match ty { TransmitIndex::Future(ty) => table.future_remove_writable(ty, writer)?, TransmitIndex::Stream(ty) => table.stream_remove_writable(ty, writer)?, }; let id = TableId::::new(transmit_rep); log::trace!("guest_drop_writable: drop writer {id:?}"); match ty { TransmitIndex::Stream(_) => store.host_drop_writer(id, None), TransmitIndex::Future(_) => store.host_drop_writer( id, Some(|| { Err(format_err!( "cannot drop future write end without first writing a value" )) }), ), } } /// Copy `count` items from `read_address` to `write_address` for the /// specified stream or future. fn copy( self, store: StoreContextMut, flat_abi: Option, write_caller_instance: RuntimeComponentInstanceIndex, write_ty: TransmitIndex, write_options: OptionsIndex, write_address: usize, read_caller_instance: RuntimeComponentInstanceIndex, read_caller_thread: QualifiedThreadId, read_ty: TransmitIndex, read_options: OptionsIndex, read_address: usize, count: usize, rep: u32, ) -> Result<()> { let (component, mut store) = self.component_and_store_mut(store.0); let types = component.types(); // Validate `write_ty` w.r.t. `write_address` to ensure it's properly // aligned and in-bounds. let write_payload_ty = write_ty.payload(types); let write_abi = match write_payload_ty { Some(ty) => types.canonical_abi(ty), None => &CanonicalAbiInfo::ZERO, }; let write_length_in_bytes = match flat_abi { Some(abi) => usize::try_from(abi.size)? * count, None => usize::try_from(write_abi.size32)? * count, }; if write_length_in_bytes > 0 { if write_address % usize::try_from(write_abi.align32)? != 0 { bail!("write pointer not aligned"); } self.options_memory(store, write_options) .get(write_address..) .and_then(|b| b.get(..write_length_in_bytes)) .ok_or_else(|| crate::format_err!("write pointer out of bounds"))?; } let read_payload_ty = read_ty.payload(types); let read_abi = match read_payload_ty { Some(ty) => types.canonical_abi(ty), None => &CanonicalAbiInfo::ZERO, }; let read_length_in_bytes = match flat_abi { Some(abi) => usize::try_from(abi.size)? * count, None => usize::try_from(read_abi.size32)? * count, }; if read_length_in_bytes > 0 { if read_address % usize::try_from(read_abi.align32)? != 0 { bail!("read pointer not aligned"); } self.options_memory(store, read_options) .get(read_address..) .and_then(|b| b.get(..read_length_in_bytes)) .ok_or_else(|| crate::format_err!("read pointer out of bounds"))?; } if write_caller_instance == read_caller_instance && !allow_intra_component_read_write(write_payload_ty) { bail!( "cannot read from and write to intra-component future/stream with non-numeric payload" ) } match (write_ty, read_ty) { (TransmitIndex::Future(_), TransmitIndex::Future(_)) => { if count != 1 { bail_bug!("futures can only send 1 item"); } let val = write_payload_ty .map(|ty| { let lift = &mut LiftContext::new(store, write_options, self); let bytes = &lift.memory()[write_address..][..write_length_in_bytes]; Val::load(lift, *ty, bytes) }) .transpose()?; if let Some(val) = val { // Serializing the value may require calling the guest's realloc function, so we // set the guest's thread context in case realloc requires it, and restore the original // thread context after the copy is complete. let old_thread = store.set_thread(read_caller_thread)?; let lower = &mut LowerContext::new(store.as_context_mut(), read_options, self); let ptr = func::validate_inbounds_dynamic( read_abi, lower.as_slice_mut(), &ValRaw::u32(read_address.try_into()?), )?; let ty = match read_payload_ty { Some(ty) => ty, None => bail_bug!("expected read payload type to be present"), }; val.store(lower, *ty, ptr)?; store.set_thread(old_thread)?; } } (TransmitIndex::Stream(_), TransmitIndex::Stream(_)) => { if write_length_in_bytes == 0 { return Ok(()); } let write_payload_ty = match write_payload_ty { Some(ty) => ty, None => bail_bug!("expected write payload type to be present"), }; let read_payload_ty = match read_payload_ty { Some(ty) => ty, None => bail_bug!("expected read payload type to be present"), }; if flat_abi.is_some() { // Fast path memcpy for "flat" (i.e. no pointers or handles) payloads: let store_opaque = store.store_opaque_mut(); assert_eq!(read_length_in_bytes, write_length_in_bytes); if write_caller_instance == read_caller_instance { let memory = self.options_memory_mut(store_opaque, read_options); memory.copy_within( write_address..write_address + write_length_in_bytes, read_address, ); } else { let src = self.options_memory(store_opaque, write_options)[write_address..] [..write_length_in_bytes] .as_ptr(); let dst = self.options_memory_mut(store_opaque, read_options) [read_address..][..read_length_in_bytes] .as_mut_ptr(); // SAFETY: Both `src` and `dst` have been validated // above to be valid pointers as they're derived from // slices that have the desired length with the desired // read/write permission. The `unsafe` bit here is that // the memories are disjoint (present in different // instances) and there's no easy way to borrow both // simultaneously from the store. Different instances // are guaranteed to be disjoint, however, so the // `unsafe` here should be ok. unsafe { src.copy_to_nonoverlapping(dst, write_length_in_bytes); } } } else { let store_opaque = store.store_opaque_mut(); let lift = &mut LiftContext::new(store_opaque, write_options, self); let bytes = &lift.memory()[write_address..][..write_length_in_bytes]; lift.consume_fuel_array(count, size_of::())?; let values = (0..count) .map(|index| { let size = usize::try_from(write_abi.size32)?; Val::load(lift, *write_payload_ty, &bytes[(index * size)..][..size]) }) .collect::>>()?; let id = TableId::::new(rep); log::trace!("copy values {values:?} for {id:?}"); // Serializing the value may require calling the guest's realloc function, so we // set the guest's thread context in case realloc requires it, and restore the original // thread context after the copy is complete. let old_thread = store.set_thread(read_caller_thread)?; let lower = &mut LowerContext::new(store.as_context_mut(), read_options, self); let mut ptr = read_address; for value in values { value.store(lower, *read_payload_ty, ptr)?; ptr += usize::try_from(read_abi.size32)?; } store.set_thread(old_thread)?; } } _ => bail_bug!("mismatched transmit types in copy"), } Ok(()) } fn check_bounds( self, store: &StoreOpaque, options: OptionsIndex, ty: TransmitIndex, address: usize, count: usize, ) -> Result<()> { let types = self.id().get(store).component().types(); let size = usize::try_from( match ty { TransmitIndex::Future(ty) => types[types[ty].ty] .payload .map(|ty| types.canonical_abi(&ty).size32), TransmitIndex::Stream(ty) => types[types[ty].ty] .payload .map(|ty| types.canonical_abi(&ty).size32), } .unwrap_or(0), )?; if count > 0 && size > 0 { self.options_memory(store, options) .get(address..) .and_then(|b| b.get(..(size * count))) .map(drop) .ok_or_else(|| crate::format_err!("read pointer out of bounds of memory")) } else { Ok(()) } } /// Write to the specified stream or future from the guest. pub(super) fn guest_write( self, mut store: StoreContextMut, caller: RuntimeComponentInstanceIndex, ty: TransmitIndex, options: OptionsIndex, flat_abi: Option, handle: u32, address: u32, count: u32, ) -> Result { if !self.options(store.0, options).async_ { // The caller may only sync call `{stream,future}.write` from an // async task (i.e. a task created via a call to an async export). // Otherwise, we'll trap. store.0.check_blocking()?; } let address = usize::try_from(address)?; let count = usize::try_from(count)?; self.check_bounds(store.0, options, ty, address, count)?; let (rep, state) = self.id().get_mut(store.0).get_mut_by_index(ty, handle)?; let TransmitLocalState::Write { done } = *state else { bail!(Trap::ConcurrentFutureStreamOp); }; if done { bail!("cannot write to stream after being notified that the readable end dropped"); } *state = TransmitLocalState::Busy; let transmit_handle = TableId::::new(rep); let concurrent_state = store.0.concurrent_state_mut(); let transmit_id = concurrent_state.get_mut(transmit_handle)?.state; let transmit = concurrent_state.get_mut(transmit_id)?; log::trace!( "guest_write {count} to {transmit_handle:?} (handle {handle}; state {transmit_id:?}); {:?}", transmit.read ); if transmit.done { bail!("cannot write to future after previous write succeeded or readable end dropped"); } let new_state = if let ReadState::Dropped = &transmit.read { ReadState::Dropped } else { ReadState::Open }; let set_guest_ready = |me: &mut ConcurrentState| { let transmit = me.get_mut(transmit_id)?; if !matches!(&transmit.write, WriteState::Open) { bail_bug!("expected `WriteState::Open`; got `{:?}`", transmit.write); } transmit.write = WriteState::GuestReady { instance: self, caller, ty, flat_abi, options, address, count, handle, }; Ok::<_, crate::Error>(()) }; let mut result = match mem::replace(&mut transmit.read, new_state) { ReadState::GuestReady { ty: read_ty, flat_abi: read_flat_abi, options: read_options, address: read_address, count: read_count, handle: read_handle, instance: read_instance, caller_instance: read_caller_instance, caller_thread: read_caller_thread, } => { if flat_abi != read_flat_abi { bail_bug!("expected flat ABI calculations to be the same"); } if let TransmitIndex::Future(_) = ty { transmit.done = true; } // Note that zero-length reads and writes are handling specially // by the spec to allow each end to signal readiness to the // other. Quoting the spec: // // ``` // The meaning of a read or write when the length is 0 is that // the caller is querying the "readiness" of the other // side. When a 0-length read/write rendezvous with a // non-0-length read/write, only the 0-length read/write // completes; the non-0-length read/write is kept pending (and // ready for a subsequent rendezvous). // // In the corner case where a 0-length read and write // rendezvous, only the writer is notified of readiness. To // avoid livelock, the Canonical ABI requires that a writer must // (eventually) follow a completed 0-length write with a // non-0-length write that is allowed to block (allowing the // reader end to run and rendezvous with its own non-0-length // read). // ``` let write_complete = count == 0 || read_count > 0; let read_complete = count > 0; let read_buffer_remaining = count < read_count; let read_handle_rep = transmit.read_handle.rep(); let count = count.min(read_count); self.copy( store.as_context_mut(), flat_abi, caller, ty, options, address, read_caller_instance, read_caller_thread, read_ty, read_options, read_address, count, rep, )?; let instance = self.id().get_mut(store.0); let types = instance.component().types(); let item_size = match ty.payload(types) { Some(ty) => usize::try_from(types.canonical_abi(ty).size32)?, None => 0, }; let concurrent_state = store.0.concurrent_state_mut(); if read_complete { let count = u32::try_from(count)?; let total = if let Some(Event::StreamRead { code: ReturnCode::Completed(old_total), .. }) = concurrent_state.take_event(read_handle_rep)? { count + old_total } else { count }; let code = ReturnCode::completed(ty.kind(), total); concurrent_state.send_read_result(read_ty, transmit_id, read_handle, code)?; } if read_buffer_remaining { let transmit = concurrent_state.get_mut(transmit_id)?; transmit.read = ReadState::GuestReady { ty: read_ty, flat_abi: read_flat_abi, options: read_options, address: read_address + (count * item_size), count: read_count - count, handle: read_handle, instance: read_instance, caller_instance: read_caller_instance, caller_thread: read_caller_thread, }; } if write_complete { ReturnCode::completed(ty.kind(), count.try_into()?) } else { set_guest_ready(concurrent_state)?; ReturnCode::Blocked } } ReadState::HostReady { consume, guest_offset, cancel, cancel_waker, } => { if cancel_waker.is_some() { bail_bug!("expected cancel_waker to be none"); } if cancel { bail_bug!("expected cancel to be false"); } if guest_offset != 0 { bail_bug!("expected guest_offset to be 0"); } if let TransmitIndex::Future(_) = ty { transmit.done = true; } set_guest_ready(concurrent_state)?; self.consume(store.0, ty.kind(), transmit_id, consume, 0, false)? } ReadState::HostToHost { .. } => bail_bug!("unexpected HostToHost"), ReadState::Open => { set_guest_ready(concurrent_state)?; ReturnCode::Blocked } ReadState::Dropped => { if let TransmitIndex::Future(_) = ty { transmit.done = true; } ReturnCode::Dropped(0) } }; if result == ReturnCode::Blocked && !self.options(store.0, options).async_ { result = self.wait_for_write(store.0, transmit_handle)?; } if result != ReturnCode::Blocked { *self.id().get_mut(store.0).get_mut_by_index(ty, handle)?.1 = TransmitLocalState::Write { done: matches!( (result, ty), (ReturnCode::Dropped(_), TransmitIndex::Stream(_)) ), }; } log::trace!( "guest_write result for {transmit_handle:?} (handle {handle}; state {transmit_id:?}): {result:?}", ); Ok(result) } /// Read from the specified stream or future from the guest. pub(super) fn guest_read( self, mut store: StoreContextMut, caller_instance: RuntimeComponentInstanceIndex, ty: TransmitIndex, options: OptionsIndex, flat_abi: Option, handle: u32, address: u32, count: u32, ) -> Result { if !self.options(store.0, options).async_ { // The caller may only sync call `{stream,future}.read` from an // async task (i.e. a task created via a call to an async export). // Otherwise, we'll trap. store.0.check_blocking()?; } let address = usize::try_from(address)?; let count = usize::try_from(count)?; self.check_bounds(store.0, options, ty, address, count)?; let (rep, state) = self.id().get_mut(store.0).get_mut_by_index(ty, handle)?; let TransmitLocalState::Read { done } = *state else { bail!(Trap::ConcurrentFutureStreamOp); }; if done { bail!("cannot read from stream after being notified that the writable end dropped"); } *state = TransmitLocalState::Busy; let transmit_handle = TableId::::new(rep); let concurrent_state = store.0.concurrent_state_mut(); let caller_thread = concurrent_state.current_guest_thread()?; let transmit_id = concurrent_state.get_mut(transmit_handle)?.state; let transmit = concurrent_state.get_mut(transmit_id)?; log::trace!( "guest_read {count} from {transmit_handle:?} (handle {handle}; state {transmit_id:?}); {:?}", transmit.write ); if transmit.done { bail!("cannot read from future after previous read succeeded"); } let new_state = if let WriteState::Dropped = &transmit.write { WriteState::Dropped } else { WriteState::Open }; let set_guest_ready = |me: &mut ConcurrentState| { let transmit = me.get_mut(transmit_id)?; if !matches!(&transmit.read, ReadState::Open) { bail_bug!("expected `ReadState::Open`; got `{:?}`", transmit.read); } transmit.read = ReadState::GuestReady { ty, flat_abi, options, address, count, handle, instance: self, caller_instance, caller_thread, }; Ok::<_, crate::Error>(()) }; let mut result = match mem::replace(&mut transmit.write, new_state) { WriteState::GuestReady { instance: _, ty: write_ty, flat_abi: write_flat_abi, options: write_options, address: write_address, count: write_count, handle: write_handle, caller: write_caller, } => { if flat_abi != write_flat_abi { bail_bug!("expected flat ABI calculations to be the same"); } if let TransmitIndex::Future(_) = ty { transmit.done = true; } let write_handle_rep = transmit.write_handle.rep(); // See the comment in `guest_write` for the // `ReadState::GuestReady` case concerning zero-length reads and // writes. let write_complete = write_count == 0 || count > 0; let read_complete = write_count > 0; let write_buffer_remaining = count < write_count; let count = count.min(write_count); self.copy( store.as_context_mut(), flat_abi, write_caller, write_ty, write_options, write_address, caller_instance, caller_thread, ty, options, address, count, rep, )?; let instance = self.id().get_mut(store.0); let types = instance.component().types(); let item_size = match ty.payload(types) { Some(ty) => usize::try_from(types.canonical_abi(ty).size32)?, None => 0, }; let concurrent_state = store.0.concurrent_state_mut(); if write_complete { let count = u32::try_from(count)?; let total = if let Some(Event::StreamWrite { code: ReturnCode::Completed(old_total), .. }) = concurrent_state.take_event(write_handle_rep)? { count + old_total } else { count }; let code = ReturnCode::completed(ty.kind(), total); concurrent_state.send_write_result( write_ty, transmit_id, write_handle, code, )?; } if write_buffer_remaining { let transmit = concurrent_state.get_mut(transmit_id)?; transmit.write = WriteState::GuestReady { instance: self, caller: write_caller, ty: write_ty, flat_abi: write_flat_abi, options: write_options, address: write_address + (count * item_size), count: write_count - count, handle: write_handle, }; } if read_complete { ReturnCode::completed(ty.kind(), count.try_into()?) } else { set_guest_ready(concurrent_state)?; ReturnCode::Blocked } } WriteState::HostReady { produce, try_into, guest_offset, cancel, cancel_waker, } => { if cancel_waker.is_some() { bail_bug!("expected cancel_waker to be none"); } if cancel { bail_bug!("expected cancel to be false"); } if guest_offset != 0 { bail_bug!("expected guest_offset to be 0"); } set_guest_ready(concurrent_state)?; let code = self.produce(store.0, ty.kind(), transmit_id, produce, try_into, 0, false)?; if let (TransmitIndex::Future(_), ReturnCode::Completed(_)) = (ty, code) { store.0.concurrent_state_mut().get_mut(transmit_id)?.done = true; } code } WriteState::Open => { set_guest_ready(concurrent_state)?; ReturnCode::Blocked } WriteState::Dropped => ReturnCode::Dropped(0), }; if result == ReturnCode::Blocked && !self.options(store.0, options).async_ { result = self.wait_for_read(store.0, transmit_handle)?; } if result != ReturnCode::Blocked { *self.id().get_mut(store.0).get_mut_by_index(ty, handle)?.1 = TransmitLocalState::Read { done: matches!( (result, ty), (ReturnCode::Dropped(_), TransmitIndex::Stream(_)) ), }; } log::trace!( "guest_read result for {transmit_handle:?} (handle {handle}; state {transmit_id:?}): {result:?}", ); Ok(result) } fn wait_for_write( self, store: &mut StoreOpaque, handle: TableId, ) -> Result { let waitable = Waitable::Transmit(handle); store.wait_for_event(waitable)?; let event = waitable.take_event(store.concurrent_state_mut())?; if let Some(event @ (Event::StreamWrite { code, .. } | Event::FutureWrite { code, .. })) = event { waitable.on_delivery(store, self, event)?; Ok(code) } else { bail_bug!("expected either a stream or future write event") } } /// Cancel a pending stream or future write. fn cancel_write( self, store: &mut StoreOpaque, transmit_id: TableId, async_: bool, ) -> Result { let state = store.concurrent_state_mut(); let transmit = state.get_mut(transmit_id)?; log::trace!( "host_cancel_write state {transmit_id:?}; write state {:?} read state {:?}", transmit.read, transmit.write ); let waitable = Waitable::Transmit(transmit.write_handle); let code = if let Some(event) = waitable.take_event(state)? { let (Event::FutureWrite { code, .. } | Event::StreamWrite { code, .. }) = event else { bail_bug!("expected either a stream or future write event") }; waitable.on_delivery(store, self, event)?; match (code, event) { (ReturnCode::Completed(count), Event::StreamWrite { .. }) => { ReturnCode::Cancelled(count) } (ReturnCode::Dropped(_) | ReturnCode::Completed(_), _) => code, _ => bail_bug!("unexpected code/event combo"), } } else if let ReadState::HostReady { cancel, cancel_waker, .. } = &mut state.get_mut(transmit_id)?.read { *cancel = true; if let Some(waker) = cancel_waker.take() { waker.wake(); } if async_ { ReturnCode::Blocked } else { let handle = store .concurrent_state_mut() .get_mut(transmit_id)? .write_handle; self.wait_for_write(store, handle)? } } else { ReturnCode::Cancelled(0) }; if !matches!(code, ReturnCode::Blocked) { let transmit = store.concurrent_state_mut().get_mut(transmit_id)?; match &transmit.write { WriteState::GuestReady { .. } => { transmit.write = WriteState::Open; } WriteState::HostReady { .. } => bail_bug!("support host write cancellation"), WriteState::Open | WriteState::Dropped => {} } } log::trace!("cancelled write {transmit_id:?}: {code:?}"); Ok(code) } fn wait_for_read( self, store: &mut StoreOpaque, handle: TableId, ) -> Result { let waitable = Waitable::Transmit(handle); store.wait_for_event(waitable)?; let event = waitable.take_event(store.concurrent_state_mut())?; if let Some(event @ (Event::StreamRead { code, .. } | Event::FutureRead { code, .. })) = event { waitable.on_delivery(store, self, event)?; Ok(code) } else { bail_bug!("expected either a stream or future read event") } } /// Cancel a pending stream or future read. fn cancel_read( self, store: &mut StoreOpaque, transmit_id: TableId, async_: bool, ) -> Result { let state = store.concurrent_state_mut(); let transmit = state.get_mut(transmit_id)?; log::trace!( "host_cancel_read state {transmit_id:?}; read state {:?} write state {:?}", transmit.read, transmit.write ); let waitable = Waitable::Transmit(transmit.read_handle); let code = if let Some(event) = waitable.take_event(state)? { let (Event::FutureRead { code, .. } | Event::StreamRead { code, .. }) = event else { bail_bug!("expected either a stream or future read event") }; waitable.on_delivery(store, self, event)?; match (code, event) { (ReturnCode::Completed(count), Event::StreamRead { .. }) => { ReturnCode::Cancelled(count) } (ReturnCode::Dropped(_) | ReturnCode::Completed(_), _) => code, _ => bail_bug!("unexpected code/event combo"), } } else if let WriteState::HostReady { cancel, cancel_waker, .. } = &mut state.get_mut(transmit_id)?.write { *cancel = true; if let Some(waker) = cancel_waker.take() { waker.wake(); } if async_ { ReturnCode::Blocked } else { let handle = store .concurrent_state_mut() .get_mut(transmit_id)? .read_handle; self.wait_for_read(store, handle)? } } else { ReturnCode::Cancelled(0) }; if !matches!(code, ReturnCode::Blocked) { let transmit = store.concurrent_state_mut().get_mut(transmit_id)?; match &transmit.read { ReadState::GuestReady { .. } => { transmit.read = ReadState::Open; } ReadState::HostReady { .. } | ReadState::HostToHost { .. } => { bail_bug!("support host read cancellation") } ReadState::Open | ReadState::Dropped => {} } } log::trace!("cancelled read {transmit_id:?}: {code:?}"); Ok(code) } /// Cancel a pending write for the specified stream or future from the guest. fn guest_cancel_write( self, store: &mut StoreOpaque, ty: TransmitIndex, async_: bool, writer: u32, ) -> Result { if !async_ { // The caller may only sync call `{stream,future}.cancel-write` from // an async task (i.e. a task created via a call to an async // export). Otherwise, we'll trap. store.check_blocking()?; } let (rep, state) = get_mut_by_index_from(self.id().get_mut(store).table_for_transmit(ty), ty, writer)?; let id = TableId::::new(rep); log::trace!("guest cancel write {id:?} (handle {writer})"); match state { TransmitLocalState::Write { .. } => { bail!("stream or future write cancelled when no write is pending") } TransmitLocalState::Read { .. } => { bail!("passed read end to `{{stream|future}}.cancel-write`") } TransmitLocalState::Busy => {} } let transmit_id = store.concurrent_state_mut().get_mut(id)?.state; let code = self.cancel_write(store, transmit_id, async_)?; if !matches!(code, ReturnCode::Blocked) { let state = get_mut_by_index_from(self.id().get_mut(store).table_for_transmit(ty), ty, writer)? .1; if let TransmitLocalState::Busy = state { *state = TransmitLocalState::Write { done: false }; } } Ok(code) } /// Cancel a pending read for the specified stream or future from the guest. fn guest_cancel_read( self, store: &mut StoreOpaque, ty: TransmitIndex, async_: bool, reader: u32, ) -> Result { if !async_ { // The caller may only sync call `{stream,future}.cancel-read` from // an async task (i.e. a task created via a call to an async // export). Otherwise, we'll trap. store.check_blocking()?; } let (rep, state) = get_mut_by_index_from(self.id().get_mut(store).table_for_transmit(ty), ty, reader)?; let id = TableId::::new(rep); log::trace!("guest cancel read {id:?} (handle {reader})"); match state { TransmitLocalState::Read { .. } => { bail!("stream or future read cancelled when no read is pending") } TransmitLocalState::Write { .. } => { bail!("passed write end to `{{stream|future}}.cancel-read`") } TransmitLocalState::Busy => {} } let transmit_id = store.concurrent_state_mut().get_mut(id)?.state; let code = self.cancel_read(store, transmit_id, async_)?; if !matches!(code, ReturnCode::Blocked) { let state = get_mut_by_index_from(self.id().get_mut(store).table_for_transmit(ty), ty, reader)? .1; if let TransmitLocalState::Busy = state { *state = TransmitLocalState::Read { done: false }; } } Ok(code) } /// Drop the readable end of the specified stream or future from the guest. fn guest_drop_readable( self, store: &mut StoreOpaque, ty: TransmitIndex, reader: u32, ) -> Result<()> { let table = self.id().get_mut(store).table_for_transmit(ty); let (rep, _is_done) = match ty { TransmitIndex::Stream(ty) => table.stream_remove_readable(ty, reader)?, TransmitIndex::Future(ty) => table.future_remove_readable(ty, reader)?, }; let kind = match ty { TransmitIndex::Stream(_) => TransmitKind::Stream, TransmitIndex::Future(_) => TransmitKind::Future, }; let id = TableId::::new(rep); log::trace!("guest_drop_readable: drop reader {id:?}"); store.host_drop_reader(id, kind) } /// Create a new error context for the given component. pub(crate) fn error_context_new( self, store: &mut StoreOpaque, ty: TypeComponentLocalErrorContextTableIndex, options: OptionsIndex, debug_msg_address: u32, debug_msg_len: u32, ) -> Result { let lift_ctx = &mut LiftContext::new(store, options, self); let debug_msg = String::linear_lift_from_flat( lift_ctx, InterfaceType::String, &[ValRaw::u32(debug_msg_address), ValRaw::u32(debug_msg_len)], )?; // Create a new ErrorContext that is tracked along with other concurrent state let err_ctx = ErrorContextState { debug_msg }; let state = store.concurrent_state_mut(); let table_id = state.push(err_ctx)?; let global_ref_count_idx = TypeComponentGlobalErrorContextTableIndex::from_u32(table_id.rep()); // Add to the global error context ref counts let _ = state .global_error_context_ref_counts .insert(global_ref_count_idx, GlobalErrorContextRefCount(1)); // Error context are tracked both locally (to a single component instance) and globally // the counts for both must stay in sync. // // Here we reflect the newly created global concurrent error context state into the // component instance's locally tracked count, along with the appropriate key into the global // ref tracking data structures to enable later lookup let local_idx = self .id() .get_mut(store) .table_for_error_context(ty) .error_context_insert(table_id.rep())?; Ok(local_idx) } /// Retrieve the debug message from the specified error context. pub(super) fn error_context_debug_message( self, store: StoreContextMut, ty: TypeComponentLocalErrorContextTableIndex, options: OptionsIndex, err_ctx_handle: u32, debug_msg_address: u32, ) -> Result<()> { // Retrieve the error context and internal debug message let handle_table_id_rep = self .id() .get_mut(store.0) .table_for_error_context(ty) .error_context_rep(err_ctx_handle)?; let state = store.0.concurrent_state_mut(); // Get the state associated with the error context let ErrorContextState { debug_msg } = state.get_mut(TableId::::new(handle_table_id_rep))?; let debug_msg = debug_msg.clone(); let lower_cx = &mut LowerContext::new(store, options, self); let debug_msg_address = usize::try_from(debug_msg_address)?; // Lower the string into the component's memory. // // Note that the "8" here is the size of a WIT `string` in linear // memory, the ptr+length. This'll need to be updated when `memory64` // comes along. (FIXME(#4311)) let offset = lower_cx .as_slice_mut() .get(debug_msg_address..) .and_then(|b| b.get(..8)) .map(|_| debug_msg_address) .ok_or_else(|| crate::format_err!("invalid debug message pointer: out of bounds"))?; debug_msg .as_str() .linear_lower_to_memory(lower_cx, InterfaceType::String, offset)?; Ok(()) } /// Implements the `future.cancel-read` intrinsic. pub(crate) fn future_cancel_read( self, store: &mut StoreOpaque, ty: TypeFutureTableIndex, async_: bool, reader: u32, ) -> Result { self.guest_cancel_read(store, TransmitIndex::Future(ty), async_, reader) .map(|v| v.encode()) } /// Implements the `future.cancel-write` intrinsic. pub(crate) fn future_cancel_write( self, store: &mut StoreOpaque, ty: TypeFutureTableIndex, async_: bool, writer: u32, ) -> Result { self.guest_cancel_write(store, TransmitIndex::Future(ty), async_, writer) .map(|v| v.encode()) } /// Implements the `stream.cancel-read` intrinsic. pub(crate) fn stream_cancel_read( self, store: &mut StoreOpaque, ty: TypeStreamTableIndex, async_: bool, reader: u32, ) -> Result { self.guest_cancel_read(store, TransmitIndex::Stream(ty), async_, reader) .map(|v| v.encode()) } /// Implements the `stream.cancel-write` intrinsic. pub(crate) fn stream_cancel_write( self, store: &mut StoreOpaque, ty: TypeStreamTableIndex, async_: bool, writer: u32, ) -> Result { self.guest_cancel_write(store, TransmitIndex::Stream(ty), async_, writer) .map(|v| v.encode()) } /// Implements the `future.drop-readable` intrinsic. pub(crate) fn future_drop_readable( self, store: &mut StoreOpaque, ty: TypeFutureTableIndex, reader: u32, ) -> Result<()> { self.guest_drop_readable(store, TransmitIndex::Future(ty), reader) } /// Implements the `stream.drop-readable` intrinsic. pub(crate) fn stream_drop_readable( self, store: &mut StoreOpaque, ty: TypeStreamTableIndex, reader: u32, ) -> Result<()> { self.guest_drop_readable(store, TransmitIndex::Stream(ty), reader) } /// Allocate a new future or stream and grant ownership of both the read and /// write ends to the (sub-)component instance to which the specified /// `TransmitIndex` belongs. fn guest_new(self, store: &mut StoreOpaque, ty: TransmitIndex) -> Result { let (write, read) = store .concurrent_state_mut() .new_transmit(TransmitOrigin::guest(self.id().instance(), ty))?; let table = self.id().get_mut(store).table_for_transmit(ty); let (read_handle, write_handle) = match ty { TransmitIndex::Future(ty) => ( table.future_insert_read(ty, read.rep())?, table.future_insert_write(ty, write.rep())?, ), TransmitIndex::Stream(ty) => ( table.stream_insert_read(ty, read.rep())?, table.stream_insert_write(ty, write.rep())?, ), }; let state = store.concurrent_state_mut(); state.get_mut(read)?.common.handle = Some(read_handle); state.get_mut(write)?.common.handle = Some(write_handle); Ok(ResourcePair { write: write_handle, read: read_handle, }) } /// Drop the specified error context. pub(crate) fn error_context_drop( self, store: &mut StoreOpaque, ty: TypeComponentLocalErrorContextTableIndex, error_context: u32, ) -> Result<()> { let instance = self.id().get_mut(store); let local_handle_table = instance.table_for_error_context(ty); let rep = local_handle_table.error_context_drop(error_context)?; let global_ref_count_idx = TypeComponentGlobalErrorContextTableIndex::from_u32(rep); let state = store.concurrent_state_mut(); let Some(GlobalErrorContextRefCount(global_ref_count)) = state .global_error_context_ref_counts .get_mut(&global_ref_count_idx) else { bail_bug!("retrieve concurrent state for error context during drop") }; // Reduce the component-global ref count, removing tracking if necessary if *global_ref_count < 1 { bail_bug!("ref count unexpectedly zero"); } *global_ref_count -= 1; if *global_ref_count == 0 { state .global_error_context_ref_counts .remove(&global_ref_count_idx); state .delete(TableId::::new(rep)) .context("deleting component-global error context data")?; } Ok(()) } /// Transfer ownership of the specified stream or future read end from one /// guest to another. fn guest_transfer( self, store: &mut StoreOpaque, src_idx: u32, src: TransmitIndex, dst: TransmitIndex, ) -> Result { let mut instance = self.id().get_mut(store); let src_table = instance.as_mut().table_for_transmit(src); let (rep, is_done) = match src { TransmitIndex::Future(idx) => src_table.future_remove_readable(idx, src_idx)?, TransmitIndex::Stream(idx) => src_table.stream_remove_readable(idx, src_idx)?, }; if is_done { bail!("cannot lift after being notified that the writable end dropped"); } let dst_table = instance.table_for_transmit(dst); let handle = match dst { TransmitIndex::Future(idx) => dst_table.future_insert_read(idx, rep), TransmitIndex::Stream(idx) => dst_table.stream_insert_read(idx, rep), }?; store .concurrent_state_mut() .get_mut(TableId::::new(rep))? .common .handle = Some(handle); Ok(handle) } /// Implements the `future.new` intrinsic. pub(crate) fn future_new( self, store: &mut StoreOpaque, ty: TypeFutureTableIndex, ) -> Result { self.guest_new(store, TransmitIndex::Future(ty)) } /// Implements the `stream.new` intrinsic. pub(crate) fn stream_new( self, store: &mut StoreOpaque, ty: TypeStreamTableIndex, ) -> Result { self.guest_new(store, TransmitIndex::Stream(ty)) } /// Transfer ownership of the specified future read end from one guest to /// another. pub(crate) fn future_transfer( self, store: &mut StoreOpaque, src_idx: u32, src: TypeFutureTableIndex, dst: TypeFutureTableIndex, ) -> Result { self.guest_transfer( store, src_idx, TransmitIndex::Future(src), TransmitIndex::Future(dst), ) } /// Transfer ownership of the specified stream read end from one guest to /// another. pub(crate) fn stream_transfer( self, store: &mut StoreOpaque, src_idx: u32, src: TypeStreamTableIndex, dst: TypeStreamTableIndex, ) -> Result { self.guest_transfer( store, src_idx, TransmitIndex::Stream(src), TransmitIndex::Stream(dst), ) } /// Copy the specified error context from one component to another. pub(crate) fn error_context_transfer( self, store: &mut StoreOpaque, src_idx: u32, src: TypeComponentLocalErrorContextTableIndex, dst: TypeComponentLocalErrorContextTableIndex, ) -> Result { let mut instance = self.id().get_mut(store); let rep = instance .as_mut() .table_for_error_context(src) .error_context_rep(src_idx)?; let dst_idx = instance .table_for_error_context(dst) .error_context_insert(rep)?; // Update the global (cross-subcomponent) count for error contexts // as the new component has essentially created a new reference that will // be dropped/handled independently let global_ref_count = store .concurrent_state_mut() .global_error_context_ref_counts .get_mut(&TypeComponentGlobalErrorContextTableIndex::from_u32(rep)) .context("global ref count present for existing (sub)component error context")?; global_ref_count.0 += 1; Ok(dst_idx) } } impl ComponentInstance { fn table_for_transmit(self: Pin<&mut Self>, ty: TransmitIndex) -> &mut HandleTable { let (states, types) = self.instance_states(); let runtime_instance = match ty { TransmitIndex::Stream(ty) => types[ty].instance, TransmitIndex::Future(ty) => types[ty].instance, }; states[runtime_instance].handle_table() } fn table_for_error_context( self: Pin<&mut Self>, ty: TypeComponentLocalErrorContextTableIndex, ) -> &mut HandleTable { let (states, types) = self.instance_states(); let runtime_instance = types[ty].instance; states[runtime_instance].handle_table() } fn get_mut_by_index( self: Pin<&mut Self>, ty: TransmitIndex, index: u32, ) -> Result<(u32, &mut TransmitLocalState)> { get_mut_by_index_from(self.table_for_transmit(ty), ty, index) } } impl ConcurrentState { fn send_write_result( &mut self, ty: TransmitIndex, id: TableId, handle: u32, code: ReturnCode, ) -> Result<()> { let write_handle = self.get_mut(id)?.write_handle.rep(); self.set_event( write_handle, match ty { TransmitIndex::Future(ty) => Event::FutureWrite { code, pending: Some((ty, handle)), }, TransmitIndex::Stream(ty) => Event::StreamWrite { code, pending: Some((ty, handle)), }, }, ) } fn send_read_result( &mut self, ty: TransmitIndex, id: TableId, handle: u32, code: ReturnCode, ) -> Result<()> { let read_handle = self.get_mut(id)?.read_handle.rep(); self.set_event( read_handle, match ty { TransmitIndex::Future(ty) => Event::FutureRead { code, pending: Some((ty, handle)), }, TransmitIndex::Stream(ty) => Event::StreamRead { code, pending: Some((ty, handle)), }, }, ) } fn take_event(&mut self, waitable: u32) -> Result> { Waitable::Transmit(TableId::::new(waitable)).take_event(self) } fn set_event(&mut self, waitable: u32, event: Event) -> Result<()> { Waitable::Transmit(TableId::::new(waitable)).set_event(self, Some(event)) } /// Set or update the event for the specified waitable. /// /// If there is already an event set for this waitable, we assert that it is /// of the same variant as the new one and reuse the `ReturnCode` count and /// the `pending` field if applicable. // TODO: This is a bit awkward due to how // `Event::{Stream,Future}{Write,Read}` and // `ReturnCode::{Completed,Dropped,Cancelled}` are currently represented. // Consider updating those representations in a way that allows this // function to be simplified. fn update_event(&mut self, waitable: u32, event: Event) -> Result<()> { let waitable = Waitable::Transmit(TableId::::new(waitable)); fn update_code(old: ReturnCode, new: ReturnCode) -> Result { let (ReturnCode::Completed(count) | ReturnCode::Dropped(count) | ReturnCode::Cancelled(count)) = old else { bail_bug!("unexpected old return code") }; Ok(match new { ReturnCode::Dropped(0) => ReturnCode::Dropped(count), ReturnCode::Cancelled(0) => ReturnCode::Cancelled(count), _ => bail_bug!("unexpected new return code"), }) } let event = match (waitable.take_event(self)?, event) { (None, _) => event, (Some(old @ Event::FutureWrite { .. }), Event::FutureWrite { .. }) => old, (Some(old @ Event::FutureRead { .. }), Event::FutureRead { .. }) => old, ( Some(Event::StreamWrite { code: old_code, pending: old_pending, }), Event::StreamWrite { code, pending }, ) => Event::StreamWrite { code: update_code(old_code, code)?, pending: old_pending.or(pending), }, ( Some(Event::StreamRead { code: old_code, pending: old_pending, }), Event::StreamRead { code, pending }, ) => Event::StreamRead { code: update_code(old_code, code)?, pending: old_pending.or(pending), }, _ => bail_bug!("unexpected event combination"), }; waitable.set_event(self, Some(event)) } /// Allocate a new future or stream, including the `TransmitState` and the /// `TransmitHandle`s corresponding to the read and write ends. fn new_transmit( &mut self, origin: TransmitOrigin, ) -> Result<(TableId, TableId)> { let state_id = self.push(TransmitState::new(origin))?; let write = self.push(TransmitHandle::new(state_id))?; let read = self.push(TransmitHandle::new(state_id))?; let state = self.get_mut(state_id)?; state.write_handle = write; state.read_handle = read; log::trace!("new transmit: state {state_id:?}; write {write:?}; read {read:?}",); Ok((write, read)) } /// Delete the specified future or stream, including the read and write ends. fn delete_transmit(&mut self, state_id: TableId) -> Result<()> { let state = self.delete(state_id)?; self.delete(state.write_handle)?; self.delete(state.read_handle)?; log::trace!( "delete transmit: state {state_id:?}; write {:?}; read {:?}", state.write_handle, state.read_handle, ); Ok(()) } } pub(crate) struct ResourcePair { pub(crate) write: u32, pub(crate) read: u32, } impl Waitable { /// Handle the imminent delivery of the specified event, e.g. by updating /// the state of the stream or future. pub(super) fn on_delivery( &self, store: &mut StoreOpaque, instance: Instance, event: Event, ) -> Result<()> { let instance = instance.id().get_mut(store); let (rep, state, code) = match event { Event::FutureRead { pending: Some((ty, handle)), code, } | Event::FutureWrite { pending: Some((ty, handle)), code, } => { let runtime_instance = instance.component().types()[ty].instance; let (rep, state) = instance.instance_states().0[runtime_instance] .handle_table() .future_rep(ty, handle)?; (rep, state, code) } Event::StreamRead { pending: Some((ty, handle)), code, } | Event::StreamWrite { pending: Some((ty, handle)), code, } => { let runtime_instance = instance.component().types()[ty].instance; let (rep, state) = instance.instance_states().0[runtime_instance] .handle_table() .stream_rep(ty, handle)?; (rep, state, code) } _ => return Ok(()), }; if rep != self.rep() { bail_bug!("unexpected rep mismatch"); } if *state != TransmitLocalState::Busy { bail_bug!("expected state to be busy"); } let done = matches!(code, ReturnCode::Dropped(_)); *state = match event { Event::FutureRead { .. } | Event::StreamRead { .. } => { TransmitLocalState::Read { done } } Event::FutureWrite { .. } | Event::StreamWrite { .. } => { TransmitLocalState::Write { done } } _ => bail_bug!("unexpected event for stream"), }; let transmit_handle = TableId::::new(rep); let state = store.concurrent_state_mut(); let transmit_id = state.get_mut(transmit_handle)?.state; let transmit = state.get_mut(transmit_id)?; match event { Event::StreamRead { .. } => { transmit.read = ReadState::Open; } Event::StreamWrite { .. } => transmit.write = WriteState::Open, _ => {} } Ok(()) } } /// Determine whether an intra-component read/write is allowed for the specified /// `stream` or `future` payload type according to the component model /// specification. fn allow_intra_component_read_write(ty: Option<&InterfaceType>) -> bool { matches!( ty, None | Some( InterfaceType::S8 | InterfaceType::U8 | InterfaceType::S16 | InterfaceType::U16 | InterfaceType::S32 | InterfaceType::U32 | InterfaceType::S64 | InterfaceType::U64 | InterfaceType::Float32 | InterfaceType::Float64 ) ) } /// Helper structure to manage moving a `T` in/out of an interior `Mutex` which /// contains an /// `Option` struct LockedState { inner: Mutex>, } impl LockedState { /// Creates a new initial state with `value` stored. fn new(value: T) -> Self { Self { inner: Mutex::new(Some(value)), } } /// Attempts to lock the inner mutex and return its guard. /// /// # Errors /// /// Fails if this lock is either poisoned or if it's currently locked. /// As-used in this file there should never actually be contention on this /// lock nor recursive access so failing to acquire the lock is a fatal /// error that gets propagated upwards. fn try_lock(&self) -> Result>> { match self.inner.try_lock() { Ok(lock) => Ok(lock), Err(_) => bail_bug!("should not have contention on state lock"), } } /// Takes the inner `T` out of this state, returning it as a guard which /// will put it back when finished. /// /// # Errors /// /// Returns an error if the state `T` isn't present. fn take(&self) -> Result> { let result = self.try_lock()?.take(); match result { Some(result) => Ok(LockedStateGuard { value: ManuallyDrop::new(result), state: self, }), None => bail_bug!("lock value unexpectedly missing"), } } /// Performs the operation `f` on the inner state `&mut T`. /// /// This will acquire the internal lock and invoke `f`, so `f` should not /// expect to be able to recursively acquire this lock. /// /// # Errors /// /// Returns an error if the state `T` isn't present. fn with(&self, f: impl FnOnce(&mut T) -> R) -> Result { let mut inner = self.try_lock()?; match &mut *inner { Some(state) => Ok(f(state)), None => bail_bug!("lock value unexpectedly missing"), } } } /// Helper structure returned from [`LockedState::take`] which will put the /// state specified by `value` back into the original lock once this is dropped. struct LockedStateGuard<'a, T> { value: ManuallyDrop, state: &'a LockedState, } impl Deref for LockedStateGuard<'_, T> { type Target = T; fn deref(&self) -> &T { &self.value } } impl DerefMut for LockedStateGuard<'_, T> { fn deref_mut(&mut self) -> &mut T { &mut self.value } } impl Drop for LockedStateGuard<'_, T> { fn drop(&mut self) { // SAFETY: `ManuallyDrop::take` requires that after invoked the // original value is not read. This is the `Drop` for this type which // means we have exclusive ownership and it is not read further in the // destructor, satisfying this requirement. let value = unsafe { ManuallyDrop::take(&mut self.value) }; // If this fails due to contention that's a bug, but we're not in a // position to panic due to this being a destructor nor return an error, // so defer the bug to showing up later. if let Ok(mut lock) = self.state.try_lock() { *lock = Some(value); } } } #[cfg(test)] mod tests { use super::*; use crate::{Engine, Store}; use core::future::pending; use core::pin::pin; use std::sync::LazyLock; static ENGINE: LazyLock = LazyLock::new(Engine::default); fn poll_future_producer(rx: Pin<&mut T>, finish: bool) -> Poll>> where T: FutureProducer<()>, { rx.poll_produce( &mut Context::from_waker(Waker::noop()), Store::new(&ENGINE, ()).as_context_mut(), finish, ) } #[test] fn future_producer() { let mut fut = pin!(async { crate::error::Ok(()) }); assert!(matches!( poll_future_producer(fut.as_mut(), false), Poll::Ready(Ok(Some(()))), )); let mut fut = pin!(async { crate::error::Ok(()) }); assert!(matches!( poll_future_producer(fut.as_mut(), true), Poll::Ready(Ok(Some(()))), )); let mut fut = pin!(pending::>()); assert!(matches!( poll_future_producer(fut.as_mut(), false), Poll::Pending, )); assert!(matches!( poll_future_producer(fut.as_mut(), true), Poll::Ready(Ok(None)), )); let (tx, rx) = oneshot::channel(); let mut rx = pin!(rx); assert!(matches!( poll_future_producer(rx.as_mut(), false), Poll::Pending, )); assert!(matches!( poll_future_producer(rx.as_mut(), true), Poll::Ready(Ok(None)), )); tx.send(()).unwrap(); assert!(matches!( poll_future_producer(rx.as_mut(), true), Poll::Ready(Ok(Some(()))), )); let (tx, rx) = oneshot::channel(); let mut rx = pin!(rx); tx.send(()).unwrap(); assert!(matches!( poll_future_producer(rx.as_mut(), false), Poll::Ready(Ok(Some(()))), )); let (tx, rx) = oneshot::channel::<()>(); let mut rx = pin!(rx); drop(tx); assert!(matches!( poll_future_producer(rx.as_mut(), false), Poll::Ready(Err(..)), )); let (tx, rx) = oneshot::channel::<()>(); let mut rx = pin!(rx); drop(tx); assert!(matches!( poll_future_producer(rx.as_mut(), true), Poll::Ready(Err(..)), )); } }