wasmtime_wasi/p2/
pipe.rs

1//! Virtual pipes.
2//!
3//! These types provide easy implementations of `WasiFile` that mimic much of the behavior of Unix
4//! pipes. These are particularly helpful for redirecting WASI stdio handles to destinations other
5//! than OS files.
6//!
7//! Some convenience constructors are included for common backing types like `Vec<u8>` and `String`,
8//! but the virtual pipes can be instantiated with any `Read` or `Write` type.
9//!
10use bytes::Bytes;
11use std::pin::{Pin, pin};
12use std::sync::{Arc, Mutex};
13use std::task::{Context, Poll};
14use tokio::io::{self, AsyncRead, AsyncWrite};
15use tokio::sync::mpsc;
16use wasmtime::format_err;
17use wasmtime_wasi_io::{
18    poll::Pollable,
19    streams::{InputStream, OutputStream, StreamError},
20};
21
22pub use crate::p2::write_stream::AsyncWriteStream;
23
24#[derive(Debug, Clone)]
25pub struct MemoryInputPipe {
26    buffer: Arc<Mutex<Bytes>>,
27}
28
29impl MemoryInputPipe {
30    pub fn new(bytes: impl Into<Bytes>) -> Self {
31        Self {
32            buffer: Arc::new(Mutex::new(bytes.into())),
33        }
34    }
35
36    pub fn is_empty(&self) -> bool {
37        self.buffer.lock().unwrap().is_empty()
38    }
39}
40
41#[async_trait::async_trait]
42impl InputStream for MemoryInputPipe {
43    fn read(&mut self, size: usize) -> Result<Bytes, StreamError> {
44        let mut buffer = self.buffer.lock().unwrap();
45        if buffer.is_empty() {
46            return Err(StreamError::Closed);
47        }
48
49        let size = size.min(buffer.len());
50        let read = buffer.split_to(size);
51        Ok(read)
52    }
53}
54
55#[async_trait::async_trait]
56impl Pollable for MemoryInputPipe {
57    async fn ready(&mut self) {}
58}
59
60impl AsyncRead for MemoryInputPipe {
61    fn poll_read(
62        self: Pin<&mut Self>,
63        _cx: &mut Context<'_>,
64        buf: &mut io::ReadBuf<'_>,
65    ) -> Poll<io::Result<()>> {
66        let mut buffer = self.buffer.lock().unwrap();
67        let size = buf.remaining().min(buffer.len());
68        let read = buffer.split_to(size);
69        buf.put_slice(&read);
70        Poll::Ready(Ok(()))
71    }
72}
73
74#[derive(Debug, Clone)]
75pub struct MemoryOutputPipe {
76    capacity: usize,
77    buffer: Arc<Mutex<bytes::BytesMut>>,
78}
79
80impl MemoryOutputPipe {
81    pub fn new(capacity: usize) -> Self {
82        MemoryOutputPipe {
83            capacity,
84            buffer: std::sync::Arc::new(std::sync::Mutex::new(bytes::BytesMut::new())),
85        }
86    }
87
88    pub fn contents(&self) -> bytes::Bytes {
89        self.buffer.lock().unwrap().clone().freeze()
90    }
91
92    pub fn try_into_inner(self) -> Option<bytes::BytesMut> {
93        std::sync::Arc::into_inner(self.buffer).map(|m| m.into_inner().unwrap())
94    }
95}
96
97#[async_trait::async_trait]
98impl OutputStream for MemoryOutputPipe {
99    fn write(&mut self, bytes: Bytes) -> Result<(), StreamError> {
100        let mut buf = self.buffer.lock().unwrap();
101        if bytes.len() > self.capacity - buf.len() {
102            return Err(StreamError::Trap(format_err!(
103                "write beyond capacity of MemoryOutputPipe"
104            )));
105        }
106        buf.extend_from_slice(bytes.as_ref());
107        // Always ready for writing
108        Ok(())
109    }
110    fn flush(&mut self) -> Result<(), StreamError> {
111        // This stream is always flushed
112        Ok(())
113    }
114    fn check_write(&mut self) -> Result<usize, StreamError> {
115        let consumed = self.buffer.lock().unwrap().len();
116        if consumed < self.capacity {
117            Ok(self.capacity - consumed)
118        } else {
119            // Since the buffer is full, no more bytes will ever be written
120            Err(StreamError::Closed)
121        }
122    }
123}
124
125#[async_trait::async_trait]
126impl Pollable for MemoryOutputPipe {
127    async fn ready(&mut self) {}
128}
129
130impl AsyncWrite for MemoryOutputPipe {
131    fn poll_write(
132        self: Pin<&mut Self>,
133        _cx: &mut Context<'_>,
134        buf: &[u8],
135    ) -> Poll<io::Result<usize>> {
136        let mut buffer = self.buffer.lock().unwrap();
137        let amt = buf.len().min(self.capacity - buffer.len());
138        buffer.extend_from_slice(&buf[..amt]);
139        Poll::Ready(Ok(amt))
140    }
141    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
142        Poll::Ready(Ok(()))
143    }
144    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
145        Poll::Ready(Ok(()))
146    }
147}
148
149/// Provides a [`InputStream`] impl from a [`tokio::io::AsyncRead`] impl
150pub struct AsyncReadStream {
151    closed: bool,
152    buffer: Option<Result<Bytes, StreamError>>,
153    receiver: mpsc::Receiver<Result<Bytes, StreamError>>,
154    join_handle: Option<crate::runtime::AbortOnDropJoinHandle<()>>,
155}
156
157impl AsyncReadStream {
158    /// Create a [`AsyncReadStream`]. In order to use the [`InputStream`] impl
159    /// provided by this struct, the argument must impl [`tokio::io::AsyncRead`].
160    pub fn new<T: AsyncRead + Send + 'static>(reader: T) -> Self {
161        let (sender, receiver) = mpsc::channel(1);
162        let join_handle = crate::runtime::spawn(async move {
163            let mut reader = pin!(reader);
164            loop {
165                use tokio::io::AsyncReadExt;
166                let mut buf = bytes::BytesMut::with_capacity(crate::MAX_READ_SIZE_ALLOC);
167                let sent = match reader.read_buf(&mut buf).await {
168                    Ok(nbytes) if nbytes == 0 => {
169                        let _ = sender.send(Err(StreamError::Closed)).await;
170                        break;
171                    }
172                    Ok(_) => sender.send(Ok(buf.freeze())).await,
173                    Err(e) => {
174                        sender
175                            .send(Err(StreamError::LastOperationFailed(e.into())))
176                            .await
177                    }
178                };
179                if sent.is_err() {
180                    // no more receiver - stop trying to read
181                    break;
182                }
183            }
184        });
185        AsyncReadStream {
186            closed: false,
187            buffer: None,
188            receiver,
189            join_handle: Some(join_handle),
190        }
191    }
192    pub(crate) fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<()> {
193        if self.buffer.is_some() || self.closed {
194            return Poll::Ready(());
195        }
196        match self.receiver.poll_recv(cx) {
197            Poll::Ready(Some(res)) => {
198                self.buffer = Some(res);
199                Poll::Ready(())
200            }
201            Poll::Ready(None) => {
202                panic!("no more sender for an open AsyncReadStream - should be impossible")
203            }
204            Poll::Pending => Poll::Pending,
205        }
206    }
207}
208
209#[async_trait::async_trait]
210impl InputStream for AsyncReadStream {
211    fn read(&mut self, size: usize) -> Result<Bytes, StreamError> {
212        use mpsc::error::TryRecvError;
213
214        match self.buffer.take() {
215            Some(Ok(mut bytes)) => {
216                // TODO: de-duplicate the buffer management with the case below
217                let len = bytes.len().min(size);
218                let rest = bytes.split_off(len);
219                if !rest.is_empty() {
220                    self.buffer = Some(Ok(rest));
221                }
222                return Ok(bytes);
223            }
224            Some(Err(e)) => {
225                self.closed = true;
226                return Err(e);
227            }
228            None => {}
229        }
230
231        match self.receiver.try_recv() {
232            Ok(Ok(mut bytes)) => {
233                let len = bytes.len().min(size);
234                let rest = bytes.split_off(len);
235                if !rest.is_empty() {
236                    self.buffer = Some(Ok(rest));
237                }
238
239                Ok(bytes)
240            }
241            Ok(Err(e)) => {
242                self.closed = true;
243                Err(e)
244            }
245            // Note: if the stream is already closed it should return an error,
246            //       returning empty list would break the wasi contract (returning 0 and ready)
247            Err(TryRecvError::Empty | TryRecvError::Disconnected) if self.closed => {
248                Err(StreamError::Closed)
249            }
250            Err(TryRecvError::Empty) => Ok(Bytes::new()),
251            Err(TryRecvError::Disconnected) => Err(StreamError::Trap(format_err!(
252                "AsyncReadStream sender died - should be impossible"
253            ))),
254        }
255    }
256
257    async fn cancel(&mut self) {
258        match self.join_handle.take() {
259            Some(task) => _ = task.cancel().await,
260            None => {}
261        }
262    }
263}
264
265#[async_trait::async_trait]
266impl Pollable for AsyncReadStream {
267    async fn ready(&mut self) {
268        std::future::poll_fn(|cx| self.poll_ready(cx)).await
269    }
270}
271
272/// An output stream that consumes all input written to it, and is always ready.
273#[derive(Copy, Clone)]
274pub struct SinkOutputStream;
275
276#[async_trait::async_trait]
277impl OutputStream for SinkOutputStream {
278    fn write(&mut self, buf: Bytes) -> Result<(), StreamError> {
279        if buf.len() > crate::MAX_READ_SIZE_ALLOC {
280            return Err(StreamError::Trap(format_err!("write exceeded budget")));
281        }
282        Ok(())
283    }
284    fn flush(&mut self) -> Result<(), StreamError> {
285        // This stream is always flushed
286        Ok(())
287    }
288
289    fn check_write(&mut self) -> Result<usize, StreamError> {
290        // This stream is always ready for writing.
291        Ok(crate::MAX_READ_SIZE_ALLOC)
292    }
293}
294
295#[async_trait::async_trait]
296impl Pollable for SinkOutputStream {
297    async fn ready(&mut self) {}
298}
299
300/// A stream that is ready immediately, but will always report that it's closed.
301#[derive(Copy, Clone)]
302pub struct ClosedInputStream;
303
304#[async_trait::async_trait]
305impl InputStream for ClosedInputStream {
306    fn read(&mut self, _size: usize) -> Result<Bytes, StreamError> {
307        Err(StreamError::Closed)
308    }
309}
310
311#[async_trait::async_trait]
312impl Pollable for ClosedInputStream {
313    async fn ready(&mut self) {}
314}
315
316/// An output stream that is always closed.
317#[derive(Copy, Clone)]
318pub struct ClosedOutputStream;
319
320#[async_trait::async_trait]
321impl OutputStream for ClosedOutputStream {
322    fn write(&mut self, _: Bytes) -> Result<(), StreamError> {
323        Err(StreamError::Closed)
324    }
325    fn flush(&mut self) -> Result<(), StreamError> {
326        Err(StreamError::Closed)
327    }
328
329    fn check_write(&mut self) -> Result<usize, StreamError> {
330        Err(StreamError::Closed)
331    }
332}
333
334#[async_trait::async_trait]
335impl Pollable for ClosedOutputStream {
336    async fn ready(&mut self) {}
337}
338
339#[cfg(test)]
340mod test {
341    use super::*;
342    use std::time::Duration;
343    use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
344
345    // This is a gross way to handle CI running under qemu for non-x86 architectures.
346    #[cfg(not(target_arch = "x86_64"))]
347    const TEST_ITERATIONS: usize = 10;
348
349    #[cfg(target_arch = "x86_64")]
350    const TEST_ITERATIONS: usize = 100;
351
352    async fn resolves_immediately<F, O>(fut: F) -> O
353    where
354        F: futures::Future<Output = O>,
355    {
356        // The input `fut` should resolve immediately, but in case it
357        // accidentally doesn't don't hang the test indefinitely. Provide a
358        // generous timeout to account for CI sensitivity and various systems.
359        tokio::time::timeout(Duration::from_secs(2), fut)
360            .await
361            .expect("operation timed out")
362    }
363
364    async fn never_resolves<F: futures::Future>(fut: F) {
365        // The input `fut` should never resolve, so only give it a small window
366        // of budget before we time out. If `fut` is actually resolved this
367        // should show up as a flaky test.
368        tokio::time::timeout(Duration::from_millis(10), fut)
369            .await
370            .err()
371            .expect("operation should time out");
372    }
373
374    pub fn simplex(size: usize) -> (impl AsyncRead, impl AsyncWrite) {
375        let (a, b) = tokio::io::duplex(size);
376        let (_read_half, write_half) = tokio::io::split(a);
377        let (read_half, _write_half) = tokio::io::split(b);
378        (read_half, write_half)
379    }
380
381    #[test_log::test(tokio::test(flavor = "multi_thread"))]
382    async fn empty_read_stream() {
383        let mut reader = AsyncReadStream::new(tokio::io::empty());
384
385        // In a multi-threaded context, the value of state is not deterministic -- the spawned
386        // reader task may run on a different thread.
387        match reader.read(10) {
388            // The reader task ran before we tried to read, and noticed that the input was empty.
389            Err(StreamError::Closed) => {}
390
391            // The reader task hasn't run yet. Call `ready` to await and fill the buffer.
392            Ok(bs) => {
393                assert!(bs.is_empty());
394                resolves_immediately(reader.ready()).await;
395                assert!(matches!(reader.read(0), Err(StreamError::Closed)));
396
397                // Try again to make sure it keeps returning `StreamError::Closed`
398                resolves_immediately(reader.ready()).await;
399                assert!(matches!(reader.read(0), Err(StreamError::Closed)));
400                assert!(matches!(reader.read(0), Err(StreamError::Closed)));
401            }
402            res => panic!("unexpected: {res:?}"),
403        }
404    }
405
406    #[test_log::test(tokio::test(flavor = "multi_thread"))]
407    async fn infinite_read_stream() {
408        let mut reader = AsyncReadStream::new(tokio::io::repeat(0));
409
410        let bs = reader.read(10).unwrap();
411        if bs.is_empty() {
412            // Reader task hasn't run yet. Call `ready` to await and fill the buffer.
413            resolves_immediately(reader.ready()).await;
414            // Now a read should succeed
415            let bs = reader.read(10).unwrap();
416            assert_eq!(bs.len(), 10);
417        } else {
418            assert_eq!(bs.len(), 10);
419        }
420
421        // Subsequent reads should succeed
422        let bs = reader.read(10).unwrap();
423        assert_eq!(bs.len(), 10);
424
425        // Even 0-length reads should succeed and show its open
426        let bs = reader.read(0).unwrap();
427        assert_eq!(bs.len(), 0);
428    }
429
430    async fn finite_async_reader(contents: &[u8]) -> impl AsyncRead + Send + 'static + use<> {
431        let (r, mut w) = simplex(contents.len());
432        w.write_all(contents).await.unwrap();
433        r
434    }
435
436    #[test_log::test(tokio::test(flavor = "multi_thread"))]
437    async fn finite_read_stream() {
438        let mut reader = AsyncReadStream::new(finite_async_reader(&[1; 123]).await);
439
440        let bs = reader.read(123).unwrap();
441        if bs.is_empty() {
442            // Reader task hasn't run yet. Call `ready` to await and fill the buffer.
443            resolves_immediately(reader.ready()).await;
444            // Now a read should succeed
445            let bs = reader.read(123).unwrap();
446            assert_eq!(bs.len(), 123);
447        } else {
448            assert_eq!(bs.len(), 123);
449        }
450
451        // The AsyncRead's should be empty now, but we have a race where the reader task hasn't
452        // yet send that to the AsyncReadStream.
453        match reader.read(0) {
454            Err(StreamError::Closed) => {} // Correct!
455            Ok(bs) => {
456                assert!(bs.is_empty());
457                // Need to await to give this side time to catch up
458                resolves_immediately(reader.ready()).await;
459                // Now a read should show closed
460                assert!(matches!(reader.read(0), Err(StreamError::Closed)));
461            }
462            res => panic!("unexpected: {res:?}"),
463        }
464
465        // Make sure it stays closed
466        resolves_immediately(reader.ready()).await;
467        assert!(matches!(reader.read(0), Err(StreamError::Closed)));
468    }
469
470    #[test_log::test(tokio::test(flavor = "multi_thread"))]
471    // Test that you can write items into the stream, and they get read out in the order they were
472    // written, with the proper indications of readiness for reading:
473    async fn multiple_chunks_read_stream() {
474        let (r, mut w) = simplex(1024);
475        let mut reader = AsyncReadStream::new(r);
476
477        w.write_all(&[123]).await.unwrap();
478
479        let bs = reader.read(1).unwrap();
480        if bs.is_empty() {
481            // Reader task hasn't run yet. Call `ready` to await and fill the buffer.
482            resolves_immediately(reader.ready()).await;
483            // Now a read should succeed
484            let bs = reader.read(1).unwrap();
485            assert_eq!(*bs, [123u8]);
486        } else {
487            assert_eq!(*bs, [123u8]);
488        }
489
490        // The stream should be empty and open now:
491        let bs = reader.read(1).unwrap();
492        assert!(bs.is_empty());
493
494        // We can wait on readiness and it will time out:
495        never_resolves(reader.ready()).await;
496
497        // Still open and empty:
498        let bs = reader.read(1).unwrap();
499        assert!(bs.is_empty());
500
501        // Put something else in the stream:
502        w.write_all(&[45]).await.unwrap();
503
504        // Wait readiness (yes we could possibly win the race and read it out faster, leaving that
505        // out of the test for simplicity)
506        resolves_immediately(reader.ready()).await;
507
508        // read the something else back out:
509        let bs = reader.read(1).unwrap();
510        assert_eq!(*bs, [45u8]);
511
512        // nothing else in there:
513        let bs = reader.read(1).unwrap();
514        assert!(bs.is_empty());
515
516        // We can wait on readiness and it will time out:
517        never_resolves(reader.ready()).await;
518
519        // nothing else in there:
520        let bs = reader.read(1).unwrap();
521        assert!(bs.is_empty());
522
523        // Now close the pipe:
524        drop(w);
525
526        // Wait readiness (yes we could possibly win the race and read it out faster, leaving that
527        // out of the test for simplicity)
528        resolves_immediately(reader.ready()).await;
529
530        // empty and now closed:
531        assert!(matches!(reader.read(1), Err(StreamError::Closed)));
532    }
533
534    #[test_log::test(tokio::test(flavor = "multi_thread"))]
535    // At the moment we are restricting AsyncReadStream from buffering more than 4k. This isn't a
536    // suitable design for all applications, and we will probably make a knob or change the
537    // behavior at some point, but this test shows the behavior as it is implemented:
538    async fn backpressure_read_stream() {
539        let (r, mut w) = simplex(4 * crate::MAX_READ_SIZE_ALLOC); // Make sure this buffer isn't a bottleneck
540        let mut reader = AsyncReadStream::new(r);
541
542        let writer_task = tokio::task::spawn(async move {
543            // Write twice as much as we can buffer up in an AsyncReadStream:
544            w.write_all(&[123; 2 * crate::MAX_READ_SIZE_ALLOC])
545                .await
546                .unwrap();
547            w
548        });
549
550        resolves_immediately(reader.ready()).await;
551
552        // Now we expect the reader task has sent 4k from the stream to the reader.
553        // Try to read out one bigger than the buffer available:
554        let bs = reader.read(crate::MAX_READ_SIZE_ALLOC + 1).unwrap();
555        assert_eq!(bs.len(), crate::MAX_READ_SIZE_ALLOC);
556
557        // Allow the crank to turn more:
558        resolves_immediately(reader.ready()).await;
559
560        // Again we expect the reader task has sent 4k from the stream to the reader.
561        // Try to read out one bigger than the buffer available:
562        let bs = reader.read(crate::MAX_READ_SIZE_ALLOC + 1).unwrap();
563        assert_eq!(bs.len(), crate::MAX_READ_SIZE_ALLOC);
564
565        // The writer task is now finished - join with it:
566        let w = resolves_immediately(writer_task).await;
567
568        // And close the pipe:
569        drop(w);
570
571        // Allow the crank to turn more:
572        resolves_immediately(reader.ready()).await;
573
574        // Now we expect the reader to be empty, and the stream.dropd:
575        assert!(matches!(
576            reader.read(crate::MAX_READ_SIZE_ALLOC + 1),
577            Err(StreamError::Closed)
578        ));
579    }
580
581    #[test_log::test(test_log::test(tokio::test(flavor = "multi_thread")))]
582    async fn sink_write_stream() {
583        let mut writer = AsyncWriteStream::new(2048, tokio::io::sink());
584        let chunk = Bytes::from_static(&[0; 1024]);
585
586        let readiness = resolves_immediately(writer.write_ready())
587            .await
588            .expect("write_ready does not trap");
589        assert_eq!(readiness, 2048);
590        // I can write whatever:
591        writer.write(chunk.clone()).expect("write does not error");
592
593        // This may consume 1k of the buffer:
594        let readiness = resolves_immediately(writer.write_ready())
595            .await
596            .expect("write_ready does not trap");
597        assert!(
598            readiness == 1024 || readiness == 2048,
599            "readiness should be 1024 or 2048, got {readiness}"
600        );
601
602        if readiness == 1024 {
603            writer.write(chunk.clone()).expect("write does not error");
604
605            let readiness = resolves_immediately(writer.write_ready())
606                .await
607                .expect("write_ready does not trap");
608            assert!(
609                readiness == 1024 || readiness == 2048,
610                "readiness should be 1024 or 2048, got {readiness}"
611            );
612        }
613    }
614
615    #[test_log::test(tokio::test(flavor = "multi_thread"))]
616    async fn closed_write_stream() {
617        // Run many times because the test is nondeterministic:
618        for n in 0..TEST_ITERATIONS {
619            closed_write_stream_(n).await
620        }
621    }
622    #[tracing::instrument]
623    async fn closed_write_stream_(n: usize) {
624        let (reader, writer) = simplex(1);
625        let mut writer = AsyncWriteStream::new(1024, writer);
626
627        // Drop the reader to allow the worker to transition to the closed state eventually.
628        drop(reader);
629
630        // First the api is going to report the last operation failed, then subsequently
631        // it will be reported as closed. We set this flag once we see LastOperationFailed.
632        let mut should_be_closed = false;
633
634        // Write some data to the stream to ensure we have data that cannot be flushed.
635        let chunk = Bytes::from_static(&[0; 1]);
636        writer
637            .write(chunk.clone())
638            .expect("first write should succeed");
639
640        // The rest of this test should be valid whether or not we check write readiness:
641        let mut write_ready_res = None;
642        if n % 2 == 0 {
643            let r = resolves_immediately(writer.write_ready()).await;
644            // Check write readiness:
645            match r {
646                // worker hasn't processed write yet:
647                Ok(1023) => {}
648                // worker reports failure:
649                Err(StreamError::LastOperationFailed(_)) => {
650                    tracing::debug!("discovered stream failure in first write_ready");
651                    should_be_closed = true;
652                }
653                r => panic!("unexpected write_ready: {r:?}"),
654            }
655            write_ready_res = Some(r);
656        }
657
658        // When we drop the simplex reader, that causes the simplex writer to return BrokenPipe on
659        // its write. Now that the buffering crank has turned, our next write will give BrokenPipe.
660        let flush_res = writer.flush();
661        match flush_res {
662            // worker reports failure:
663            Err(StreamError::LastOperationFailed(_)) => {
664                tracing::debug!("discovered stream failure trying to flush");
665                assert!(!should_be_closed);
666                should_be_closed = true;
667            }
668            // Already reported failure, now closed
669            Err(StreamError::Closed) => {
670                assert!(
671                    should_be_closed,
672                    "expected a LastOperationFailed before we see Closed. {write_ready_res:?}"
673                );
674            }
675            // Also possible the worker hasn't processed write yet:
676            Ok(()) => {}
677            Err(e) => panic!("unexpected flush error: {e:?} {write_ready_res:?}"),
678        }
679
680        // Waiting for the flush to complete should always indicate that the channel has been
681        // closed.
682        match resolves_immediately(writer.write_ready()).await {
683            // worker reports failure:
684            Err(StreamError::LastOperationFailed(_)) => {
685                tracing::debug!("discovered stream failure trying to flush");
686                assert!(!should_be_closed);
687            }
688            // Already reported failure, now closed
689            Err(StreamError::Closed) => {
690                assert!(should_be_closed);
691            }
692            r => {
693                panic!(
694                    "stream should be reported closed by the end of write_ready after flush, got {r:?}. {write_ready_res:?} {flush_res:?}"
695                )
696            }
697        }
698    }
699
700    #[test_log::test(tokio::test(flavor = "multi_thread"))]
701    async fn multiple_chunks_write_stream() {
702        // Run many times because the test is nondeterministic:
703        for n in 0..TEST_ITERATIONS {
704            multiple_chunks_write_stream_aux(n).await
705        }
706    }
707    #[tracing::instrument]
708    async fn multiple_chunks_write_stream_aux(_: usize) {
709        use std::ops::Deref;
710
711        let (mut reader, writer) = simplex(1024);
712        let mut writer = AsyncWriteStream::new(1024, writer);
713
714        // Write a chunk:
715        let chunk = Bytes::from_static(&[123; 1]);
716
717        let permit = resolves_immediately(writer.write_ready())
718            .await
719            .expect("write should be ready");
720        assert_eq!(permit, 1024);
721
722        writer.write(chunk.clone()).expect("write does not trap");
723
724        // At this point the message will either be waiting for the worker to process the write, or
725        // it will be buffered in the simplex channel.
726        let permit = resolves_immediately(writer.write_ready())
727            .await
728            .expect("write should be ready");
729        assert!(matches!(permit, 1023 | 1024));
730
731        let mut read_buf = vec![0; chunk.len()];
732        let read_len = reader.read_exact(&mut read_buf).await.unwrap();
733        assert_eq!(read_len, chunk.len());
734        assert_eq!(read_buf.as_slice(), chunk.deref());
735
736        // Write a second, different chunk:
737        let chunk2 = Bytes::from_static(&[45; 1]);
738
739        // We're only guaranteed to see a consistent write budget if we flush.
740        writer.flush().expect("channel is still alive");
741
742        let permit = resolves_immediately(writer.write_ready())
743            .await
744            .expect("write should be ready");
745        assert_eq!(permit, 1024);
746
747        writer.write(chunk2.clone()).expect("write does not trap");
748
749        // At this point the message will either be waiting for the worker to process the write, or
750        // it will be buffered in the simplex channel.
751        let permit = resolves_immediately(writer.write_ready())
752            .await
753            .expect("write should be ready");
754        assert!(matches!(permit, 1023 | 1024));
755
756        let mut read2_buf = vec![0; chunk2.len()];
757        let read2_len = reader.read_exact(&mut read2_buf).await.unwrap();
758        assert_eq!(read2_len, chunk2.len());
759        assert_eq!(read2_buf.as_slice(), chunk2.deref());
760
761        // We're only guaranteed to see a consistent write budget if we flush.
762        writer.flush().expect("channel is still alive");
763
764        let permit = resolves_immediately(writer.write_ready())
765            .await
766            .expect("write should be ready");
767        assert_eq!(permit, 1024);
768    }
769
770    #[test_log::test(tokio::test(flavor = "multi_thread"))]
771    async fn backpressure_write_stream() {
772        // Run many times because the test is nondeterministic:
773        for n in 0..TEST_ITERATIONS {
774            backpressure_write_stream_aux(n).await
775        }
776    }
777    #[tracing::instrument]
778    async fn backpressure_write_stream_aux(_: usize) {
779        use futures::future::poll_immediate;
780
781        // The channel can buffer up to 1k, plus another 1k in the stream, before not
782        // accepting more input:
783        let (mut reader, writer) = simplex(1024);
784        let mut writer = AsyncWriteStream::new(1024, writer);
785
786        let chunk = Bytes::from_static(&[0; 1024]);
787
788        let permit = resolves_immediately(writer.write_ready())
789            .await
790            .expect("write should be ready");
791        assert_eq!(permit, 1024);
792
793        writer.write(chunk.clone()).expect("write succeeds");
794
795        // We might still be waiting for the worker to process the message, or the worker may have
796        // processed it and released all the budget back to us.
797        let permit = poll_immediate(writer.write_ready()).await;
798        assert!(matches!(permit, None | Some(Ok(1024))));
799
800        // Given a little time, the worker will process the message and release all the budget
801        // back.
802        let permit = resolves_immediately(writer.write_ready())
803            .await
804            .expect("write should be ready");
805        assert_eq!(permit, 1024);
806
807        // Now fill the buffer between here and the writer task. This should always indicate
808        // back-pressure because now both buffers (simplex and worker) are full.
809        writer.write(chunk.clone()).expect("write does not trap");
810
811        // Try shoving even more down there, and it shouldn't accept more input:
812        writer
813            .write(chunk.clone())
814            .err()
815            .expect("unpermitted write does trap");
816
817        // No amount of waiting will resolve the situation, as nothing is emptying the simplex
818        // buffer.
819        never_resolves(writer.write_ready()).await;
820
821        // There is 2k buffered between the simplex and worker buffers. I should be able to read
822        // all of it out:
823        let mut buf = [0; 2048];
824        reader.read_exact(&mut buf).await.unwrap();
825
826        // and no more:
827        never_resolves(reader.read(&mut buf)).await;
828
829        // Now the backpressure should be cleared, and an additional write should be accepted.
830        let permit = resolves_immediately(writer.write_ready())
831            .await
832            .expect("ready is ok");
833        assert_eq!(permit, 1024);
834
835        // and the write succeeds:
836        writer.write(chunk.clone()).expect("write does not trap");
837    }
838
839    #[test_log::test(tokio::test(flavor = "multi_thread"))]
840    async fn backpressure_write_stream_with_flush() {
841        for n in 0..TEST_ITERATIONS {
842            backpressure_write_stream_with_flush_aux(n).await;
843        }
844    }
845
846    async fn backpressure_write_stream_with_flush_aux(_: usize) {
847        // The channel can buffer up to 1k, plus another 1k in the stream, before not
848        // accepting more input:
849        let (mut reader, writer) = simplex(1024);
850        let mut writer = AsyncWriteStream::new(1024, writer);
851
852        let chunk = Bytes::from_static(&[0; 1024]);
853
854        let permit = resolves_immediately(writer.write_ready())
855            .await
856            .expect("write should be ready");
857        assert_eq!(permit, 1024);
858
859        writer.write(chunk.clone()).expect("write succeeds");
860
861        writer.flush().expect("flush succeeds");
862
863        // Waiting for write_ready to resolve after a flush should always show that we have the
864        // full budget available, as the message will have flushed to the simplex channel.
865        let permit = resolves_immediately(writer.write_ready())
866            .await
867            .expect("write_ready succeeds");
868        assert_eq!(permit, 1024);
869
870        // Write enough to fill the simplex buffer:
871        writer.write(chunk.clone()).expect("write does not trap");
872
873        // Writes should be refused until this flush succeeds.
874        writer.flush().expect("flush succeeds");
875
876        // Try shoving even more down there, and it shouldn't accept more input:
877        writer
878            .write(chunk.clone())
879            .err()
880            .expect("unpermitted write does trap");
881
882        // No amount of waiting will resolve the situation, as nothing is emptying the simplex
883        // buffer.
884        never_resolves(writer.write_ready()).await;
885
886        // There is 2k buffered between the simplex and worker buffers. I should be able to read
887        // all of it out:
888        let mut buf = [0; 2048];
889        reader.read_exact(&mut buf).await.unwrap();
890
891        // and no more:
892        never_resolves(reader.read(&mut buf)).await;
893
894        // Now the backpressure should be cleared, and an additional write should be accepted.
895        let permit = resolves_immediately(writer.write_ready())
896            .await
897            .expect("ready is ok");
898        assert_eq!(permit, 1024);
899
900        // and the write succeeds:
901        writer.write(chunk.clone()).expect("write does not trap");
902
903        writer.flush().expect("flush succeeds");
904
905        let permit = resolves_immediately(writer.write_ready())
906            .await
907            .expect("ready is ok");
908        assert_eq!(permit, 1024);
909    }
910}
911
912#[cfg(test)]
913mod write_budget_tests {
914    use super::*;
915
916    /// A write larger than the budget `check_write` reports has to trap instead
917    /// of being silently discarded.
918    #[test]
919    fn sink_write_larger_than_budget_traps() {
920        let mut stream = SinkOutputStream;
921        assert_eq!(stream.check_write().unwrap(), crate::MAX_READ_SIZE_ALLOC);
922        stream
923            .write(Bytes::from(vec![0; crate::MAX_READ_SIZE_ALLOC]))
924            .unwrap();
925        let err = stream
926            .write(Bytes::from(vec![0; crate::MAX_READ_SIZE_ALLOC + 1]))
927            .unwrap_err();
928        assert!(matches!(err, StreamError::Trap(_)), "{err:?}");
929    }
930}
wasmtime_wasi/p2/pipe.rs

wasmtime_wasi/p2/
pipe.rs
