From f8bd6e5a4cf86332c50204e25779d2c56bb24829 Mon Sep 17 00:00:00 2001 From: Alex Crichton Date: Fri, 30 Oct 2026 12:46:25 -0700 Subject: [PATCH] Add configuration for buffer sizes in WASI This commit adds a configuration knob on `WasiCtxBuilder` to control the default and maximum sizes of buffers used throughout wasmtime-wasi where necessary. These sizes have historically been limited in the past to ensure that guest memory usage is capped but this can cause real-world programs to break such as in WebAssembly/wasi-libc#894 where Python expects a `write` with a large size to succeed all in one go and doesn't internally have a retry loop for remaining bytes. The default value of 64KiB is not changed in this commit, but the `wasmtime` CLI now has a default buffer size of 2MiB to work with more programs by default. --- crates/cli-flags/src/lib.rs | 4 ++ crates/wasi/src/ctx.rs | 20 +++++++ crates/wasi/src/filesystem.rs | 14 ++++- crates/wasi/src/lib.rs | 22 +++----- crates/wasi/src/p2/filesystem.rs | 22 +++++--- crates/wasi/src/p2/host/filesystem.rs | 15 ++--- crates/wasi/src/p2/host/tcp.rs | 4 +- crates/wasi/src/p2/host/udp.rs | 2 +- crates/wasi/src/p2/tcp.rs | 59 +++++++++++++------- crates/wasi/src/p3/filesystem/host.rs | 51 +++++++++++------ crates/wasi/src/p3/sockets/host/types/tcp.rs | 26 ++++++--- crates/wasi/src/sockets/mod.rs | 14 ++++- src/common.rs | 11 ++++ 13 files changed, 181 insertions(+), 83 deletions(-) diff --git a/crates/cli-flags/src/lib.rs b/crates/cli-flags/src/lib.rs index 590ead919660..65cecf0fdc5c 100644 --- a/crates/cli-flags/src/lib.rs +++ b/crates/cli-flags/src/lib.rs @@ -635,6 +635,10 @@ wasmtime_option_group! { pub max_resources: Option, /// Fuel to use for all hostcalls to limit guest<->host data transfer. pub hostcall_fuel: Option, + /// Maximum size, in bytes, of a single WASI filesystem or socket read + /// or write, which bounds the host buffer allocated per I/O operation. + /// (default: 2MiB) + pub max_io_size: Option, /// Maximum number of component-model handle-table slots (e.g. for /// resources, streams, futures) per store; allocating past it traps. /// (default: 100000) diff --git a/crates/wasi/src/ctx.rs b/crates/wasi/src/ctx.rs index a54ed070af52..776330866aa9 100644 --- a/crates/wasi/src/ctx.rs +++ b/crates/wasi/src/ctx.rs @@ -8,6 +8,7 @@ use rand::Rng; use std::future::Future; use std::mem; use std::net::SocketAddr; +use std::num::NonZeroUsize; use std::path::Path; use std::pin::Pin; use tokio::io::{stderr, stdin, stdout}; @@ -163,6 +164,25 @@ impl WasiCtxBuilder { self } + /// Configures the maximum number of bytes that a single read or write of a + /// file or socket may transfer. + /// + /// This bounds the size of the host-side buffer allocated for each + /// individual filesystem or socket I/O operation, which is how much data a + /// guest can cause the host to allocate at once. For example file + /// reads/writes, UDP bytes received, and TCP reads/writes will use this + /// value when limiting a buffering being allocated is appropriate. + /// + /// Larger values may improve throughput of large transfers at the cost of + /// larger per-operation allocations. + /// + /// The default is 64 KiB. + pub fn max_io_size(&mut self, size: NonZeroUsize) -> &mut Self { + self.filesystem.max_io_size = size; + self.sockets.max_io_size = size; + self + } + /// Appends multiple environment variables at once for this builder. /// /// All environment variables are appended to the list of environment diff --git a/crates/wasi/src/filesystem.rs b/crates/wasi/src/filesystem.rs index 7b8052a9a344..f5e2eadca7e5 100644 --- a/crates/wasi/src/filesystem.rs +++ b/crates/wasi/src/filesystem.rs @@ -4,6 +4,7 @@ use crate::runtime::{AbortOnDropJoinHandle, spawn_blocking}; use crate::{NamedId, WasiCtxNamedView}; use std::collections::hash_map; use std::marker; +use std::num::NonZeroUsize; use std::sync::Arc; use std::time::SystemTime; use tracing::debug; @@ -65,10 +66,21 @@ impl HasData for WasiFilesystem { type Data<'a> = WasiFilesystemCtxView<'a>; } -#[derive(Clone, Default)] +#[derive(Clone)] pub struct WasiFilesystemCtx { pub(crate) allow_blocking_current_thread: bool, pub(crate) preopens: Vec<(Dir, String)>, + pub(crate) max_io_size: NonZeroUsize, +} + +impl Default for WasiFilesystemCtx { + fn default() -> Self { + Self { + allow_blocking_current_thread: false, + preopens: Vec::new(), + max_io_size: crate::DEFAULT_MAX_IO_SIZE, + } + } } pub struct WasiFilesystemCtxView<'a> { diff --git a/crates/wasi/src/lib.rs b/crates/wasi/src/lib.rs index 483d97ab15a4..94fbf41dadd6 100644 --- a/crates/wasi/src/lib.rs +++ b/crates/wasi/src/lib.rs @@ -12,23 +12,15 @@ //! //! For WASIp3, see [`p3`]. WASIp3 support is experimental, unstable and incomplete. -/// The maximum size, in bytes, that this crate will allocate on the host on a -/// per-read basis. +/// The default maximum size, in bytes, that this crate will allocate on the +/// host on a per-read basis. /// -/// This is used to limit the size of `wasi:io/streams.input-stream#read`, for -/// example, along with a variety of other read-style APIs. All of these APIs -/// have the shape where the guest asks the host to read some data for it, and -/// then it's returned. The data is allocated on the host, however, meaning that -/// when the guest asks for an extremely large read that could cause a very -/// large allocations on the host. For now this constant serves as a hard limit -/// on the size of the allocation a host will make. -/// -/// TODO: make this configurable per-I/O? Per-WASI? It's not great that this -/// just a hard and unconfigurable limit. There are probably situations where -/// performance will be improved if this limit were higher or removed. Doing so -/// without exposing hosts to guest-controlled resource exhaustion is the -/// important part, though. +/// This is used throughout this crate as the default limiting factor for +/// buffers allocated on the host for I/O reads, writes, etc. This can be +/// overridden through `WasiCtxBuilder` in some situations, too. const MAX_READ_SIZE_ALLOC: usize = 64 * 1024; +const DEFAULT_MAX_IO_SIZE: std::num::NonZeroUsize = + std::num::NonZeroUsize::new(MAX_READ_SIZE_ALLOC).unwrap(); pub mod cli; pub mod clocks; diff --git a/crates/wasi/src/p2/filesystem.rs b/crates/wasi/src/p2/filesystem.rs index 0802b913b179..b5b0ac5a6fc4 100644 --- a/crates/wasi/src/p2/filesystem.rs +++ b/crates/wasi/src/p2/filesystem.rs @@ -67,6 +67,7 @@ pub struct FileInputStream { file: File, position: u64, state: ReadState, + max_read_size: usize, } enum ReadState { Idle, @@ -76,16 +77,17 @@ enum ReadState { Closed, } impl FileInputStream { - pub fn new(file: &File, position: u64) -> Self { + pub fn new(file: &File, position: u64, max_read_size: usize) -> Self { Self { file: file.clone(), position, state: ReadState::Idle, + max_read_size, } } fn blocking_read(file: &std::fs::File, offset: u64, size: usize) -> ReadState { - let mut buf = BytesMut::zeroed(size.min(crate::MAX_READ_SIZE_ALLOC)); + let mut buf = BytesMut::zeroed(size); loop { match sys::read_at_cursor_unspecified(file, &mut buf, offset) { Ok(0) => return ReadState::Closed, @@ -121,6 +123,7 @@ impl InputStream for FileInputStream { } let p = self.position; + let size = size.min(self.max_read_size); self.state = ReadState::Waiting( self.file .spawn_blocking(move |f| Self::blocking_read(f, p, size)), @@ -152,6 +155,7 @@ impl InputStream for FileInputStream { // Before we defer to the regular `read`, make sure it has data ready to go: if let ReadState::Idle = self.state { let p = self.position; + let size = size.min(self.max_read_size); self.state = self .file .run_blocking(move |f| Self::blocking_read(f, p, size)) @@ -208,6 +212,7 @@ pub(crate) struct FileOutputStream { file: File, mode: FileOutputMode, state: OutputState, + max_write_size: usize, } enum OutputState { @@ -221,19 +226,21 @@ enum OutputState { } impl FileOutputStream { - pub fn write_at(file: &File, position: u64) -> Self { + pub fn write_at(file: &File, position: u64, max_write_size: usize) -> Self { Self { file: file.clone(), mode: FileOutputMode::Position(position), state: OutputState::Ready, + max_write_size, } } - pub fn append(file: &File) -> Self { + pub fn append(file: &File, max_write_size: usize) -> Self { Self { file: file.clone(), mode: FileOutputMode::Append, state: OutputState::Ready, + max_write_size, } } @@ -273,14 +280,11 @@ impl FileOutputStream { } } -// FIXME: configurable? determine from how much space left in file? -const FILE_WRITE_CAPACITY: usize = crate::MAX_READ_SIZE_ALLOC; - #[async_trait::async_trait] impl OutputStream for FileOutputStream { fn write(&mut self, buf: Bytes) -> Result<(), StreamError> { match self.state { - OutputState::Ready if buf.len() <= FILE_WRITE_CAPACITY => {} + OutputState::Ready if buf.len() <= self.max_write_size => {} OutputState::Ready => { return Err(StreamError::Trap(format_err!("write exceeded budget"))); } @@ -350,7 +354,7 @@ impl OutputStream for FileOutputStream { } fn check_write(&mut self) -> Result { match self.state { - OutputState::Ready => Ok(FILE_WRITE_CAPACITY), + OutputState::Ready => Ok(self.max_write_size), OutputState::Closed => Err(StreamError::Closed), OutputState::Error(_) => match mem::replace(&mut self.state, OutputState::Closed) { OutputState::Error(e) => Err(StreamError::LastOperationFailed(e.into())), diff --git a/crates/wasi/src/p2/host/filesystem.rs b/crates/wasi/src/p2/host/filesystem.rs index 1988cd107217..4e0b2b25590f 100644 --- a/crates/wasi/src/p2/host/filesystem.rs +++ b/crates/wasi/src/p2/host/filesystem.rs @@ -107,15 +107,11 @@ impl HostDescriptor for WasiFilesystemCtxView<'_> { offset: types::Filesize, ) -> FsResult<(Vec, bool)> { let f = self.table.get(&fd)?.file()?; + let max = self.ctx.max_io_size.get(); let (mut buffer, r) = f .run_blocking(move |f| { - let mut buffer = vec![ - 0; - len.try_into() - .unwrap_or(usize::MAX) - .min(crate::MAX_READ_SIZE_ALLOC) - ]; + let mut buffer = vec![0; len.try_into().unwrap_or(usize::MAX).min(max)]; let r = sys::read_at_cursor_unspecified(f, &mut buffer, offset); (buffer, r) }) @@ -372,7 +368,8 @@ impl HostDescriptor for WasiFilesystemCtxView<'_> { }; // Create a stream view for it. - let reader: DynInputStream = Box::new(FileInputStream::new(f, offset)); + let reader: DynInputStream = + Box::new(FileInputStream::new(f, offset, self.ctx.max_io_size.get())); // Insert the stream view into the table. Trap if the table is full. let index = self.table.push(reader)?; @@ -393,7 +390,7 @@ impl HostDescriptor for WasiFilesystemCtxView<'_> { } // Create a stream view for it. - let writer = FileOutputStream::write_at(f, offset); + let writer = FileOutputStream::write_at(f, offset, self.ctx.max_io_size.get()); let writer: DynOutputStream = Box::new(writer); // Insert the stream view into the table. Trap if the table is full. @@ -414,7 +411,7 @@ impl HostDescriptor for WasiFilesystemCtxView<'_> { } // Create a stream view for it. - let appender = FileOutputStream::append(f); + let appender = FileOutputStream::append(f, self.ctx.max_io_size.get()); let appender: DynOutputStream = Box::new(appender); // Insert the stream view into the table. Trap if the table is full. diff --git a/crates/wasi/src/p2/host/tcp.rs b/crates/wasi/src/p2/host/tcp.rs index 13d6dd1dfb13..052defdbad69 100644 --- a/crates/wasi/src/p2/host/tcp.rs +++ b/crates/wasi/src/p2/host/tcp.rs @@ -89,7 +89,7 @@ impl crate::p2::host::tcp::tcp::HostTcpSocket for WasiSocketsCtxView<'_> { return Err(e.into()); } - let (input, output) = socket.take_streams()?; + let (input, output) = socket.take_streams(self.ctx.max_io_size.get())?; let input = self.table.push_child(input, &this)?; let output = self.table.push_child(output, &this)?; Ok((input, output)) @@ -139,7 +139,7 @@ impl crate::p2::host::tcp::tcp::HostTcpSocket for WasiSocketsCtxView<'_> { None => return Err(ErrorCode::WouldBlock.into()), }; let mut tcp_socket = TcpSocket::new(accepted); - let (input, output) = tcp_socket.take_streams()?; + let (input, output) = tcp_socket.take_streams(self.ctx.max_io_size.get())?; let tcp_socket = self.table.push(tcp_socket)?; let input_stream = self.table.push_child(input, &tcp_socket)?; diff --git a/crates/wasi/src/p2/host/udp.rs b/crates/wasi/src/p2/host/udp.rs index 029d97e56863..d731492595a0 100644 --- a/crates/wasi/src/p2/host/udp.rs +++ b/crates/wasi/src/p2/host/udp.rs @@ -190,7 +190,7 @@ impl udp::HostIncomingDatagramStream for WasiSocketsCtxView<'_> { let mut datagrams = vec![]; let mut sum = 0; - while datagrams.len() < max_results && sum < crate::MAX_READ_SIZE_ALLOC { + while datagrams.len() < max_results && sum < self.ctx.max_io_size.get() { match stream.try_recv() { Err(ErrorCode::WouldBlock) => break, Ok((data, remote_addr)) => { diff --git a/crates/wasi/src/p2/tcp.rs b/crates/wasi/src/p2/tcp.rs index 3589b473ba07..675701b99c66 100644 --- a/crates/wasi/src/p2/tcp.rs +++ b/crates/wasi/src/p2/tcp.rs @@ -1,4 +1,3 @@ -use crate::MAX_READ_SIZE_ALLOC; use crate::p2::bindings::sockets::network::ErrorCode; use crate::p2::{ DynInputStream, DynOutputStream, InputStream, OutputStream, Pollable, SocketResult, StreamError, @@ -42,9 +41,12 @@ impl TcpSocket { writer: None, } } - pub(crate) fn take_streams(&mut self) -> SocketResult<(DynInputStream, DynOutputStream)> { - let reader = TcpReader::new(self.inner.take_receive_stream()?); - let writer = TcpWriter::new(self.inner.take_send_stream()?); + pub(crate) fn take_streams( + &mut self, + max_io_size: usize, + ) -> SocketResult<(DynInputStream, DynOutputStream)> { + let reader = TcpReader::new(self.inner.take_receive_stream()?, max_io_size); + let writer = TcpWriter::new(self.inner.take_send_stream()?, max_io_size); self.reader = Some(reader.clone()); self.writer = Some(writer.clone()); let input: DynInputStream = Box::new(reader); @@ -56,11 +58,11 @@ impl TcpSocket { let writer = self.writer.as_mut().ok_or(ErrorCode::InvalidState)?; if let Shutdown::Both | Shutdown::Read = how { - reader.0.lock().unwrap().shutdown(); + reader.state.lock().unwrap().shutdown(); } if let Shutdown::Both | Shutdown::Write = how { - writer.0.lock().unwrap().shutdown(); + writer.state.lock().unwrap().shutdown(); } Ok(()) @@ -79,7 +81,7 @@ impl ReadState { if size == 0 { return Ok(bytes::Bytes::new()); } - let mut buf = bytes::BytesMut::zeroed(size.min(crate::MAX_READ_SIZE_ALLOC)); + let mut buf = bytes::BytesMut::zeroed(size); let n = match poll_now(|cx| stream.poll_read(cx, &mut buf)) { None => 0, Some(Ok(0)) => { @@ -110,24 +112,33 @@ impl ReadState { } #[derive(Clone)] -struct TcpReader(Arc>); +struct TcpReader { + state: Arc>, + max_read_size: usize, +} impl TcpReader { - fn new(stream: TcpReceiveStream) -> Self { - Self(Arc::new(Mutex::new(ReadState::Open(stream)))) + fn new(stream: TcpReceiveStream, max_read_size: usize) -> Self { + Self { + state: Arc::new(Mutex::new(ReadState::Open(stream))), + max_read_size, + } } } #[async_trait::async_trait] impl InputStream for TcpReader { fn read(&mut self, size: usize) -> StreamResult { - self.0.lock().unwrap().read(size) + self.state + .lock() + .unwrap() + .read(size.min(self.max_read_size)) } } #[async_trait::async_trait] impl Pollable for TcpReader { async fn ready(&mut self) { - std::future::poll_fn(|cx| self.0.lock().unwrap().poll_ready(cx)).await + std::future::poll_fn(|cx| self.state.lock().unwrap().poll_ready(cx)).await } } @@ -158,11 +169,11 @@ impl WriteState { mem::replace(self, WriteState::Closed(WriteError::Closed)) } - fn check_write(&mut self) -> StreamResult { + fn check_write(&mut self, max_write_size: usize) -> StreamResult { match poll_now(|cx| self.poll_ready(cx)) { None => Ok(0), Some(Ok((_, permit))) => { - *permit = MAX_READ_SIZE_ALLOC; + *permit = max_write_size; Ok(*permit) } Some(Err(e)) => Err(e), @@ -284,25 +295,31 @@ impl WriteState { } #[derive(Clone)] -struct TcpWriter(Arc>); +struct TcpWriter { + state: Arc>, + max_write_size: usize, +} impl TcpWriter { - fn new(stream: TcpSendStream) -> Self { - Self(Arc::new(Mutex::new(WriteState::Ready(stream, 0)))) + fn new(stream: TcpSendStream, max_write_size: usize) -> Self { + Self { + state: Arc::new(Mutex::new(WriteState::Ready(stream, 0))), + max_write_size, + } } } #[async_trait::async_trait] impl OutputStream for TcpWriter { fn write(&mut self, bytes: bytes::Bytes) -> StreamResult<()> { - self.0.lock().unwrap().write(bytes) + self.state.lock().unwrap().write(bytes) } fn flush(&mut self) -> StreamResult<()> { - self.0.lock().unwrap().flush() + self.state.lock().unwrap().flush() } fn check_write(&mut self) -> StreamResult { - self.0.lock().unwrap().check_write() + self.state.lock().unwrap().check_write(self.max_write_size) } async fn cancel(&mut self) { @@ -315,6 +332,6 @@ impl OutputStream for TcpWriter { #[async_trait::async_trait] impl Pollable for TcpWriter { async fn ready(&mut self) { - poll_fn(|cx| self.0.lock().unwrap().poll_ready(cx).map(|_| ())).await; + poll_fn(|cx| self.state.lock().unwrap().poll_ready(cx).map(|_| ())).await; } } diff --git a/crates/wasi/src/p3/filesystem/host.rs b/crates/wasi/src/p3/filesystem/host.rs index 3f70bcbfc9fe..469e9bfd71a5 100644 --- a/crates/wasi/src/p3/filesystem/host.rs +++ b/crates/wasi/src/p3/filesystem/host.rs @@ -1,12 +1,12 @@ use crate::filesystem::sys; use crate::filesystem::{Descriptor, Dir, File, WasiFilesystem, WasiFilesystemCtxView}; +use crate::p3::FallibleIteratorProducer; use crate::p3::bindings::clocks::system_clock; use crate::p3::bindings::filesystem::types::{ self, Advice, DescriptorFlags, DescriptorStat, DescriptorType, DirectoryEntry, ErrorCode, Filesize, MetadataHashValue, NewTimestamp, OpenFlags, PathFlags, }; use crate::p3::filesystem::{FilesystemError, FilesystemResult, preopens}; -use crate::p3::{DEFAULT_BUFFER_CAPACITY, FallibleIteratorProducer}; use bytes::BytesMut; use core::pin::Pin; use core::task::{Context, Poll, ready}; @@ -97,6 +97,7 @@ struct ReadStreamProducer { offset: u64, result: Option>>, task: Option>>, + max_read_size: usize, } impl Drop for ReadStreamProducer { @@ -134,14 +135,14 @@ impl StreamProducer for ReadStreamProducer { fn poll_produce<'a>( mut self: Pin<&mut Self>, cx: &mut Context<'_>, - store: StoreContextMut<'a, D>, + mut store: StoreContextMut<'a, D>, mut dst: Destination<'a, Self::Item, Self::Buffer>, finish: bool, ) -> Poll> { if let Some(file) = self.file.as_blocking_file() { // Once a blocking file, always a blocking file, so assert as such. assert!(self.task.is_none()); - let mut dst = dst.as_direct(store, DEFAULT_BUFFER_CAPACITY); + let mut dst = dst.as_direct(store, self.max_read_size); let buf = dst.remaining()?; if buf.is_empty() { return Poll::Ready(Ok(StreamResult::Completed)); @@ -164,18 +165,23 @@ impl StreamProducer for ReadStreamProducer { // Lazily spawn a read task if one hasn't already been spawned yet. let me = &mut *self; - let task = me.task.get_or_insert_with(|| { + if me.task.is_none() { + let size = match dst.remaining(store.as_context_mut())? { + Some(0) | None => me.max_read_size, + Some(n) => n.min(me.max_read_size), + }; let mut buf = dst.take_buffer(); - buf.resize(DEFAULT_BUFFER_CAPACITY, 0); + buf.resize(size, 0); let file = Arc::clone(me.file.as_file()); let offset = me.offset; - spawn_blocking(move || { + me.task = Some(spawn_blocking(move || { sys::read_at_cursor_unspecified(&file, &mut buf, offset).map(|n| { buf.truncate(n); buf }) - }) - }); + })); + } + let task = me.task.as_mut().unwrap(); // Await the completion of the read task. Note that this is not a // cancellable await point because we can't cancel the other task, so @@ -345,6 +351,7 @@ struct WriteStreamConsumer { result: Option>>, buffer: BytesMut, task: Option>>, + max_write_size: usize, } #[derive(Copy, Clone)] @@ -358,6 +365,7 @@ impl WriteStreamConsumer { file: File, location: WriteLocation, result: oneshot::Sender>, + max_write_size: usize, ) -> Self { Self { file, @@ -365,6 +373,7 @@ impl WriteStreamConsumer { result: Some(result), buffer: BytesMut::default(), task: None, + max_write_size, } } @@ -430,7 +439,7 @@ impl StreamConsumer for WriteStreamConsumer { let remaining = src.remaining()?; let task = me.task.get_or_insert_with(|| { debug_assert!(me.buffer.is_empty()); - let n = remaining.len().min(DEFAULT_BUFFER_CAPACITY); + let n = remaining.len().min(me.max_write_size); me.buffer.extend_from_slice(&remaining[..n]); let buf = mem::take(&mut me.buffer); let file = Arc::clone(me.file.as_file()); @@ -487,6 +496,7 @@ fn read_via_stream( mut store: impl AsContextMut, fd: Descriptor, offset: Filesize, + max_read_size: usize, ) -> wasmtime::Result<(StreamReader, FutureReader>)> { let mut store = store.as_context_mut(); let file = match fd { @@ -509,6 +519,7 @@ fn read_via_stream( offset, result: Some(result_tx), task: None, + max_read_size, }, )?, FutureReader::new(&mut store, result_rx)?, @@ -520,6 +531,7 @@ fn write_via_stream( file: FilesystemResult, mut data: StreamReader, location: WriteLocation, + max_write_size: usize, ) -> wasmtime::Result>> { let mut store = store.as_context_mut(); let (result_tx, result_rx) = oneshot::channel(); @@ -527,7 +539,7 @@ fn write_via_stream( Ok(file) => { data.pipe( &mut store, - WriteStreamConsumer::new(file, location, result_tx), + WriteStreamConsumer::new(file, location, result_tx, max_write_size), )?; } Err(err) => { @@ -863,7 +875,8 @@ impl types::HostDescriptorWithStore for WasiFilesystem { offset: Filesize, ) -> wasmtime::Result<(StreamReader, FutureReader>)> { let fd = get_descriptor(store.get().table, &fd)?.clone(); - read_via_stream(&mut store, fd, offset) + let max = store.get().ctx.max_io_size.get(); + read_via_stream(&mut store, fd, offset, max) } fn write_via_stream( @@ -873,7 +886,8 @@ impl types::HostDescriptorWithStore for WasiFilesystem { offset: Filesize, ) -> wasmtime::Result>> { let file = get_writable_file(store.get().table, &fd); - write_via_stream(&mut store, file, data, WriteLocation::Offset(offset)) + let max = store.get().ctx.max_io_size.get(); + write_via_stream(&mut store, file, data, WriteLocation::Offset(offset), max) } fn append_via_stream( @@ -882,7 +896,8 @@ impl types::HostDescriptorWithStore for WasiFilesystem { data: StreamReader, ) -> wasmtime::Result>> { let file = get_writable_file(store.get().table, &fd); - write_via_stream(&mut store, file, data, WriteLocation::End) + let max = store.get().ctx.max_io_size.get(); + write_via_stream(&mut store, file, data, WriteLocation::End, max) } async fn advise( @@ -1159,7 +1174,8 @@ mod named { offset: Filesize, ) -> wasmtime::Result<(StreamReader, FutureReader>)> { let fd = super::get_descriptor(store.get().0.filesystem(id).table, &fd)?.clone(); - super::read_via_stream(&mut store, fd, offset) + let max = store.get().0.filesystem(id).ctx.max_io_size.get(); + super::read_via_stream(&mut store, fd, offset, max) } fn write_via_stream( @@ -1170,7 +1186,9 @@ mod named { offset: Filesize, ) -> wasmtime::Result>> { let file = super::get_writable_file(store.get().0.filesystem(id).table, &fd); - super::write_via_stream(&mut store, file, data, super::WriteLocation::Offset(offset)) + let max = store.get().0.filesystem(id).ctx.max_io_size.get(); + let location = super::WriteLocation::Offset(offset); + super::write_via_stream(&mut store, file, data, location, max) } fn append_via_stream( @@ -1180,7 +1198,8 @@ mod named { data: StreamReader, ) -> wasmtime::Result>> { let file = super::get_writable_file(store.get().0.filesystem(id).table, &fd); - super::write_via_stream(&mut store, file, data, super::WriteLocation::End) + let max = store.get().0.filesystem(id).ctx.max_io_size.get(); + super::write_via_stream(&mut store, file, data, super::WriteLocation::End, max) } async fn advise( diff --git a/crates/wasi/src/p3/sockets/host/types/tcp.rs b/crates/wasi/src/p3/sockets/host/types/tcp.rs index c85373b44170..3777c891b6fe 100644 --- a/crates/wasi/src/p3/sockets/host/types/tcp.rs +++ b/crates/wasi/src/p3/sockets/host/types/tcp.rs @@ -1,4 +1,3 @@ -use crate::p3::DEFAULT_BUFFER_CAPACITY; use crate::p3::bindings::sockets::types::{ Duration, ErrorCode, HostTcpSocket, HostTcpSocketWithStore, IpAddressFamily, IpSocketAddress, TcpSocket, @@ -83,7 +82,10 @@ where } } -struct ReceiveStreamProducer(Option<(TcpReceiveStream, oneshot::Sender>)>); +struct ReceiveStreamProducer { + stream: Option<(TcpReceiveStream, oneshot::Sender>)>, + max_read_size: usize, +} impl Drop for ReceiveStreamProducer { fn drop(&mut self) { @@ -93,7 +95,7 @@ impl Drop for ReceiveStreamProducer { impl ReceiveStreamProducer { fn close(&mut self, res: Result<(), ErrorCode>) { - if let Some((_, tx)) = self.0.take() { + if let Some((_, tx)) = self.stream.take() { _ = tx.send(res); } } @@ -110,7 +112,8 @@ impl StreamProducer for ReceiveStreamProducer { dst: Destination<'a, Self::Item, Self::Buffer>, finish: bool, ) -> Poll> { - let Some((stream, _)) = self.0.as_mut() else { + let max_read_size = self.max_read_size; + let Some((stream, _)) = self.stream.as_mut() else { return Poll::Ready(Ok(StreamResult::Dropped)); }; @@ -123,7 +126,7 @@ impl StreamProducer for ReceiveStreamProducer { }; } - let mut dst = dst.as_direct(store, DEFAULT_BUFFER_CAPACITY); + let mut dst = dst.as_direct(store, max_read_size); let buf = dst.remaining()?; match stream.poll_read(cx, buf) { Poll::Ready(Ok(0)) => { @@ -226,13 +229,18 @@ fn send( fn receive( mut store: impl AsContextMut, stream: Result, + max_read_size: usize, ) -> wasmtime::Result<(StreamReader, FutureReader>)> { let mut store = store.as_context_mut(); match stream { Ok(stream) => { let (result_tx, result_rx) = oneshot::channel(); + let producer = ReceiveStreamProducer { + stream: Some((stream, result_tx)), + max_read_size, + }; Ok(( - StreamReader::new(&mut store, ReceiveStreamProducer(Some((stream, result_tx))))?, + StreamReader::new(&mut store, producer)?, FutureReader::new(&mut store, result_rx)?, )) } @@ -320,7 +328,8 @@ impl HostTcpSocketWithStore for WasiSockets { ) -> wasmtime::Result<(StreamReader, FutureReader>)> { let socket = get_socket_mut(store.get().table, &socket)?; let stream = socket.take_receive_stream(); - receive(&mut store, stream) + let max = store.get().ctx.max_io_size.get(); + receive(&mut store, stream, max) } } @@ -549,7 +558,8 @@ mod named { ) -> wasmtime::Result<(StreamReader, FutureReader>)> { let socket = super::get_socket_mut(store.get().0.sockets(id).table, &socket)?; let stream = socket.take_receive_stream(); - super::receive(&mut store, stream) + let max = store.get().0.sockets(id).ctx.max_io_size.get(); + super::receive(&mut store, stream, max) } } diff --git a/crates/wasi/src/sockets/mod.rs b/crates/wasi/src/sockets/mod.rs index 39f8689491e0..df6fa8b1cc50 100644 --- a/crates/wasi/src/sockets/mod.rs +++ b/crates/wasi/src/sockets/mod.rs @@ -8,6 +8,7 @@ use rustix::fd::AsFd; use rustix::io::Errno; use rustix::net::sockopt; use std::marker; +use std::num::NonZeroUsize; use std::pin::Pin; use std::sync::Arc; use std::task::Poll; @@ -65,10 +66,21 @@ impl HasData for WasiSockets { type Data<'a> = WasiSocketsCtxView<'a>; } -#[derive(Clone, Default)] +#[derive(Clone)] pub struct WasiSocketsCtx { pub(crate) socket_addr_check: SocketAddrCheck, pub(crate) allowed_network_uses: AllowedNetworkUses, + pub(crate) max_io_size: NonZeroUsize, +} + +impl Default for WasiSocketsCtx { + fn default() -> Self { + Self { + socket_addr_check: Default::default(), + allowed_network_uses: Default::default(), + max_io_size: crate::DEFAULT_MAX_IO_SIZE, + } + } } pub struct WasiSocketsCtxView<'a> { diff --git a/src/common.rs b/src/common.rs index b56be9bd95a2..d86dd1415b9f 100644 --- a/src/common.rs +++ b/src/common.rs @@ -2,6 +2,7 @@ use clap::Parser; use std::net::TcpListener; +use std::num::NonZeroUsize; use std::{fs::File, path::Path, time::Duration}; use wasmtime::{ Engine, Module, Precompiled, Result, Store, StoreLimits, StoreLimitsBuilder, bail, @@ -372,6 +373,16 @@ impl RunCommon { if let Some(enable) = self.common.wasi.allow_ip_name_lookup { builder.allow_ip_name_lookup(enable); } + // The default of 64KiB for this value is a bit small for a CLI + // embedding which has more leeway with resources and just a single + // guest, so raise it by default if no configuration is otherwise + // applied. + builder.max_io_size( + self.common + .wasi + .max_io_size + .unwrap_or(NonZeroUsize::new(2 << 20).unwrap()), + ); if let Some(enable) = self.common.wasi.tcp { builder.allow_tcp(enable); }