Skip to content

Commit 28f4272

Browse files
authored
Optimize BufWriter::poll_* method (#407)
- Simplify their code - Only poll when there isn't enough buffer - Only do internal `copy_within` to remove written data when necessary, since `memmove` is expensive
1 parent 52e1043 commit 28f4272

File tree

1 file changed

+25
-18
lines changed

1 file changed

+25
-18
lines changed

crates/async-compression/src/generic/write/buf_writer.rs

Lines changed: 25 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ impl BufWriter {
4747
}
4848

4949
/// Remove the already written data
50-
fn reshuffle_and_remove_written(&mut self) {
50+
fn remove_written(&mut self) {
5151
self.buf.copy_within(self.written..self.buffered, 0);
5252
self.buffered -= self.written;
5353
self.written = 0;
@@ -82,11 +82,7 @@ impl BufWriter {
8282
Ok(())
8383
};
8484

85-
if self.written > 0 {
86-
self.reshuffle_and_remove_written();
87-
88-
Poll::Ready(ret)
89-
} else if self.buffered == 0 {
85+
if self.written > 0 || self.buffered < self.buf.len() {
9086
Poll::Ready(ret)
9187
} else {
9288
ret?;
@@ -99,7 +95,11 @@ impl BufWriter {
9995
poll_write: &mut dyn FnMut(&[u8]) -> Poll<io::Result<usize>>,
10096
) -> Poll<io::Result<()>> {
10197
let ret = ready!(self.do_flush(poll_write));
102-
self.reshuffle_and_remove_written();
98+
99+
debug_assert_eq!(self.buffered, self.written);
100+
self.buffered = 0;
101+
self.written = 0;
102+
103103
Poll::Ready(ret)
104104
}
105105

@@ -108,22 +108,24 @@ impl BufWriter {
108108
buf: &[u8],
109109
poll_write: &mut dyn FnMut(&[u8]) -> Poll<io::Result<usize>>,
110110
) -> Poll<io::Result<usize>> {
111-
if self.buffered + buf.len() > self.buf.len() {
112-
ready!(self.partial_flush_buf(poll_write))?;
113-
}
114-
115111
if buf.len() >= self.buf.len() {
116-
if self.buffered == 0 {
117-
poll_write(buf)
118-
} else {
119-
// The only way that `partial_flush_buf` would have returned with
120-
// `this.buffered != 0` is if it were Pending, so our waker was already queued
121-
Poll::Pending
122-
}
112+
ready!(self.flush_buf(poll_write))?;
113+
poll_write(buf)
114+
} else if (self.buf.len() - self.buffered) >= buf.len() {
115+
self.buf[self.buffered..].copy_from_slice(buf);
116+
self.buffered += buf.len();
117+
118+
Poll::Ready(Ok(buf.len()))
123119
} else {
120+
ready!(self.partial_flush_buf(poll_write))?;
121+
if self.written > 0 {
122+
self.remove_written();
123+
}
124+
124125
let len = buf.len().min(self.buf.len() - self.buffered);
125126
self.buf[self.buffered..self.buffered + len].copy_from_slice(&buf[..len]);
126127
self.buffered += len;
128+
127129
Poll::Ready(Ok(len))
128130
}
129131
}
@@ -133,6 +135,11 @@ impl BufWriter {
133135
poll_write: &mut dyn FnMut(&[u8]) -> Poll<io::Result<usize>>,
134136
) -> Poll<io::Result<&mut [u8]>> {
135137
ready!(self.partial_flush_buf(poll_write))?;
138+
139+
if self.written > 0 {
140+
self.remove_written();
141+
}
142+
136143
Poll::Ready(Ok(&mut self.buf[self.buffered..]))
137144
}
138145

0 commit comments

Comments
 (0)