@@ -7,8 +7,9 @@ use bytes::{Buf, BufMut, BytesMut};
7
7
use std:: pin:: Pin ;
8
8
use std:: task:: { Context , Poll } ;
9
9
use tokio:: io:: { AsyncRead , AsyncWrite , ReadBuf } ;
10
+ use tokio_util:: io:: poll_write_buf;
10
11
11
- use std:: io:: { self , Cursor , IoSlice } ;
12
+ use std:: io:: { self , Cursor } ;
12
13
13
14
// A macro to get around a method needing to borrow &mut self
14
15
macro_rules! limited_write_buf {
@@ -45,8 +46,11 @@ struct Encoder<B> {
45
46
/// Max frame size, this is specified by the peer
46
47
max_frame_size : FrameSize ,
47
48
48
- /// Whether or not the wrapped `AsyncWrite` supports vectored IO.
49
- is_write_vectored : bool ,
49
+ /// Chain payloads bigger than this.
50
+ chain_threshold : usize ,
51
+
52
+ /// Min buffer required to attempt to write a frame
53
+ min_buffer_capacity : usize ,
50
54
}
51
55
52
56
#[ derive( Debug ) ]
@@ -61,22 +65,28 @@ enum Next<B> {
61
65
/// frame that big.
62
66
const DEFAULT_BUFFER_CAPACITY : usize = 16 * 1_024 ;
63
67
64
- /// Min buffer required to attempt to write a frame
65
- const MIN_BUFFER_CAPACITY : usize = frame:: HEADER_LEN + CHAIN_THRESHOLD ;
66
-
67
- /// Chain payloads bigger than this. The remote will never advertise a max frame
68
- /// size less than this (well, the spec says the max frame size can't be less
69
- /// than 16kb, so not even close).
68
+ /// Chain payloads bigger than this when vectored I/O is enabled. The remote
69
+ /// will never advertise a max frame size less than this (well, the spec says
70
+ /// the max frame size can't be less than 16kb, so not even close).
70
71
const CHAIN_THRESHOLD : usize = 256 ;
71
72
73
+ /// Chain payloads bigger than this when vectored I/O is **not** enabled.
74
+ /// A larger value in this scenario will reduce the number of small and
75
+ /// fragmented data being sent, and hereby improve the throughput.
76
+ const CHAIN_THRESHOLD_WITHOUT_VECTORED_IO : usize = 1024 ;
77
+
72
78
// TODO: Make generic
73
79
impl < T , B > FramedWrite < T , B >
74
80
where
75
81
T : AsyncWrite + Unpin ,
76
82
B : Buf ,
77
83
{
78
84
pub fn new ( inner : T ) -> FramedWrite < T , B > {
79
- let is_write_vectored = inner. is_write_vectored ( ) ;
85
+ let chain_threshold = if inner. is_write_vectored ( ) {
86
+ CHAIN_THRESHOLD
87
+ } else {
88
+ CHAIN_THRESHOLD_WITHOUT_VECTORED_IO
89
+ } ;
80
90
FramedWrite {
81
91
inner,
82
92
encoder : Encoder {
85
95
next : None ,
86
96
last_data_frame : None ,
87
97
max_frame_size : frame:: DEFAULT_MAX_FRAME_SIZE ,
88
- is_write_vectored,
98
+ chain_threshold,
99
+ min_buffer_capacity : chain_threshold + frame:: HEADER_LEN ,
89
100
} ,
90
101
}
91
102
}
@@ -126,23 +137,17 @@ where
126
137
Some ( Next :: Data ( ref mut frame) ) => {
127
138
tracing:: trace!( queued_data_frame = true ) ;
128
139
let mut buf = ( & mut self . encoder . buf ) . chain ( frame. payload_mut ( ) ) ;
129
- ready ! ( write(
130
- & mut self . inner,
131
- self . encoder. is_write_vectored,
132
- & mut buf,
133
- cx,
134
- ) ) ?
140
+ ready ! ( poll_write_buf( Pin :: new( & mut self . inner) , cx, & mut buf) ) ?
135
141
}
136
142
_ => {
137
143
tracing:: trace!( queued_data_frame = false ) ;
138
- ready ! ( write(
139
- & mut self . inner,
140
- self . encoder. is_write_vectored,
141
- & mut self . encoder. buf,
144
+ ready ! ( poll_write_buf(
145
+ Pin :: new( & mut self . inner) ,
142
146
cx,
147
+ & mut self . encoder. buf
143
148
) ) ?
144
149
}
145
- }
150
+ } ;
146
151
}
147
152
148
153
match self . encoder . unset_frame ( ) {
@@ -165,30 +170,6 @@ where
165
170
}
166
171
}
167
172
168
- fn write < T , B > (
169
- writer : & mut T ,
170
- is_write_vectored : bool ,
171
- buf : & mut B ,
172
- cx : & mut Context < ' _ > ,
173
- ) -> Poll < io:: Result < ( ) > >
174
- where
175
- T : AsyncWrite + Unpin ,
176
- B : Buf ,
177
- {
178
- // TODO(eliza): when tokio-util 0.5.1 is released, this
179
- // could just use `poll_write_buf`...
180
- const MAX_IOVS : usize = 64 ;
181
- let n = if is_write_vectored {
182
- let mut bufs = [ IoSlice :: new ( & [ ] ) ; MAX_IOVS ] ;
183
- let cnt = buf. chunks_vectored ( & mut bufs) ;
184
- ready ! ( Pin :: new( writer) . poll_write_vectored( cx, & bufs[ ..cnt] ) ) ?
185
- } else {
186
- ready ! ( Pin :: new( writer) . poll_write( cx, buf. chunk( ) ) ) ?
187
- } ;
188
- buf. advance ( n) ;
189
- Ok ( ( ) ) . into ( )
190
- }
191
-
192
173
#[ must_use]
193
174
enum ControlFlow {
194
175
Continue ,
@@ -240,12 +221,17 @@ where
240
221
return Err ( PayloadTooBig ) ;
241
222
}
242
223
243
- if len >= CHAIN_THRESHOLD {
224
+ if len >= self . chain_threshold {
244
225
let head = v. head ( ) ;
245
226
246
227
// Encode the frame head to the buffer
247
228
head. encode ( len, self . buf . get_mut ( ) ) ;
248
229
230
+ if self . buf . get_ref ( ) . remaining ( ) < self . chain_threshold {
231
+ let extra_bytes = self . chain_threshold - self . buf . remaining ( ) ;
232
+ self . buf . get_mut ( ) . put ( v. payload_mut ( ) . take ( extra_bytes) ) ;
233
+ }
234
+
249
235
// Save the data frame
250
236
self . next = Some ( Next :: Data ( v) ) ;
251
237
} else {
@@ -305,7 +291,9 @@ where
305
291
}
306
292
307
293
fn has_capacity ( & self ) -> bool {
308
- self . next . is_none ( ) && self . buf . get_ref ( ) . remaining_mut ( ) >= MIN_BUFFER_CAPACITY
294
+ self . next . is_none ( )
295
+ && ( self . buf . get_ref ( ) . capacity ( ) - self . buf . get_ref ( ) . len ( )
296
+ >= self . min_buffer_capacity )
309
297
}
310
298
311
299
fn is_empty ( & self ) -> bool {
0 commit comments