azure_core/
seekable_stream.rs

1use bytes::Bytes;
2use dyn_clone::DynClone;
3use futures::{io::AsyncRead, stream::Stream, task::Poll};
4use std::{pin::Pin, task::Context};
5
6/// Amount of the stream to buffer in memory during streaming uploads
7pub(crate) const DEFAULT_BUFFER_SIZE: usize = 1024 * 64;
8
9/// Enable a type implementing `AsyncRead` to be consumed as if it were
10/// a `Stream` of `Bytes`.
11#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
12#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
13pub trait SeekableStream: AsyncRead + Unpin + std::fmt::Debug + Send + Sync + DynClone {
14    async fn reset(&mut self) -> crate::error::Result<()>;
15    fn len(&self) -> usize;
16
17    fn is_empty(&self) -> bool {
18        self.len() == 0
19    }
20
21    fn buffer_size(&self) -> usize {
22        DEFAULT_BUFFER_SIZE
23    }
24}
25
26dyn_clone::clone_trait_object!(SeekableStream);
27
28impl Stream for dyn SeekableStream {
29    type Item = crate::error::Result<Bytes>;
30
31    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
32        let mut buffer = vec![0_u8; self.buffer_size()];
33
34        match self.poll_read(cx, &mut buffer) {
35            Poll::Ready(Ok(0)) => Poll::Ready(None),
36            Poll::Ready(Ok(bytes_read)) => {
37                let bytes: Bytes = buffer.into();
38                let bytes = bytes.slice(0..bytes_read);
39                Poll::Ready(Some(Ok(bytes)))
40            }
41            Poll::Ready(Err(err)) => Poll::Ready(Some(Err(crate::error::Error::full(
42                crate::error::ErrorKind::Io,
43                err,
44                "an error was encountered when trying to read from a stream",
45            )))),
46            Poll::Pending => Poll::Pending,
47        }
48    }
49}