1
- use std:: cell:: Cell ;
2
1
use std:: cmp;
3
2
use std:: fmt;
4
3
use std:: io:: { self , IoSlice } ;
@@ -57,13 +56,14 @@ where
57
56
B : Buf ,
58
57
{
59
58
pub fn new ( io : T ) -> Buffered < T , B > {
59
+ let write_buf = WriteBuf :: new ( & io) ;
60
60
Buffered {
61
61
flush_pipeline : false ,
62
62
io,
63
63
read_blocked : false ,
64
64
read_buf : BytesMut :: with_capacity ( 0 ) ,
65
65
read_buf_strategy : ReadStrategy :: default ( ) ,
66
- write_buf : WriteBuf :: new ( ) ,
66
+ write_buf,
67
67
}
68
68
}
69
69
98
98
self . write_buf . set_strategy ( WriteStrategy :: Flatten ) ;
99
99
}
100
100
101
- pub fn set_write_strategy_queue ( & mut self ) {
102
- // this should always be called only at construction time,
103
- // so this assert is here to catch myself
104
- debug_assert ! ( self . write_buf. queue. bufs_cnt( ) == 0 ) ;
105
- self . write_buf . set_strategy ( WriteStrategy :: Queue ) ;
106
- }
107
-
108
101
pub fn read_buf ( & self ) -> & [ u8 ] {
109
102
self . read_buf . as_ref ( )
110
103
}
@@ -237,13 +230,13 @@ where
237
230
if let WriteStrategy :: Flatten = self . write_buf . strategy {
238
231
return self . poll_flush_flattened ( cx) ;
239
232
}
233
+
240
234
loop {
241
- // TODO(eliza): this basically ignores all of `WriteBuf`...put
242
- // back vectored IO and `poll_write_buf` when the appropriate Tokio
243
- // changes land...
244
- let n = ready ! ( Pin :: new( & mut self . io)
245
- // .poll_write_buf(cx, &mut self.write_buf.auto()))?;
246
- . poll_write( cx, self . write_buf. auto( ) . bytes( ) ) ) ?;
235
+ let n = {
236
+ let mut iovs = [ IoSlice :: new ( & [ ] ) ; crate :: common:: io:: MAX_WRITEV_BUFS ] ;
237
+ let len = self . write_buf . bytes_vectored ( & mut iovs) ;
238
+ ready ! ( Pin :: new( & mut self . io) . poll_write_vectored( cx, & iovs[ ..len] ) ) ?
239
+ } ;
247
240
// TODO(eliza): we have to do this manually because
248
241
// `poll_write_buf` doesn't exist in Tokio 0.3 yet...when
249
242
// `poll_write_buf` comes back, the manual advance will need to leave!
@@ -462,12 +455,17 @@ pub(super) struct WriteBuf<B> {
462
455
}
463
456
464
457
impl < B : Buf > WriteBuf < B > {
465
- fn new ( ) -> WriteBuf < B > {
458
+ fn new ( io : & impl AsyncWrite ) -> WriteBuf < B > {
459
+ let strategy = if io. is_write_vectored ( ) {
460
+ WriteStrategy :: Queue
461
+ } else {
462
+ WriteStrategy :: Flatten
463
+ } ;
466
464
WriteBuf {
467
465
headers : Cursor :: new ( Vec :: with_capacity ( INIT_BUFFER_SIZE ) ) ,
468
466
max_buf_size : DEFAULT_MAX_BUFFER_SIZE ,
469
467
queue : BufList :: new ( ) ,
470
- strategy : WriteStrategy :: Auto ,
468
+ strategy,
471
469
}
472
470
}
473
471
}
@@ -480,12 +478,6 @@ where
480
478
self . strategy = strategy;
481
479
}
482
480
483
- // TODO(eliza): put back writev!
484
- #[ inline]
485
- fn auto ( & mut self ) -> WriteBufAuto < ' _ , B > {
486
- WriteBufAuto :: new ( self )
487
- }
488
-
489
481
pub ( super ) fn buffer < BB : Buf + Into < B > > ( & mut self , mut buf : BB ) {
490
482
debug_assert ! ( buf. has_remaining( ) ) ;
491
483
match self . strategy {
@@ -505,7 +497,7 @@ where
505
497
buf. advance ( adv) ;
506
498
}
507
499
}
508
- WriteStrategy :: Auto | WriteStrategy :: Queue => {
500
+ WriteStrategy :: Queue => {
509
501
self . queue . push ( buf. into ( ) ) ;
510
502
}
511
503
}
@@ -514,7 +506,7 @@ where
514
506
fn can_buffer ( & self ) -> bool {
515
507
match self . strategy {
516
508
WriteStrategy :: Flatten => self . remaining ( ) < self . max_buf_size ,
517
- WriteStrategy :: Auto | WriteStrategy :: Queue => {
509
+ WriteStrategy :: Queue => {
518
510
self . queue . bufs_cnt ( ) < MAX_BUF_LIST_BUFFERS && self . remaining ( ) < self . max_buf_size
519
511
}
520
512
}
@@ -573,65 +565,8 @@ impl<B: Buf> Buf for WriteBuf<B> {
573
565
}
574
566
}
575
567
576
- /// Detects when wrapped `WriteBuf` is used for vectored IO, and
577
- /// adjusts the `WriteBuf` strategy if not.
578
- struct WriteBufAuto < ' a , B : Buf > {
579
- bytes_called : Cell < bool > ,
580
- bytes_vec_called : Cell < bool > ,
581
- inner : & ' a mut WriteBuf < B > ,
582
- }
583
-
584
- impl < ' a , B : Buf > WriteBufAuto < ' a , B > {
585
- fn new ( inner : & ' a mut WriteBuf < B > ) -> WriteBufAuto < ' a , B > {
586
- WriteBufAuto {
587
- bytes_called : Cell :: new ( false ) ,
588
- bytes_vec_called : Cell :: new ( false ) ,
589
- inner,
590
- }
591
- }
592
- }
593
-
594
- impl < ' a , B : Buf > Buf for WriteBufAuto < ' a , B > {
595
- #[ inline]
596
- fn remaining ( & self ) -> usize {
597
- self . inner . remaining ( )
598
- }
599
-
600
- #[ inline]
601
- fn bytes ( & self ) -> & [ u8 ] {
602
- self . bytes_called . set ( true ) ;
603
- self . inner . bytes ( )
604
- }
605
-
606
- #[ inline]
607
- fn advance ( & mut self , cnt : usize ) {
608
- self . inner . advance ( cnt)
609
- }
610
-
611
- #[ inline]
612
- fn bytes_vectored < ' t > ( & ' t self , dst : & mut [ IoSlice < ' t > ] ) -> usize {
613
- self . bytes_vec_called . set ( true ) ;
614
- self . inner . bytes_vectored ( dst)
615
- }
616
- }
617
-
618
- impl < ' a , B : Buf + ' a > Drop for WriteBufAuto < ' a , B > {
619
- fn drop ( & mut self ) {
620
- if let WriteStrategy :: Auto = self . inner . strategy {
621
- if self . bytes_vec_called . get ( ) {
622
- self . inner . strategy = WriteStrategy :: Queue ;
623
- } else if self . bytes_called . get ( ) {
624
- trace ! ( "detected no usage of vectored write, flattening" ) ;
625
- self . inner . strategy = WriteStrategy :: Flatten ;
626
- self . inner . headers . bytes . put ( & mut self . inner . queue ) ;
627
- }
628
- }
629
- }
630
- }
631
-
632
568
#[ derive( Debug ) ]
633
569
enum WriteStrategy {
634
- Auto ,
635
570
Flatten ,
636
571
Queue ,
637
572
}
@@ -643,8 +578,8 @@ mod tests {
643
578
644
579
use tokio_test:: io:: Builder as Mock ;
645
580
646
- #[ cfg( feature = "nightly" ) ]
647
- use test:: Bencher ;
581
+ // #[cfg(feature = "nightly")]
582
+ // use test::Bencher;
648
583
649
584
/*
650
585
impl<T: Read> MemRead for AsyncIo<T> {
@@ -873,33 +808,6 @@ mod tests {
873
808
buffered. flush ( ) . await . expect ( "flush" ) ;
874
809
}
875
810
876
- #[ tokio:: test]
877
- async fn write_buf_auto_flatten ( ) {
878
- let _ = pretty_env_logger:: try_init ( ) ;
879
-
880
- let mock = Mock :: new ( )
881
- // Expects write_buf to only consume first buffer
882
- . write ( b"hello " )
883
- // And then the Auto strategy will have flattened
884
- . write ( b"world, it's hyper!" )
885
- . build ( ) ;
886
-
887
- let mut buffered = Buffered :: < _ , Cursor < Vec < u8 > > > :: new ( mock) ;
888
-
889
- // we have 4 buffers, but hope to detect that vectored IO isn't
890
- // being used, and switch to flattening automatically,
891
- // resulting in only 2 writes
892
- buffered. headers_buf ( ) . extend ( b"hello " ) ;
893
- buffered. buffer ( Cursor :: new ( b"world, " . to_vec ( ) ) ) ;
894
- buffered. buffer ( Cursor :: new ( b"it's " . to_vec ( ) ) ) ;
895
- buffered. buffer ( Cursor :: new ( b"hyper!" . to_vec ( ) ) ) ;
896
- assert_eq ! ( buffered. write_buf. queue. bufs_cnt( ) , 3 ) ;
897
-
898
- buffered. flush ( ) . await . expect ( "flush" ) ;
899
-
900
- assert_eq ! ( buffered. write_buf. queue. bufs_cnt( ) , 0 ) ;
901
- }
902
-
903
811
#[ tokio:: test]
904
812
async fn write_buf_queue_disable_auto ( ) {
905
813
let _ = pretty_env_logger:: try_init ( ) ;
@@ -928,19 +836,19 @@ mod tests {
928
836
assert_eq ! ( buffered. write_buf. queue. bufs_cnt( ) , 0 ) ;
929
837
}
930
838
931
- #[ cfg( feature = "nightly" ) ]
932
- #[ bench]
933
- fn bench_write_buf_flatten_buffer_chunk ( b : & mut Bencher ) {
934
- let s = "Hello, World!" ;
935
- b. bytes = s. len ( ) as u64 ;
936
-
937
- let mut write_buf = WriteBuf :: < bytes:: Bytes > :: new ( ) ;
938
- write_buf. set_strategy ( WriteStrategy :: Flatten ) ;
939
- b. iter ( || {
940
- let chunk = bytes:: Bytes :: from ( s) ;
941
- write_buf. buffer ( chunk) ;
942
- :: test:: black_box ( & write_buf) ;
943
- write_buf. headers . bytes . clear ( ) ;
944
- } )
945
- }
839
+ // #[cfg(feature = "nightly")]
840
+ // #[bench]
841
+ // fn bench_write_buf_flatten_buffer_chunk(b: &mut Bencher) {
842
+ // let s = "Hello, World!";
843
+ // b.bytes = s.len() as u64;
844
+
845
+ // let mut write_buf = WriteBuf::<bytes::Bytes>::new();
846
+ // write_buf.set_strategy(WriteStrategy::Flatten);
847
+ // b.iter(|| {
848
+ // let chunk = bytes::Bytes::from(s);
849
+ // write_buf.buffer(chunk);
850
+ // ::test::black_box(&write_buf);
851
+ // write_buf.headers.bytes.clear();
852
+ // })
853
+ // }
946
854
}
0 commit comments