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:
Markus Westerlind
2021-02-18 20:17:49 +01:00
committed by GitHub
parent 6357e3256a
commit 30ca832790
7 changed files with 1097 additions and 822 deletions

View File

@@ -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 {