Make some functions less-generic to reduce binary bloat (#503)
* refactor: Extract FramedWrite::buffer to a less generic function Should cut out another 23 KiB (since I see it duplicated) * refactor: Extract some duplicated code to a function * refactor: Extract part of flush into a less generic function * refactor: Extract a less generic part of connection * refactor: Factor out a less generic part of Connection::poll2 * refactor: Extract a non-generic part of handshake2 * refactor: Don't duplicate Streams code on Peer (-3.5%) The `P: Peer` parameter is rarely used and there is already a mechanism for using it dynamically. * refactor: Make recv_frame less generic (-2.3%) * Move out part of Connection::poll * refactor: Extract parts of Connection * refactor: Extract a non-generic part of reclaim_frame * comments
This commit is contained in:
committed by
GitHub
parent
6357e3256a
commit
30ca832790
@@ -23,6 +23,11 @@ pub struct FramedWrite<T, B> {
|
||||
/// Upstream `AsyncWrite`
|
||||
inner: T,
|
||||
|
||||
encoder: Encoder<B>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Encoder<B> {
|
||||
/// HPACK encoder
|
||||
hpack: hpack::Encoder,
|
||||
|
||||
@@ -74,12 +79,14 @@ where
|
||||
let is_write_vectored = inner.is_write_vectored();
|
||||
FramedWrite {
|
||||
inner,
|
||||
hpack: hpack::Encoder::default(),
|
||||
buf: Cursor::new(BytesMut::with_capacity(DEFAULT_BUFFER_CAPACITY)),
|
||||
next: None,
|
||||
last_data_frame: None,
|
||||
max_frame_size: frame::DEFAULT_MAX_FRAME_SIZE,
|
||||
is_write_vectored,
|
||||
encoder: Encoder {
|
||||
hpack: hpack::Encoder::default(),
|
||||
buf: Cursor::new(BytesMut::with_capacity(DEFAULT_BUFFER_CAPACITY)),
|
||||
next: None,
|
||||
last_data_frame: None,
|
||||
max_frame_size: frame::DEFAULT_MAX_FRAME_SIZE,
|
||||
is_write_vectored,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -88,11 +95,11 @@ where
|
||||
/// Calling this function may result in the current contents of the buffer
|
||||
/// to be flushed to `T`.
|
||||
pub fn poll_ready(&mut self, cx: &mut Context) -> Poll<io::Result<()>> {
|
||||
if !self.has_capacity() {
|
||||
if !self.encoder.has_capacity() {
|
||||
// Try flushing
|
||||
ready!(self.flush(cx))?;
|
||||
|
||||
if !self.has_capacity() {
|
||||
if !self.encoder.has_capacity() {
|
||||
return Poll::Pending;
|
||||
}
|
||||
}
|
||||
@@ -105,6 +112,128 @@ where
|
||||
/// `poll_ready` must be called first to ensure that a frame may be
|
||||
/// accepted.
|
||||
pub fn buffer(&mut self, item: Frame<B>) -> Result<(), UserError> {
|
||||
self.encoder.buffer(item)
|
||||
}
|
||||
|
||||
/// Flush buffered data to the wire
|
||||
pub fn flush(&mut self, cx: &mut Context) -> Poll<io::Result<()>> {
|
||||
let span = tracing::trace_span!("FramedWrite::flush");
|
||||
let _e = span.enter();
|
||||
|
||||
loop {
|
||||
while !self.encoder.is_empty() {
|
||||
match self.encoder.next {
|
||||
Some(Next::Data(ref mut frame)) => {
|
||||
tracing::trace!(queued_data_frame = true);
|
||||
let mut buf = (&mut self.encoder.buf).chain(frame.payload_mut());
|
||||
ready!(write(
|
||||
&mut self.inner,
|
||||
self.encoder.is_write_vectored,
|
||||
&mut buf,
|
||||
cx,
|
||||
))?
|
||||
}
|
||||
_ => {
|
||||
tracing::trace!(queued_data_frame = false);
|
||||
ready!(write(
|
||||
&mut self.inner,
|
||||
self.encoder.is_write_vectored,
|
||||
&mut self.encoder.buf,
|
||||
cx,
|
||||
))?
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
match self.encoder.unset_frame() {
|
||||
ControlFlow::Continue => (),
|
||||
ControlFlow::Break => break,
|
||||
}
|
||||
}
|
||||
|
||||
tracing::trace!("flushing buffer");
|
||||
// Flush the upstream
|
||||
ready!(Pin::new(&mut self.inner).poll_flush(cx))?;
|
||||
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
/// Close the codec
|
||||
pub fn shutdown(&mut self, cx: &mut Context) -> Poll<io::Result<()>> {
|
||||
ready!(self.flush(cx))?;
|
||||
Pin::new(&mut self.inner).poll_shutdown(cx)
|
||||
}
|
||||
}
|
||||
|
||||
fn write<T, B>(
|
||||
writer: &mut T,
|
||||
is_write_vectored: bool,
|
||||
buf: &mut B,
|
||||
cx: &mut Context<'_>,
|
||||
) -> Poll<io::Result<()>>
|
||||
where
|
||||
T: AsyncWrite + Unpin,
|
||||
B: Buf,
|
||||
{
|
||||
// TODO(eliza): when tokio-util 0.5.1 is released, this
|
||||
// could just use `poll_write_buf`...
|
||||
const MAX_IOVS: usize = 64;
|
||||
let n = if is_write_vectored {
|
||||
let mut bufs = [IoSlice::new(&[]); MAX_IOVS];
|
||||
let cnt = buf.chunks_vectored(&mut bufs);
|
||||
ready!(Pin::new(writer).poll_write_vectored(cx, &bufs[..cnt]))?
|
||||
} else {
|
||||
ready!(Pin::new(writer).poll_write(cx, buf.chunk()))?
|
||||
};
|
||||
buf.advance(n);
|
||||
Ok(()).into()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
enum ControlFlow {
|
||||
Continue,
|
||||
Break,
|
||||
}
|
||||
|
||||
impl<B> Encoder<B>
|
||||
where
|
||||
B: Buf,
|
||||
{
|
||||
fn unset_frame(&mut self) -> ControlFlow {
|
||||
// Clear internal buffer
|
||||
self.buf.set_position(0);
|
||||
self.buf.get_mut().clear();
|
||||
|
||||
// The data frame has been written, so unset it
|
||||
match self.next.take() {
|
||||
Some(Next::Data(frame)) => {
|
||||
self.last_data_frame = Some(frame);
|
||||
debug_assert!(self.is_empty());
|
||||
ControlFlow::Break
|
||||
}
|
||||
Some(Next::Continuation(frame)) => {
|
||||
// Buffer the continuation frame, then try to write again
|
||||
let mut buf = limited_write_buf!(self);
|
||||
if let Some(continuation) = frame.encode(&mut self.hpack, &mut buf) {
|
||||
// We previously had a CONTINUATION, and after encoding
|
||||
// it, we got *another* one? Let's just double check
|
||||
// that at least some progress is being made...
|
||||
if self.buf.get_ref().len() == frame::HEADER_LEN {
|
||||
// If *only* the CONTINUATION frame header was
|
||||
// written, and *no* header fields, we're stuck
|
||||
// in a loop...
|
||||
panic!("CONTINUATION frame write loop; header value too big to encode");
|
||||
}
|
||||
|
||||
self.next = Some(Next::Continuation(continuation));
|
||||
}
|
||||
ControlFlow::Continue
|
||||
}
|
||||
None => ControlFlow::Break,
|
||||
}
|
||||
}
|
||||
|
||||
fn buffer(&mut self, item: Frame<B>) -> Result<(), UserError> {
|
||||
// Ensure that we have enough capacity to accept the write.
|
||||
assert!(self.has_capacity());
|
||||
let span = tracing::trace_span!("FramedWrite::buffer", frame = ?item);
|
||||
@@ -185,93 +314,6 @@ where
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Flush buffered data to the wire
|
||||
pub fn flush(&mut self, cx: &mut Context) -> Poll<io::Result<()>> {
|
||||
const MAX_IOVS: usize = 64;
|
||||
|
||||
let span = tracing::trace_span!("FramedWrite::flush");
|
||||
let _e = span.enter();
|
||||
|
||||
loop {
|
||||
while !self.is_empty() {
|
||||
match self.next {
|
||||
Some(Next::Data(ref mut frame)) => {
|
||||
tracing::trace!(queued_data_frame = true);
|
||||
let mut buf = (&mut self.buf).chain(frame.payload_mut());
|
||||
// TODO(eliza): when tokio-util 0.5.1 is released, this
|
||||
// could just use `poll_write_buf`...
|
||||
let n = if self.is_write_vectored {
|
||||
let mut bufs = [IoSlice::new(&[]); MAX_IOVS];
|
||||
let cnt = buf.chunks_vectored(&mut bufs);
|
||||
ready!(Pin::new(&mut self.inner).poll_write_vectored(cx, &bufs[..cnt]))?
|
||||
} else {
|
||||
ready!(Pin::new(&mut self.inner).poll_write(cx, buf.chunk()))?
|
||||
};
|
||||
buf.advance(n);
|
||||
}
|
||||
_ => {
|
||||
tracing::trace!(queued_data_frame = false);
|
||||
let n = if self.is_write_vectored {
|
||||
let mut iovs = [IoSlice::new(&[]); MAX_IOVS];
|
||||
let cnt = self.buf.chunks_vectored(&mut iovs);
|
||||
ready!(
|
||||
Pin::new(&mut self.inner).poll_write_vectored(cx, &mut iovs[..cnt])
|
||||
)?
|
||||
} else {
|
||||
ready!(Pin::new(&mut self.inner).poll_write(cx, &mut self.buf.chunk()))?
|
||||
};
|
||||
self.buf.advance(n);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Clear internal buffer
|
||||
self.buf.set_position(0);
|
||||
self.buf.get_mut().clear();
|
||||
|
||||
// The data frame has been written, so unset it
|
||||
match self.next.take() {
|
||||
Some(Next::Data(frame)) => {
|
||||
self.last_data_frame = Some(frame);
|
||||
debug_assert!(self.is_empty());
|
||||
break;
|
||||
}
|
||||
Some(Next::Continuation(frame)) => {
|
||||
// Buffer the continuation frame, then try to write again
|
||||
let mut buf = limited_write_buf!(self);
|
||||
if let Some(continuation) = frame.encode(&mut self.hpack, &mut buf) {
|
||||
// We previously had a CONTINUATION, and after encoding
|
||||
// it, we got *another* one? Let's just double check
|
||||
// that at least some progress is being made...
|
||||
if self.buf.get_ref().len() == frame::HEADER_LEN {
|
||||
// If *only* the CONTINUATION frame header was
|
||||
// written, and *no* header fields, we're stuck
|
||||
// in a loop...
|
||||
panic!("CONTINUATION frame write loop; header value too big to encode");
|
||||
}
|
||||
|
||||
self.next = Some(Next::Continuation(continuation));
|
||||
}
|
||||
}
|
||||
None => {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tracing::trace!("flushing buffer");
|
||||
// Flush the upstream
|
||||
ready!(Pin::new(&mut self.inner).poll_flush(cx))?;
|
||||
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
/// Close the codec
|
||||
pub fn shutdown(&mut self, cx: &mut Context) -> Poll<io::Result<()>> {
|
||||
ready!(self.flush(cx))?;
|
||||
Pin::new(&mut self.inner).poll_shutdown(cx)
|
||||
}
|
||||
|
||||
fn has_capacity(&self) -> bool {
|
||||
self.next.is_none() && self.buf.get_ref().remaining_mut() >= MIN_BUFFER_CAPACITY
|
||||
}
|
||||
@@ -284,26 +326,32 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
impl<B> Encoder<B> {
|
||||
fn max_frame_size(&self) -> usize {
|
||||
self.max_frame_size as usize
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, B> FramedWrite<T, B> {
|
||||
/// Returns the max frame size that can be sent
|
||||
pub fn max_frame_size(&self) -> usize {
|
||||
self.max_frame_size as usize
|
||||
self.encoder.max_frame_size()
|
||||
}
|
||||
|
||||
/// Set the peer's max frame size.
|
||||
pub fn set_max_frame_size(&mut self, val: usize) {
|
||||
assert!(val <= frame::MAX_MAX_FRAME_SIZE as usize);
|
||||
self.max_frame_size = val as FrameSize;
|
||||
self.encoder.max_frame_size = val as FrameSize;
|
||||
}
|
||||
|
||||
/// Set the peer's header table size.
|
||||
pub fn set_header_table_size(&mut self, val: usize) {
|
||||
self.hpack.update_max_size(val);
|
||||
self.encoder.hpack.update_max_size(val);
|
||||
}
|
||||
|
||||
/// Retrieve the last data frame that has been sent
|
||||
pub fn take_last_data_frame(&mut self) -> Option<frame::Data<B>> {
|
||||
self.last_data_frame.take()
|
||||
self.encoder.last_data_frame.take()
|
||||
}
|
||||
|
||||
pub fn get_mut(&mut self) -> &mut T {
|
||||
|
||||
Reference in New Issue
Block a user