diff --git a/src/common/buf.rs b/src/common/buf.rs index d00071551b..91d9f22250 100644 --- a/src/common/buf.rs +++ b/src/common/buf.rs @@ -5,18 +5,21 @@ use bytes::{Buf, BufMut, Bytes, BytesMut}; pub(crate) struct BufList { bufs: VecDeque, + remaining: usize, } impl BufList { pub(crate) fn new() -> BufList { BufList { bufs: VecDeque::new(), + remaining: 0, } } #[inline] pub(crate) fn push(&mut self, buf: T) { debug_assert!(buf.has_remaining()); + self.remaining += buf.remaining(); self.bufs.push_back(buf); } @@ -29,7 +32,7 @@ impl BufList { impl Buf for BufList { #[inline] fn remaining(&self) -> usize { - self.bufs.iter().map(|buf| buf.remaining()).sum() + self.remaining } #[inline] @@ -39,6 +42,8 @@ impl Buf for BufList { #[inline] fn advance(&mut self, mut cnt: usize) { + assert!(cnt <= self.remaining, "`cnt` greater than remaining"); + self.remaining -= cnt; while cnt > 0 { { let front = &mut self.bufs[0]; @@ -78,12 +83,18 @@ impl Buf for BufList { Some(front) if front.remaining() == len => { let b = front.copy_to_bytes(len); self.bufs.pop_front(); + self.remaining -= len; b } - Some(front) if front.remaining() > len => front.copy_to_bytes(len), + Some(front) if front.remaining() > len => { + self.remaining -= len; + front.copy_to_bytes(len) + } _ => { - assert!(len <= self.remaining(), "`len` greater than remaining"); + assert!(len <= self.remaining, "`len` greater than remaining"); let mut bm = BytesMut::with_capacity(len); + // Note: `bm.put(self.take(len))` calls `self.advance()` internally, + // which already decrements `self.remaining`. bm.put(self.take(len)); bm.freeze() } @@ -100,6 +111,7 @@ mod tests { fn hello_world_buf() -> BufList { BufList { bufs: vec![Bytes::from("Hello"), Bytes::from(" "), Bytes::from("World")].into(), + remaining: 11, } }