Skip to main content

zerofs_client/
stream.rs

1//! `stream` feature: a [`futures_core::Stream`] adapter over [`Dir`] listing,
2//! so a directory can be consumed with `StreamExt` (`.next()`, `.collect()`,
3//! `try_for_each`, ...). Rust-only; never crosses the FFI boundary.
4
5use crate::dir::Dir;
6use crate::error::ZeroFsError;
7use crate::types::DirEntry;
8use futures_core::Stream;
9use std::collections::VecDeque;
10use std::future::Future;
11use std::pin::Pin;
12use std::sync::Arc;
13use std::task::{Context, Poll, ready};
14
15#[cfg(not(target_arch = "wasm32"))]
16type BoxFut<T> = Pin<Box<dyn Future<Output = T> + Send>>;
17#[cfg(target_arch = "wasm32")]
18type BoxFut<T> = Pin<Box<dyn Future<Output = T>>>;
19
20/// A pull-based stream of [`DirEntry`] over an open [`Dir`], yielding entries
21/// one at a time (fetched a server batch at a time underneath). It holds the
22/// `Dir` open for its lifetime; an error ends the stream after it is yielded.
23pub struct DirStream {
24    dir: Arc<Dir>,
25    buf: VecDeque<DirEntry>,
26    fetch: Option<BoxFut<Result<Vec<DirEntry>, ZeroFsError>>>,
27    done: bool,
28}
29
30impl DirStream {
31    pub(crate) fn new(dir: Arc<Dir>) -> Self {
32        Self {
33            dir,
34            buf: VecDeque::new(),
35            fetch: None,
36            done: false,
37        }
38    }
39}
40
41impl Stream for DirStream {
42    type Item = Result<DirEntry, ZeroFsError>;
43
44    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
45        loop {
46            if let Some(entry) = self.buf.pop_front() {
47                return Poll::Ready(Some(Ok(entry)));
48            }
49            if self.done {
50                return Poll::Ready(None);
51            }
52            if self.fetch.is_none() {
53                let dir = Arc::clone(&self.dir);
54                self.fetch = Some(Box::pin(async move { dir.next_batch(None).await }));
55            }
56            let batch = ready!(self.fetch.as_mut().unwrap().as_mut().poll(cx));
57            self.fetch = None;
58            match batch {
59                // An empty batch is end-of-directory.
60                Ok(batch) if batch.is_empty() => {
61                    self.done = true;
62                    return Poll::Ready(None);
63                }
64                Ok(batch) => self.buf.extend(batch),
65                Err(e) => {
66                    self.done = true;
67                    return Poll::Ready(Some(Err(e)));
68                }
69            }
70        }
71    }
72}