@@ -20,13 +20,18 @@ use std::ops::DerefMut;
2020use std:: pin:: Pin ;
2121use std:: task:: { Context , Poll } ;
2222use tokio:: io:: AsyncRead ;
23+ use tokio:: stream:: Stream ;
24+ use tokio:: time:: { delay_for, Delay , Duration , Instant } ;
25+
26+ const MAX_READ_AMOUNT : usize = 500 * 1000 ; // 0.5MB
27+ const GRACE_DURATION : Duration = Duration :: from_millis ( 1 ) ;
2328
2429pub struct AvailableReader < ' a , R : AsyncRead + Unpin > {
25- // TODO: come up with a way to avoid using RefCell (not sure if possible though)
30+ // TODO: come up with a way to avoid using RefCell (not sure if possible though due to having to
31+ // mutably borrow both inner reader and buffer at the same time)
2632 buf : RefCell < BytesMut > ,
2733 inner : RefCell < & ' a mut R > ,
28- // idea for the future: tiny delay that allows to prevent unnecessary extra fragmentation
29- // grace_period: Option<Delay>,
34+ grace_period : Option < Delay > ,
3035}
3136
3237impl < ' a , R > AvailableReader < ' a , R >
@@ -39,20 +44,15 @@ where
3944 AvailableReader {
4045 buf : RefCell :: new ( BytesMut :: with_capacity ( Self :: BUF_INCREMENT ) ) ,
4146 inner : RefCell :: new ( reader) ,
42- // grace_period: None ,
47+ grace_period : Some ( delay_for ( GRACE_DURATION ) ) ,
4348 }
4449 }
4550}
4651
47- // TODO: change this guy to a stream? Seems waaay more appropriate considering
48- // we're getting new Bytes items regularly rather than calling it once.
49-
50- impl < ' a , R : AsyncRead + Unpin > Future for AvailableReader < ' a , R > {
51- type Output = io:: Result < ( Bytes , bool ) > ;
52+ impl < ' a , R : AsyncRead + Unpin > Stream for AvailableReader < ' a , R > {
53+ type Item = io:: Result < Bytes > ;
5254
53- // this SHOULD stay mutable, because we rely on runtime checks inside the method
54- #[ allow( unused_mut) ]
55- fn poll ( mut self : Pin < & mut Self > , cx : & mut Context < ' _ > ) -> Poll < Self :: Output > {
55+ fn poll_next ( mut self : Pin < & mut Self > , cx : & mut Context < ' _ > ) -> Poll < Option < Self :: Item > > {
5656 // if we have no space in buffer left - expand it
5757 if !self . buf . borrow ( ) . has_remaining_mut ( ) {
5858 self . buf . borrow_mut ( ) . reserve ( Self :: BUF_INCREMENT ) ;
@@ -68,19 +68,43 @@ impl<'a, R: AsyncRead + Unpin> Future for AvailableReader<'a, R> {
6868 if self . buf . borrow ( ) . is_empty ( ) {
6969 Poll :: Pending
7070 } else {
71+ // if exists - check grace period
72+ if let Some ( grace_period) = self . grace_period . as_mut ( ) {
73+ if Pin :: new ( grace_period) . poll ( cx) . is_pending ( ) {
74+ return Poll :: Pending ;
75+ }
76+ }
77+
7178 let buf = self . buf . replace ( BytesMut :: new ( ) ) ;
72- Poll :: Ready ( Ok ( ( buf. freeze ( ) , false ) ) )
79+ Poll :: Ready ( Some ( Ok ( buf. freeze ( ) ) ) )
7380 }
7481 }
75- Poll :: Ready ( Err ( err) ) => Poll :: Ready ( Err ( err) ) ,
82+ Poll :: Ready ( Err ( err) ) => Poll :: Ready ( Some ( Err ( err) ) ) ,
7683 Poll :: Ready ( Ok ( n) ) => {
84+ // if exists - reset grace period
85+ if let Some ( grace_period) = self . grace_period . as_mut ( ) {
86+ let now = Instant :: now ( ) ;
87+ grace_period. reset ( now + GRACE_DURATION ) ;
88+ }
89+
7790 // if we read a non-0 amount, we're not done yet!
7891 if n == 0 {
7992 let buf = self . buf . replace ( BytesMut :: new ( ) ) ;
80- Poll :: Ready ( Ok ( ( buf. freeze ( ) , true ) ) )
93+ if buf. len ( ) > 0 {
94+ Poll :: Ready ( Some ( Ok ( buf. freeze ( ) ) ) )
95+ } else {
96+ Poll :: Ready ( None )
97+ }
8198 } else {
8299 // tell the waker we should be polled again!
83100 cx. waker ( ) . wake_by_ref ( ) ;
101+
102+ // if we reached our maximum amount - return it
103+ let read_bytes_len = self . buf . borrow ( ) . len ( ) ;
104+ if read_bytes_len >= MAX_READ_AMOUNT {
105+ let buf = self . buf . replace ( BytesMut :: new ( ) ) ;
106+ return Poll :: Ready ( Some ( Ok ( buf. freeze ( ) ) ) ) ;
107+ }
84108 Poll :: Pending
85109 }
86110 }
@@ -91,31 +115,34 @@ impl<'a, R: AsyncRead + Unpin> Future for AvailableReader<'a, R> {
91115#[ cfg( test) ]
92116mod tests {
93117 use super :: * ;
118+ use futures:: poll;
94119 use std:: io:: Cursor ;
95120 use std:: time:: Duration ;
121+ use tokio:: stream:: StreamExt ;
122+ use tokio_test:: assert_pending;
96123
97124 #[ tokio:: test]
98125 async fn available_reader_reads_all_available_data_smaller_than_its_buf ( ) {
99126 let data = vec ! [ 42u8 ; 100 ] ;
100127 let mut reader = Cursor :: new ( data. clone ( ) ) ;
101128
102- let available_reader = AvailableReader :: new ( & mut reader) ;
103- let ( read_data, is_finished ) = available_reader. await . unwrap ( ) ;
129+ let mut available_reader = AvailableReader :: new ( & mut reader) ;
130+ let read_data = available_reader. next ( ) . await . unwrap ( ) . unwrap ( ) ;
104131
105132 assert_eq ! ( read_data, data) ;
106- assert ! ( is_finished )
133+ assert ! ( available_reader . next ( ) . await . is_none ( ) )
107134 }
108135
109136 #[ tokio:: test]
110137 async fn available_reader_reads_all_available_data_bigger_than_its_buf ( ) {
111138 let data = vec ! [ 42u8 ; AvailableReader :: <Cursor <Vec <u8 >>>:: BUF_INCREMENT + 100 ] ;
112139 let mut reader = Cursor :: new ( data. clone ( ) ) ;
113140
114- let available_reader = AvailableReader :: new ( & mut reader) ;
115- let ( read_data, is_finished ) = available_reader. await . unwrap ( ) ;
141+ let mut available_reader = AvailableReader :: new ( & mut reader) ;
142+ let read_data = available_reader. next ( ) . await . unwrap ( ) . unwrap ( ) ;
116143
117144 assert_eq ! ( read_data, data) ;
118- assert ! ( is_finished )
145+ assert ! ( available_reader . next ( ) . await . is_none ( ) )
119146 }
120147
121148 #[ tokio:: test]
@@ -129,11 +156,11 @@ mod tests {
129156 . read ( & second_data_chunk)
130157 . build ( ) ;
131158
132- let available_reader = AvailableReader :: new ( & mut reader_mock) ;
133- let ( read_data, is_finished ) = available_reader. await . unwrap ( ) ;
159+ let mut available_reader = AvailableReader :: new ( & mut reader_mock) ;
160+ let read_data = available_reader. next ( ) . await . unwrap ( ) . unwrap ( ) ;
134161
135162 assert_eq ! ( read_data, first_data_chunk) ;
136- assert ! ( !is_finished )
163+ assert_pending ! ( poll! ( available_reader . next ( ) ) ) ;
137164 }
138165
139166 #[ tokio:: test]
@@ -145,10 +172,40 @@ mod tests {
145172 . read ( & data)
146173 . build ( ) ;
147174
148- let available_reader = AvailableReader :: new ( & mut reader_mock) ;
149- let ( read_data, is_finished ) = available_reader. await . unwrap ( ) ;
175+ let mut available_reader = AvailableReader :: new ( & mut reader_mock) ;
176+ let read_data = available_reader. next ( ) . await . unwrap ( ) . unwrap ( ) ;
150177
151178 assert_eq ! ( read_data, data) ;
152- assert ! ( is_finished )
179+ assert ! ( available_reader . next ( ) . await . is_none ( ) )
153180 }
181+
182+ // perhaps the issue of tokio io builder will be resolved in tokio 0.3?
183+ // #[tokio::test]
184+ // async fn available_reader_will_wait_for_more_data_if_its_within_grace_period() {
185+ // let first_data_chunk = vec![42u8; 100];
186+ // let second_data_chunk = vec![123u8; 100];
187+ //
188+ // let combined_chunks: Vec<_> = first_data_chunk
189+ // .iter()
190+ // .cloned()
191+ // .chain(second_data_chunk.iter().cloned())
192+ // .collect();
193+ //
194+ // let mut reader_mock = tokio_test::io::Builder::new()
195+ // .read(&first_data_chunk)
196+ // .wait(Duration::from_millis(2))
197+ // .read(&second_data_chunk)
198+ // .build();
199+ //
200+ // let mut available_reader = AvailableReader {
201+ // buf: RefCell::new(BytesMut::with_capacity(4096)),
202+ // inner: RefCell::new(&mut reader_mock),
203+ // grace_period: Some(delay_for(Duration::from_millis(5))),
204+ // };
205+ //
206+ // let read_data = available_reader.next().await.unwrap().unwrap();
207+ //
208+ // assert_eq!(read_data, combined_chunks);
209+ // assert!(available_reader.next().await.is_none())
210+ // }
154211}
0 commit comments