sausage/vfs/src/streams.rs

//! Nonblocking WASI streams that await layer operations through readiness.
use crate::*;
use bytes::Bytes;
use std::{future::Future, pin::Pin};
type Pending<T> = Pin<Box<dyn Future<Output = Result<T>> + Send>>;
use wasmtime_wasi::p2::{InputStream, OutputStream, Pollable, StreamError, StreamResult};
fn failure(e: io::Error) -> StreamError {
    StreamError::LastOperationFailed(e.into())
}
pub(crate) struct Reader {
    file: Arc<dyn File>,
    offset: u64,
    pending: Option<Pending<Vec<u8>>>,
    data: Bytes,
    eof: bool,
    error: Option<io::Error>,
}
impl Reader {
    pub fn new(file: Arc<dyn File>, offset: u64) -> Self {
        Self {
            file,
            offset,
            pending: None,
            data: Bytes::new(),
            eof: false,
            error: None,
        }
    }
    fn start(&mut self, size: usize) {
        if self.pending.is_none() && self.data.is_empty() && !self.eof {
            let f = self.file.clone();
            let offset = self.offset;
            self.pending = Some(Box::pin(
                async move { f.read(offset, size.min(65536)).await },
            ));
        }
    }
}
#[async_trait]
impl Pollable for Reader {
    async fn ready(&mut self) {
        self.start(4096);
        if let Some(task) = &mut self.pending {
            match task.await {
                Ok(data) => {
                    self.eof = data.is_empty();
                    self.offset += data.len() as u64;
                    self.data = data.into();
                }
                Err(e) => self.error = Some(e),
            }
            self.pending = None;
        }
    }
}
#[async_trait]
impl InputStream for Reader {
    async fn blocking_read(&mut self, size: usize) -> StreamResult<Bytes> {
        if size == 0 {
            return Ok(Bytes::new());
        }
        self.start(size);
        self.ready().await;
        self.read(size)
    }
    fn read(&mut self, size: usize) -> StreamResult<Bytes> {
        if let Some(e) = self.error.take() {
            return Err(failure(e));
        }
        if self.eof && self.data.is_empty() {
            return Err(StreamError::Closed);
        }
        let n = size.min(self.data.len());
        let data = self.data.split_to(n);
        if data.is_empty() && size != 0 {
            self.start(size);
        }
        Ok(data)
    }
}
pub(crate) struct Writer {
    file: Arc<dyn File>,
    offset: Option<u64>,
    pending: Option<Pending<usize>>,
    error: Option<io::Error>,
}
impl Writer {
    pub fn new(file: Arc<dyn File>, offset: Option<u64>) -> Self {
        Self {
            file,
            offset,
            pending: None,
            error: None,
        }
    }
}
#[async_trait]
impl Pollable for Writer {
    async fn ready(&mut self) {
        if let Some(task) = &mut self.pending {
            match task.await {
                Ok(n) => {
                    if let Some(p) = &mut self.offset {
                        *p += n as u64
                    }
                }
                Err(e) => self.error = Some(e),
            }
            self.pending = None;
        }
    }
}
#[async_trait]
impl OutputStream for Writer {
    async fn cancel(&mut self) {
        self.ready().await;
    }
    fn write(&mut self, bytes: Bytes) -> StreamResult<()> {
        if self.check_write()? == 0 {
            return Err(StreamError::trap("write without a permit"));
        }
        let f = self.file.clone();
        let offset = self.offset;
        self.pending = Some(Box::pin(async move {
            let mut done = 0;
            while done < bytes.len() {
                let n = match offset {
                    Some(p) => f.write(p + done as u64, bytes[done..].to_vec()).await?,
                    None => f.append(bytes[done..].to_vec()).await?,
                };
                if n == 0 {
                    return Err(error(io::ErrorKind::WriteZero));
                }
                done += n;
            }
            Ok(done)
        }));
        Ok(())
    }
    fn flush(&mut self) -> StreamResult<()> {
        Ok(())
    }
    fn check_write(&mut self) -> StreamResult<usize> {
        if let Some(e) = self.error.take() {
            return Err(failure(e));
        }
        Ok(if self.pending.is_some() { 0 } else { 65536 })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    /// Both read APIs request only the bytes needed instead of caching later mutations.
    #[tokio::test]
    async fn small_reads() {
        for blocking in [true, false] {
            let memory = Memory::default();
            let file = memory
                .open(
                    "/file",
                    Open {
                        write: true,
                        create: true,
                        ..Open::default()
                    },
                )
                .await
                .unwrap();
            file.write(0, b"abcd".to_vec()).await.unwrap();
            let mut reader = Reader::new(file.clone(), 0);
            let first = if blocking {
                reader.blocking_read(1).await.unwrap()
            } else {
                assert!(reader.read(1).unwrap().is_empty());
                reader.ready().await;
                reader.read(1).unwrap()
            };
            assert_eq!(first.as_ref(), b"a");
            file.write(1, b"Z".to_vec()).await.unwrap();
            assert_eq!(reader.blocking_read(1).await.unwrap().as_ref(), b"Z");
        }
    }
}