1use 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#[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
103pub struct Measure<T> {
108 task: T,
109 inner: MeasureInner,
110}
111impl<T> Measure<T> {
112 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 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 pub fn metrics(&self) -> Var<Metrics> {
129 self.inner.metrics.read_only()
130 }
131
132 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 let self_ = unsafe { self.get_unchecked_mut() };
149
150 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 let self_ = unsafe { self.get_unchecked_mut() };
164
165 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 let self_ = unsafe { self.get_unchecked_mut() };
178
179 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 let self_ = unsafe { self.get_unchecked_mut() };
186
187 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 let self_ = unsafe { self.get_unchecked_mut() };
195
196 unsafe { Pin::new_unchecked(&mut self_.task) }.poll_fill_buf(cx)
198 }
199
200 fn consume(self: Pin<&mut Self>, amt: usize) {
201 let self_ = unsafe { self.get_unchecked_mut() };
203 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#[derive(Debug, Clone, PartialEq, Eq)]
252#[non_exhaustive]
253pub struct Metrics {
254 pub read_progress: (ByteLength, ByteLength),
256
257 pub read_speed: ByteLength,
259
260 pub write_progress: (ByteLength, ByteLength),
262
263 pub write_speed: ByteLength,
265
266 pub total_time: Duration,
269}
270impl Metrics {
271 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 pub static ref METRICS_ID: zng_state_map::StateId<Metrics>;
338}
339
340pub trait McBufErrorExt {
342 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
355pub 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 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 pub fn is_lazy(&self) -> bool {
420 self.lazy
421 }
422
423 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 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 let mut i = inner.clones[self_.index];
479 let mut ready;
480
481 match &inner.result {
482 ReadState::Running => {
483 ready = &inner.buf[i..];
486
487 if ready.is_empty() {
488 if self.lazy {
489 if inner.non_lazy_count == 0 {
490 return Poll::Ready(Err(Error::other(ONLY_NON_LAZY_ERROR_MSG)));
492 } else {
493 inner.lazy_wakers.push(cx.waker().clone());
495
496 return Poll::Pending;
498 }
499 }
500
501 ready = &[];
504
505 let waker = match inner.waker.push(cx.waker().clone()) {
506 Some(w) => w,
507 None => {
508 return Poll::Pending;
510 }
511 };
512
513 let min_i = inner.clones.iter().copied().min().unwrap();
514 if min_i > 0 {
515 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 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 for waker in inner.lazy_wakers.drain(..) {
539 waker.wake();
540 }
541
542 match result {
543 Ok(0) => {
544 inner.waker.cancel();
545
546 inner.buf.truncate(new_start);
548 inner.result = ReadState::Eof;
549 inner.source = None;
550
551 }
553 Ok(read) => {
554 inner.waker.cancel();
555
556 inner.buf.truncate(new_start + read);
558 ready = &inner.buf[i..];
559
560 }
562 Err(e) => {
563 inner.waker.cancel();
564
565 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 }
587 ReadState::Err(e) => return Poll::Ready(e.err()),
588 }
589
590 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#[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 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 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
647pub struct ReadLimited<S> {
651 source: S,
652 limit: usize,
653 on_limit: fn() -> std::io::Error,
654}
655impl<S> ReadLimited<S> {
656 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 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 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 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 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 unsafe { Pin::new_unchecked(&mut self_.source) }.poll_fill_buf(cx)
716 }
717
718 fn consume(self: Pin<&mut Self>, amt: usize) {
719 let self_ = unsafe { self.get_unchecked_mut() };
721 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 #[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}