ufile_rus3/api/
stream.rs

1use std::{
2    pin::Pin,
3    sync::{Arc, atomic::AtomicUsize},
4    task::Poll,
5};
6
7use bytes::Bytes;
8use futures_util::{
9    AsyncRead, Stream,
10    io::{BufReader, Cursor},
11};
12use pin_project_lite::pin_project;
13
14pin_project! {
15    pub struct ByteStream {
16        #[pin]
17        inner: Inner
18    }
19}
20
21struct Inner(bytes::Bytes);
22
23impl ByteStream {
24    pub fn from_bytes(bytes: Bytes) -> Self {
25        Self {
26            inner: Inner(bytes),
27        }
28    }
29}
30
31/// struct to wrap file reader with progress
32pub struct ProgressStream<T> {
33    reader: BufReader<T>,
34    progress: Arc<AtomicUsize>,
35    size: usize,
36}
37
38impl<T: AsyncRead + Unpin> ProgressStream<T> {
39    pub fn new(reader: T, size: usize) -> Self {
40        Self {
41            reader: BufReader::new(reader),
42            progress: Arc::new(AtomicUsize::new(0)),
43            size,
44        }
45    }
46
47    pub fn get_progress(&self) -> usize {
48        self.progress.load(std::sync::atomic::Ordering::Relaxed)
49    }
50}
51
52impl<T: AsyncRead + Unpin> Stream for ProgressStream<T> {
53    type Item = Result<bytes::Bytes, std::io::Error>;
54
55    fn poll_next(
56        self: std::pin::Pin<&mut Self>,
57        cx: &mut std::task::Context<'_>,
58    ) -> std::task::Poll<Option<Self::Item>> {
59        // 8kb buffer
60        let mut buffer = [0u8; 8092];
61        let this = self.get_mut();
62        let reader = Pin::new(&mut this.reader);
63        match reader.poll_read(cx, &mut buffer) {
64            std::task::Poll::Ready(Ok(n)) => {
65                if n == 0 {
66                    // we are at the end of file.
67                    return Poll::Ready(None);
68                }
69                let bytes = buffer[0..n].to_vec();
70                let num_bytes_read = bytes.len();
71                let prev = this
72                    .progress
73                    .fetch_add(num_bytes_read, std::sync::atomic::Ordering::Relaxed);
74                let current = prev + num_bytes_read;
75                // 计算并打印进度
76                let percent = (current as f64 / this.size as f64) * 100.0;
77                tracing::debug!(
78                    "Upload progress: {:.2}% ({} bytes/{} bytes)",
79                    percent,
80                    current,
81                    this.size
82                );
83                Poll::Ready(Some(Ok(Bytes::from_iter(bytes))))
84            }
85            std::task::Poll::Ready(Err(error)) => {
86                tracing::error!("Failed to read file, error: {:?}", error);
87                Poll::Ready(Some(Err(error)))
88            }
89            std::task::Poll::Pending => Poll::Pending,
90        }
91    }
92}
93
94impl From<ByteStream> for ProgressStream<Cursor<Bytes>> {
95    fn from(stream: ByteStream) -> Self {
96        let bytes = stream.inner.0;
97        let size = bytes.len();
98        let cursor = Cursor::new(bytes);
99        Self::new(cursor, size)
100    }
101}
ufile_rus3/api/stream.rs

ufile_rus3/api/
stream.rs
