Skip to main content

zng_task/
io.rs

1//! IO tasks.
2//!
3//! Most of the types in this module are re-exported from [`futures_lite::io`].
4//!
5//! [`futures_lite::io`]: https://docs.rs/futures-lite/latest/futures_lite/io/index.html
6
7use std::{
8    fmt,
9    io::{BufRead, ErrorKind, Read, Write},
10    pin::Pin,
11    sync::Arc,
12    task::{self, Poll},
13    time::Duration,
14};
15
16use crate::{McWaker, Progress};
17
18#[doc(no_inline)]
19pub use futures_lite::io::{
20    AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt, AsyncWrite, AsyncWriteExt, BoxedReader, BoxedWriter,
21    BufReader, BufWriter, Cursor, ReadHalf, WriteHalf, copy, empty, repeat, sink, split,
22};
23// TODO(breaking) add StreamExt to support lines().next()
24
25#[doc(no_inline)]
26pub use blocking::Unblock;
27
28use parking_lot::Mutex;
29use std::io::{Error, Result};
30use zng_time::{DInstant, INSTANT};
31use zng_txt::formatx;
32use zng_unit::{ByteLength, ByteUnits};
33use zng_var::{Var, impl_from_and_into_var, var};
34
35struct MeasureInner {
36    metrics: Var<Metrics>,
37    start_time: DInstant,
38    last_write: DInstant,
39    last_read: DInstant,
40}
41impl MeasureInner {
42    fn new(read_progress: (ByteLength, ByteLength), write_progress: (ByteLength, ByteLength)) -> Self {
43        let now = INSTANT.now();
44        Self {
45            metrics: var(Metrics {
46                read_progress,
47                read_speed: 0.bytes(),
48                write_progress,
49                write_speed: 0.bytes(),
50                total_time: Duration::ZERO,
51            }),
52            start_time: now,
53            last_write: now,
54            last_read: now,
55        }
56    }
57
58    fn on_read(&mut self, bytes: u64) {
59        if bytes == 0 {
60            return;
61        }
62
63        let bytes = bytes.bytes();
64
65        let now = INSTANT.now();
66        let elapsed = now - self.last_read;
67
68        self.last_read = now;
69        let read_speed = bytes_per_sec(bytes, elapsed);
70
71        let total_time = now - self.start_time;
72
73        self.metrics.modify(move |m| {
74            m.read_progress.0 += bytes;
75            m.read_speed = read_speed;
76            m.total_time = total_time;
77        });
78    }
79
80    fn on_write(&mut self, bytes: u64) {
81        if bytes == 0 {
82            return;
83        }
84
85        let bytes = bytes.bytes();
86
87        let now = INSTANT.now();
88        let elapsed = now - self.last_write;
89
90        self.last_write = now;
91        let write_speed = bytes_per_sec(bytes, elapsed);
92
93        let total_time = now - self.start_time;
94
95        self.metrics.modify(move |m| {
96            m.write_progress.0 += bytes;
97            m.write_speed = write_speed;
98            m.total_time = total_time;
99        });
100    }
101}
102
103/// Measure read/write task.
104///
105/// Metrics are updated after each read/write, if you read/write all bytes in one call
106/// the metrics will only update once.
107pub struct Measure<T> {
108    task: T,
109    inner: MeasureInner,
110}
111impl<T> Measure<T> {
112    /// Start measuring a new read/write task.
113    pub fn new(task: T, total_read: ByteLength, total_write: ByteLength) -> Self {
114        Self::new_ongoing(task, (0.bytes(), total_read), (0.bytes(), total_write))
115    }
116
117    /// Continue measuring a read/write task.
118    pub fn new_ongoing(task: T, read_progress: (ByteLength, ByteLength), write_progress: (ByteLength, ByteLength)) -> Self {
119        Measure {
120            task,
121            inner: MeasureInner::new(read_progress, write_progress),
122        }
123    }
124
125    /// Current metrics.
126    ///
127    /// This value is updated after every read/write.
128    pub fn metrics(&self) -> Var<Metrics> {
129        self.inner.metrics.read_only()
130    }
131
132    /// Unwrap the inner task and final metrics.
133    pub fn finish(self) -> (T, Metrics) {
134        let mut metrics = self.inner.metrics.get();
135        metrics.total_time = self.inner.start_time.elapsed();
136        (self.task, metrics)
137    }
138}
139
140fn bytes_per_sec(bytes: ByteLength, elapsed: Duration) -> ByteLength {
141    let bytes_per_sec = bytes.0 as u128 / elapsed.as_nanos() / Duration::from_secs(1).as_nanos();
142    ByteLength(bytes_per_sec as u64)
143}
144
145impl<T: AsyncRead> AsyncRead for Measure<T> {
146    fn poll_read(self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut [u8]) -> Poll<Result<usize>> {
147        // SAFETY: we don't move anything.
148        let self_ = unsafe { self.get_unchecked_mut() };
149
150        // SAFETY: we don't move task
151        match unsafe { Pin::new_unchecked(&mut self_.task) }.poll_read(cx, buf) {
152            Poll::Ready(Ok(bytes)) => {
153                self_.inner.on_read(bytes as u64);
154                Poll::Ready(Ok(bytes))
155            }
156            p => p,
157        }
158    }
159}
160impl<T: AsyncWrite> AsyncWrite for Measure<T> {
161    fn poll_write(self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &[u8]) -> Poll<Result<usize>> {
162        // SAFETY: we don't move anything.
163        let self_ = unsafe { self.get_unchecked_mut() };
164
165        // SAFETY: we don't move task
166        match unsafe { Pin::new_unchecked(&mut self_.task) }.poll_write(cx, buf) {
167            Poll::Ready(Ok(bytes)) => {
168                self_.inner.on_write(bytes as u64);
169                Poll::Ready(Ok(bytes))
170            }
171            p => p,
172        }
173    }
174
175    fn poll_flush(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<()>> {
176        // SAFETY: we don't move anything.
177        let self_ = unsafe { self.get_unchecked_mut() };
178
179        // SAFETY: we don't move task
180        unsafe { Pin::new_unchecked(&mut self_.task) }.poll_flush(cx)
181    }
182
183    fn poll_close(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<()>> {
184        // SAFETY: we don't move anything.
185        let self_ = unsafe { self.get_unchecked_mut() };
186
187        // SAFETY: we don't move task
188        unsafe { Pin::new_unchecked(&mut self_.task) }.poll_flush(cx)
189    }
190}
191impl<T: AsyncBufRead> AsyncBufRead for Measure<T> {
192    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<&[u8]>> {
193        // SAFETY: we don't move anything.
194        let self_ = unsafe { self.get_unchecked_mut() };
195
196        // SAFETY: we don't move task
197        unsafe { Pin::new_unchecked(&mut self_.task) }.poll_fill_buf(cx)
198    }
199
200    fn consume(self: Pin<&mut Self>, amt: usize) {
201        // SAFETY: we don't move anything.
202        let self_ = unsafe { self.get_unchecked_mut() };
203        // SAFETY: we don't move task
204        unsafe { Pin::new_unchecked(&mut self_.task) }.consume(amt);
205        self_.inner.on_read(amt as u64);
206    }
207}
208impl<T: Read> Read for Measure<T> {
209    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
210        match self.task.read(buf) {
211            Ok(bytes) => {
212                self.inner.on_read(bytes as u64);
213                Ok(bytes)
214            }
215            r => r,
216        }
217    }
218}
219impl<T: Write> Write for Measure<T> {
220    fn write(&mut self, buf: &[u8]) -> Result<usize> {
221        match self.task.write(buf) {
222            Ok(bytes) => {
223                self.inner.on_write(bytes as u64);
224                Ok(bytes)
225            }
226            r => r,
227        }
228    }
229
230    fn flush(&mut self) -> Result<()> {
231        self.task.flush()
232    }
233}
234impl<T: BufRead> BufRead for Measure<T> {
235    fn fill_buf(&mut self) -> Result<&[u8]> {
236        self.task.fill_buf()
237    }
238
239    fn consume(&mut self, amount: usize) {
240        self.task.consume(amount);
241        self.inner.on_read(amount as u64);
242    }
243}
244
245/// Information about the state of an async IO task.
246///
247/// Read is also called *receive* or *download*. Write is also called *send* or *upload*. The default
248/// display print uses arrows ↓ and ↑ for read and write.
249///
250/// Use [`Measure`] to measure a task.
251#[derive(Debug, Clone, PartialEq, Eq)]
252#[non_exhaustive]
253pub struct Metrics {
254    /// Number of bytes read / estimated total.
255    pub read_progress: (ByteLength, ByteLength),
256
257    /// Average read speed in bytes/second.
258    pub read_speed: ByteLength,
259
260    /// Number of bytes written / estimated total.
261    pub write_progress: (ByteLength, ByteLength),
262
263    /// Average write speed in bytes/second.
264    pub write_speed: ByteLength,
265
266    /// Total time for the entire task. This will continuously increase until
267    /// the task is finished.
268    pub total_time: Duration,
269}
270impl Metrics {
271    /// All zeros.
272    pub fn zero() -> Self {
273        Self {
274            read_progress: (0.bytes(), 0.bytes()),
275            read_speed: 0.bytes(),
276            write_progress: (0.bytes(), 0.bytes()),
277            write_speed: 0.bytes(),
278            total_time: Duration::ZERO,
279        }
280    }
281}
282impl fmt::Display for Metrics {
283    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
284        let mut nl = false;
285        if self.read_progress.1 > 0.bytes() {
286            nl = true;
287            if self.read_progress.0 != self.read_progress.1 {
288                if self.read_progress.0 <= self.read_progress.1 {
289                    write!(f, "↓ {}-{}, {}/s", self.read_progress.0, self.read_progress.1, self.read_speed)?;
290                } else {
291                    write!(f, "↓ {}, {}/s", self.read_progress.0, self.read_speed)?;
292                }
293            } else {
294                write!(f, "↓ {} . {:?}", self.read_progress.0, self.total_time)?;
295            }
296        }
297        if self.write_progress.1 > 0.bytes() {
298            if nl {
299                writeln!(f)?;
300            }
301            if self.write_progress.0 != self.write_progress.1 {
302                if self.write_progress.0 <= self.write_progress.1 {
303                    write!(f, "↑ {}-{}, {}/s", self.write_progress.0, self.write_progress.1, self.write_speed)?;
304                } else {
305                    write!(f, "↑ {}, {}/s", self.write_progress.0, self.write_speed)?;
306                }
307            } else {
308                write!(f, "↑ {} . {:?}", self.write_progress.0, self.total_time)?;
309            }
310        }
311
312        Ok(())
313    }
314}
315impl_from_and_into_var! {
316    fn from(metrics: Metrics) -> Progress {
317        let mut status = Progress::indeterminate();
318        if metrics.read_progress.1 > 0.bytes() {
319            status = Progress::from_n_of(metrics.read_progress.0.0, metrics.read_progress.1.0);
320        }
321        if metrics.write_progress.1 > 0.bytes() {
322            let w_status = Progress::from_n_of(metrics.write_progress.0.0, metrics.write_progress.1.0);
323            if status.is_indeterminate() {
324                status = w_status;
325            } else {
326                status = status.and_fct(w_status.fct());
327            }
328        }
329        status.with_msg(formatx!("{metrics}")).with_meta_mut(|mut m| {
330            m.set(*METRICS_ID, metrics);
331        })
332    }
333}
334
335zng_state_map::static_id! {
336    /// Metrics in a [`Progress::with_meta`] metadata.
337    pub static ref METRICS_ID: zng_state_map::StateId<Metrics>;
338}
339
340/// Extension methods for [`std::io::Error`] to be used with errors returned by [`McBufReader`].
341pub trait McBufErrorExt {
342    /// Returns `true` if this error represents the condition where there are only [`McBufReader::is_lazy`] readers
343    /// left, the buffer is drained and the inner reader is not EOF.
344    ///
345    /// You can recover from this error by turning the reader non-lazy using [`McBufReader::set_lazy`].
346    fn is_only_lazy_left(&self) -> bool;
347}
348impl McBufErrorExt for std::io::Error {
349    fn is_only_lazy_left(&self) -> bool {
350        matches!(self.kind(), ErrorKind::Other) && format!("{self:?}").contains(ONLY_NON_LAZY_ERROR_MSG)
351    }
352}
353const ONLY_NON_LAZY_ERROR_MSG: &str = "no non-lazy readers left to read";
354
355/// Multiple consumer buffered read.
356///
357/// Clone an instance to create a new consumer, already read bytes stay in the buffer until all clones have read it,
358/// clones continue reading from the same offset as the reader they cloned.
359///
360/// A single instance of this reader behaves like a `BufReader`.
361///
362/// # Result
363///
364/// The result is *repeats* ready when `EOF` or an [`Error`] occurs, unfortunately the IO error is not cloneable
365/// so the error is recreated using [`CloneableError`] for subsequent poll attempts.
366///
367/// The inner reader is dropped as soon as it finishes.
368///
369/// # Lazy Clones
370///
371/// You can mark clones as [lazy], lazy clones don't pull from the inner reader, only advance when another clone reads, if
372/// all living clones are lazy they stop reading with an error. You can identify this custom error using the [`McBufErrorExt::is_only_lazy_left`]
373/// extension method.
374///
375/// [lazy]: Self::set_lazy
376pub struct McBufReader<S: AsyncRead> {
377    inner: Arc<Mutex<McBufInner<S>>>,
378    index: usize,
379    lazy: bool,
380}
381struct McBufInner<S: AsyncRead> {
382    source: Option<S>,
383    waker: McWaker,
384    lazy_wakers: Vec<task::Waker>,
385
386    buf: Vec<u8>,
387
388    clones: Vec<usize>,
389    non_lazy_count: usize,
390
391    result: ReadState,
392}
393impl<S: AsyncRead> McBufReader<S> {
394    /// Creates a buffered reader.
395    pub fn new(source: S) -> Self {
396        let mut clones = Vec::with_capacity(2);
397        clones.push(0);
398        McBufReader {
399            inner: Arc::new(Mutex::new(McBufInner {
400                source: Some(source),
401                waker: McWaker::empty(),
402                lazy_wakers: vec![],
403
404                buf: Vec::with_capacity(10.kilobytes().0 as usize),
405
406                clones,
407                non_lazy_count: 1,
408
409                result: ReadState::Running,
410            })),
411            index: 0,
412            lazy: false,
413        }
414    }
415
416    /// Returns `true` if this reader does not pull from the inner reader, only advancing when a non-lazy reader advances.
417    ///
418    /// The initial reader is not lazy, only clones of lazy readers are lazy by default.
419    pub fn is_lazy(&self) -> bool {
420        self.lazy
421    }
422
423    /// Sets [`is_lazy`].
424    ///
425    /// [`is_lazy`]: Self::is_lazy
426    pub fn set_lazy(&mut self, lazy: bool) {
427        if self.lazy != lazy {
428            if lazy {
429                self.inner.lock().non_lazy_count -= 1;
430            } else {
431                self.inner.lock().non_lazy_count += 1;
432            }
433            self.lazy = lazy;
434        }
435    }
436}
437impl<S: AsyncRead> Clone for McBufReader<S> {
438    fn clone(&self) -> Self {
439        let mut inner = self.inner.lock();
440
441        let offset = inner.clones[self.index];
442        let index = inner.clones.len();
443        inner.clones.push(offset);
444
445        if !self.lazy {
446            inner.non_lazy_count += 1;
447        }
448
449        Self {
450            inner: self.inner.clone(),
451            index,
452            lazy: self.lazy,
453        }
454    }
455}
456impl<S: AsyncRead> Drop for McBufReader<S> {
457    fn drop(&mut self) {
458        let mut inner = self.inner.lock();
459        inner.clones[self.index] = usize::MAX;
460        if !self.lazy {
461            inner.non_lazy_count -= 1;
462            if inner.non_lazy_count == 0 {
463                // notify lazy so they get the error.
464                for waker in inner.lazy_wakers.drain(..) {
465                    waker.wake();
466                }
467            }
468        }
469    }
470}
471impl<S: AsyncRead> AsyncRead for McBufReader<S> {
472    fn poll_read(self: Pin<&mut Self>, cx: &mut task::Context<'_>, buf: &mut [u8]) -> Poll<Result<usize>> {
473        let self_ = self.as_ref();
474        let mut inner = self_.inner.lock();
475        let inner = &mut *inner;
476
477        // ready data for this clone.
478        let mut i = inner.clones[self_.index];
479        let mut ready;
480
481        match &inner.result {
482            ReadState::Running => {
483                // source has not finished yet.
484
485                ready = &inner.buf[i..];
486
487                if ready.is_empty() {
488                    if self.lazy {
489                        if inner.non_lazy_count == 0 {
490                            // user can make this reader non-lazy and try again.
491                            return Poll::Ready(Err(Error::other(ONLY_NON_LAZY_ERROR_MSG)));
492                        } else {
493                            // register waker for after non-lazy poll.
494                            inner.lazy_wakers.push(cx.waker().clone());
495
496                            // wait non-lazy to pull.
497                            return Poll::Pending;
498                        }
499                    }
500
501                    // time to poll source.
502
503                    ready = &[];
504
505                    let waker = match inner.waker.push(cx.waker().clone()) {
506                        Some(w) => w,
507                        None => {
508                            // already polling from another clone.
509                            return Poll::Pending;
510                        }
511                    };
512
513                    let min_i = inner.clones.iter().copied().min().unwrap();
514                    if min_i > 0 {
515                        // reuse front.
516                        inner.buf.copy_within(min_i.., 0);
517                        inner.buf.truncate(inner.buf.len() - min_i);
518
519                        i -= min_i;
520                        for i in &mut inner.clones {
521                            *i -= min_i;
522                        }
523                    }
524
525                    let new_start = inner.buf.len();
526
527                    inner.buf.resize(inner.buf.len() + buf.len().max(10.kilobytes().0 as usize), 0);
528
529                    let mut inner_cx = task::Context::from_waker(&waker);
530
531                    // SAFETY: we don't move `source`.
532                    let source = unsafe { Pin::new_unchecked(inner.source.as_mut().unwrap()) };
533                    let result = source.poll_read(&mut inner_cx, &mut inner.buf[new_start..]);
534
535                    match result {
536                        Poll::Ready(result) => {
537                            // notify lazy readers.
538                            for waker in inner.lazy_wakers.drain(..) {
539                                waker.wake();
540                            }
541
542                            match result {
543                                Ok(0) => {
544                                    inner.waker.cancel();
545
546                                    // EOF
547                                    inner.buf.truncate(new_start);
548                                    inner.result = ReadState::Eof;
549                                    inner.source = None;
550
551                                    // continue 'copy ready
552                                }
553                                Ok(read) => {
554                                    inner.waker.cancel();
555
556                                    // Read > 0
557                                    inner.buf.truncate(new_start + read);
558                                    ready = &inner.buf[i..];
559
560                                    // continue 'copy ready
561                                }
562                                Err(e) => {
563                                    inner.waker.cancel();
564
565                                    // Error
566                                    inner.result = ReadState::Err(CloneableError::new(&e));
567                                    inner.buf = vec![];
568                                    inner.source = None;
569
570                                    return Poll::Ready(Err(e));
571                                }
572                            }
573                        }
574
575                        Poll::Pending => {
576                            inner.buf.truncate(new_start);
577                            return Poll::Pending;
578                        }
579                    }
580                }
581            }
582            ReadState::Eof => {
583                ready = &inner.buf[i..];
584
585                // continue 'copy ready
586            }
587            ReadState::Err(e) => return Poll::Ready(e.err()),
588        }
589
590        // 'copy ready
591
592        let max_ready = buf.len().min(ready.len());
593        buf[..max_ready].copy_from_slice(&ready[..max_ready]);
594
595        i += max_ready;
596        inner.clones[self_.index] = i;
597
598        Poll::Ready(Ok(max_ready))
599    }
600}
601
602/// Represents the cloneable parts of an [`Error`].
603///
604/// Unfortunately [`Error`] does not implement clone, this is needed to implemented
605/// IO futures that repeat the ready result after subsequent polls. This type partially
606/// works around the issue by copying enough information to recreate an error that is still useful.
607///
608/// The OS error code, [`ErrorKind`] and display message are preserved. Note that this not an error type,
609/// it must be converted to [`Error`] using `into` or [`err`].
610///
611/// [`err`]: Self::err
612#[derive(Clone)]
613pub struct CloneableError {
614    info: ErrorInfo,
615}
616#[derive(Clone)]
617enum ErrorInfo {
618    OsError(i32),
619    Other(ErrorKind, String),
620}
621impl CloneableError {
622    /// Copy the cloneable information from the [`Error`].
623    pub fn new(e: &Error) -> Self {
624        let info = if let Some(code) = e.raw_os_error() {
625            ErrorInfo::OsError(code)
626        } else {
627            ErrorInfo::Other(e.kind(), format!("{e}"))
628        };
629
630        Self { info }
631    }
632
633    /// Returns an `Err(Error)` generated from the cloneable information.
634    pub fn err<T>(&self) -> Result<T> {
635        Err(self.clone().into())
636    }
637}
638impl From<CloneableError> for Error {
639    fn from(e: CloneableError) -> Self {
640        match e.info {
641            ErrorInfo::OsError(code) => Error::from_raw_os_error(code),
642            ErrorInfo::Other(kind, msg) => Error::new(kind, msg),
643        }
644    }
645}
646
647/// Represents a stream reader that generates an error if the source stream exceeds a limit.
648///
649/// Note that some bytes over the limit may be read once if the source stream is buffered.
650pub struct ReadLimited<S> {
651    source: S,
652    limit: usize,
653    on_limit: fn() -> std::io::Error,
654}
655impl<S> ReadLimited<S> {
656    /// Construct a limited reader.
657    ///
658    /// The `on_limit` closure is called for every read attempt after the limit is reached.
659    pub fn new(source: S, limit: ByteLength, on_limit: fn() -> std::io::Error) -> Self {
660        Self {
661            source,
662            limit: limit.0.try_into().unwrap_or(usize::MAX),
663            on_limit,
664        }
665    }
666
667    /// New with default on limit error.
668    pub fn new_default_err(source: S, limit: ByteLength) -> Self {
669        Self::new(source, limit, || {
670            std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "source exceeded read limit")
671        })
672    }
673}
674impl<S> AsyncRead for ReadLimited<S>
675where
676    S: AsyncRead,
677{
678    fn poll_read(self: Pin<&mut Self>, cx: &mut task::Context<'_>, mut buf: &mut [u8]) -> Poll<Result<usize>> {
679        // SAFETY: we don't move anything.
680        let self_ = unsafe { self.get_unchecked_mut() };
681
682        if self_.limit == 0 {
683            let err = (self_.on_limit)();
684            return Poll::Ready(Err(err));
685        }
686
687        if buf.len() > self_.limit {
688            buf = &mut buf[..self_.limit];
689        }
690
691        // SAFETY: we never move `source`.
692        match unsafe { Pin::new_unchecked(&mut self_.source) }.poll_read(cx, buf) {
693            Poll::Ready(Ok(n)) => {
694                self_.limit = self_.limit.saturating_sub(n);
695                Poll::Ready(Ok(n))
696            }
697            r => r,
698        }
699    }
700}
701impl<S> AsyncBufRead for ReadLimited<S>
702where
703    S: AsyncBufRead,
704{
705    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Result<&[u8]>> {
706        // SAFETY: we don't move anything.
707        let self_ = unsafe { self.get_unchecked_mut() };
708
709        if self_.limit == 0 {
710            let err = (self_.on_limit)();
711            return Poll::Ready(Err(err));
712        }
713
714        // SAFETY: we never move `source`.
715        unsafe { Pin::new_unchecked(&mut self_.source) }.poll_fill_buf(cx)
716    }
717
718    fn consume(self: Pin<&mut Self>, amt: usize) {
719        // SAFETY: we don't move anything.
720        let self_ = unsafe { self.get_unchecked_mut() };
721        // SAFETY: we never move `source`.
722        unsafe { Pin::new_unchecked(&mut self_.source) }.consume(amt);
723        self_.limit = self_.limit.saturating_sub(amt);
724    }
725}
726impl<S> Read for ReadLimited<S>
727where
728    S: Read,
729{
730    fn read(&mut self, mut buf: &mut [u8]) -> Result<usize> {
731        if self.limit == 0 {
732            let err = (self.on_limit)();
733            return Err(err);
734        }
735
736        if buf.len() > self.limit {
737            buf = &mut buf[..self.limit];
738        }
739
740        match self.source.read(buf) {
741            Ok(n) => {
742                self.limit = self.limit.saturating_sub(n);
743                Ok(n)
744            }
745            r => r,
746        }
747    }
748}
749impl<S> BufRead for ReadLimited<S>
750where
751    S: BufRead,
752{
753    fn fill_buf(&mut self) -> Result<&[u8]> {
754        if self.limit == 0 {
755            let err = (self.on_limit)();
756            return Err(err);
757        }
758
759        self.source.fill_buf()
760    }
761
762    fn consume(&mut self, amount: usize) {
763        self.source.consume(amount);
764        self.limit = self.limit.saturating_sub(amount);
765    }
766}
767
768enum ReadState {
769    Running,
770    Eof,
771    Err(CloneableError),
772}
773
774#[cfg(test)]
775mod tests {
776    use super::*;
777    use crate as task;
778    use zng_unit::TimeUnits;
779
780    #[test]
781    pub fn mc_buf_reader_parallel() {
782        let data = Data::new(60.kilobytes().0 as _);
783
784        let mut expected = vec![0; data.len];
785        let _ = data.clone().blocking_read(&mut expected[..]);
786
787        let mut a = McBufReader::new(data);
788        let mut b = a.clone();
789        let mut c = a.clone();
790
791        let (a, b, c) = async_test(async move {
792            let a = task::run(async move {
793                let mut buf = vec![];
794                a.read_to_end(&mut buf).await.unwrap();
795                buf
796            });
797            let b = task::run(async move {
798                let mut buf: Vec<u8> = vec![];
799                b.read_to_end(&mut buf).await.unwrap();
800                buf
801            });
802            let c = task::run(async move {
803                let mut buf: Vec<u8> = vec![];
804                c.read_to_end(&mut buf).await.unwrap();
805                buf
806            });
807
808            task::all!(a, b, c).await
809        });
810
811        crate::assert_vec_eq!(expected, a);
812        crate::assert_vec_eq!(expected, b);
813        crate::assert_vec_eq!(expected, c);
814    }
815
816    #[test]
817    pub fn mc_buf_reader_single() {
818        let data = Data::new(60.kilobytes().0 as _);
819
820        let mut expected = vec![0; data.len];
821        let _ = data.clone().blocking_read(&mut expected[..]);
822
823        let mut a = McBufReader::new(data);
824
825        let a = async_test(async move {
826            let a = task::run(async move {
827                let mut buf = vec![];
828                a.read_to_end(&mut buf).await.unwrap();
829                buf
830            });
831
832            a.await
833        });
834
835        crate::assert_vec_eq!(expected, a);
836    }
837
838    #[test]
839    pub fn mc_buf_reader_sequential() {
840        let data = Data::new(60.kilobytes().0 as _);
841
842        let mut expected = vec![0; data.len];
843        let _ = data.clone().blocking_read(&mut expected[..]);
844
845        let mut clones = vec![McBufReader::new(data)];
846        for _ in 0..5 {
847            clones.push(clones[0].clone());
848        }
849
850        let r = async_test(async move {
851            let mut r = vec![];
852
853            for mut clone in clones {
854                let mut buf = vec![];
855                clone.read_to_end(&mut buf).await.unwrap();
856                r.push(buf);
857            }
858
859            r
860        });
861
862        for r in r {
863            crate::assert_vec_eq!(expected, r);
864        }
865    }
866
867    #[test]
868    pub fn mc_buf_reader_completed() {
869        let data = Data::new(60.kilobytes().0 as _);
870        let mut buf = Vec::with_capacity(data.len);
871        let mut a = McBufReader::new(data);
872
873        let r = async_test(async move {
874            a.read_to_end(&mut buf).await.unwrap();
875
876            let mut b = a.clone();
877            buf.clear();
878
879            b.read_to_end(&mut buf).await.unwrap();
880            buf.len()
881        });
882
883        assert_eq!(0, r);
884    }
885
886    #[test]
887    pub fn mc_buf_reader_error() {
888        let mut data = Data::new(20.kilobytes().0 as _);
889        data.set_error();
890
891        let mut expected = vec![0; data.len];
892        let _ = data.clone().blocking_read(&mut expected[..]);
893
894        let mut a = McBufReader::new(data);
895        let mut b = a.clone();
896
897        let (a, b) = async_test(async move {
898            let a = task::run(async move {
899                let mut buf = vec![];
900                a.read_to_end(&mut buf).await.unwrap_err()
901            });
902            let b = task::run(async move {
903                let mut buf: Vec<u8> = vec![];
904                b.read_to_end(&mut buf).await.unwrap_err()
905            });
906
907            task::all!(a, b).await
908        });
909
910        assert_eq!(ErrorKind::InvalidData, a.kind());
911        assert_eq!(ErrorKind::InvalidData, b.kind());
912    }
913
914    #[test]
915    pub fn mc_buf_reader_error_completed() {
916        let mut data = Data::new(20.kilobytes().0 as _);
917        data.set_error();
918
919        let mut buf = Vec::with_capacity(data.len);
920        let mut a = McBufReader::new(data);
921
922        let (a, b) = async_test(async move {
923            let a_err = a.read_to_end(&mut buf).await.unwrap_err();
924
925            let mut b = a.clone();
926            buf.clear();
927
928            let b_err = b.read_to_end(&mut buf).await.unwrap_err();
929
930            (a_err, b_err)
931        });
932
933        assert_eq!(ErrorKind::InvalidData, a.kind());
934        assert_eq!(ErrorKind::InvalidData, b.kind());
935    }
936
937    #[test]
938    pub fn mc_buf_reader_parallel_with_delay1() {
939        let mut data = Data::new(60.kilobytes().0 as _);
940        data.enable_pending();
941
942        let mut expected = vec![0; data.len];
943        let _ = data.clone().blocking_read(&mut expected[..]);
944
945        let mut a = McBufReader::new(data);
946        let mut b = a.clone();
947        let mut c = a.clone();
948
949        let (a, b, c) = async_test(async move {
950            let a = task::run(async move {
951                let mut buf = vec![];
952                a.read_to_end(&mut buf).await.unwrap();
953                buf
954            });
955            let b = task::run(async move {
956                let mut buf: Vec<u8> = vec![];
957                b.read_to_end(&mut buf).await.unwrap();
958                buf
959            });
960            let c = task::run(async move {
961                let mut buf: Vec<u8> = vec![];
962                c.read_to_end(&mut buf).await.unwrap();
963                buf
964            });
965
966            task::all!(a, b, c).await
967        });
968
969        crate::assert_vec_eq!(expected, a);
970        crate::assert_vec_eq!(expected, b);
971        crate::assert_vec_eq!(expected, c);
972    }
973
974    #[test]
975    pub fn mc_buf_reader_parallel_with_delay2() {
976        let mut data = Data::new(60.kilobytes().0 as _);
977        data.enable_pending();
978
979        let mut expected = vec![0; data.len];
980        let _ = data.clone().blocking_read(&mut expected[..]);
981
982        let mut a = McBufReader::new(data);
983        let mut b = a.clone();
984        let mut c = a.clone();
985
986        let (a, b, c) = async_test(async move {
987            let a = task::run(async move {
988                let mut buf = vec![];
989                a.read_to_end(&mut buf).await.unwrap();
990                buf
991            });
992            let b = task::run(async move {
993                let mut buf: Vec<u8> = vec![];
994                task::deadline(5.ms()).await;
995                b.read_to_end(&mut buf).await.unwrap();
996                buf
997            });
998            let c = task::run(async move {
999                let mut buf: Vec<u8> = vec![];
1000                c.read_to_end(&mut buf).await.unwrap();
1001                buf
1002            });
1003
1004            task::all!(a, b, c).await
1005        });
1006
1007        crate::assert_vec_eq!(expected, a);
1008        crate::assert_vec_eq!(expected, b);
1009        crate::assert_vec_eq!(expected, c);
1010    }
1011
1012    #[derive(Clone)]
1013    struct Data {
1014        b: u8,
1015        len: usize,
1016        error: Option<CloneableError>,
1017        delay: Duration,
1018        pending: bool,
1019    }
1020    impl Data {
1021        pub fn new(len: usize) -> Self {
1022            Self {
1023                b: 0,
1024                len,
1025                error: None,
1026                delay: 0.ms(),
1027                pending: false,
1028            }
1029        }
1030        pub fn blocking_read(&mut self, buf: &mut [u8]) -> Result<usize> {
1031            let len = self.len;
1032            for b in buf.iter_mut().take(len) {
1033                *b = self.b;
1034                self.len -= 1;
1035                self.b = self.b.wrapping_add(1);
1036            }
1037
1038            if len == 0
1039                && let Some(e) = &self.error
1040            {
1041                return e.err();
1042            }
1043
1044            Ok(buf.len().min(len))
1045        }
1046        pub fn set_error(&mut self) {
1047            self.error = Some(CloneableError::new(&Error::new(ErrorKind::InvalidData, "test error")));
1048        }
1049
1050        pub fn enable_pending(&mut self) {
1051            self.delay = 3.ms();
1052        }
1053    }
1054    impl AsyncRead for Data {
1055        fn poll_read(mut self: Pin<&mut Self>, cx: &mut std::task::Context<'_>, buf: &mut [u8]) -> Poll<Result<usize>> {
1056            if self.delay > Duration::ZERO {
1057                self.pending = !self.pending;
1058                if self.pending {
1059                    let waker = cx.waker().clone();
1060                    let delay = self.delay;
1061                    task::spawn(async move {
1062                        task::deadline(delay).await;
1063                        waker.wake();
1064                    });
1065                    return Poll::Pending;
1066                }
1067            }
1068
1069            let r = self.as_mut().blocking_read(buf);
1070            Poll::Ready(r)
1071        }
1072    }
1073
1074    #[track_caller]
1075    fn async_test<F>(test: F) -> F::Output
1076    where
1077        F: Future,
1078    {
1079        task::block_on(task::with_deadline(test, 5.secs())).unwrap()
1080    }
1081
1082    /// Assert vector equality with better error message.
1083    #[macro_export]
1084    macro_rules! assert_vec_eq {
1085        ($a:expr, $b: expr) => {
1086            match (&$a, &$b) {
1087                (ref a, ref b) => {
1088                    let len_not_eq = a.len() != b.len();
1089                    let mut data_not_eq = None;
1090                    for (i, (a, b)) in a.iter().zip(b.iter()).enumerate() {
1091                        if a != b {
1092                            data_not_eq = Some(i);
1093                            break;
1094                        }
1095                    }
1096
1097                    if len_not_eq || data_not_eq.is_some() {
1098                        use std::fmt::*;
1099
1100                        let mut error = format!("`{}` != `{}`", stringify!($a), stringify!($b));
1101                        if len_not_eq {
1102                            let _ = write!(&mut error, "\n  lengths not equal: {} != {}", a.len(), b.len());
1103                        }
1104                        if let Some(i) = data_not_eq {
1105                            let _ = write!(&mut error, "\n  data not equal at index {}: {} != {:?}", i, a[i], b[i]);
1106                        }
1107                        panic!("{error}")
1108                    }
1109                }
1110            }
1111        };
1112    }
1113}