mirror of https://github.com/fafhrd91/actix-web
125 lines
3.7 KiB
Rust
125 lines
3.7 KiB
Rust
use std::collections::VecDeque;
|
|
|
|
use bytes::{Buf, BufMut, Bytes, BytesMut};
|
|
|
|
// 64KB max capacity (arbitrarily chosen)
|
|
const MAX_CAPACITY: usize = 1024 * 64;
|
|
|
|
pub struct BigBytes {
|
|
buffer: BytesMut,
|
|
frozen: VecDeque<Bytes>,
|
|
frozen_len: usize,
|
|
}
|
|
|
|
impl BigBytes {
|
|
/// Initialize a new BigBytes with the internal buffer set to `capacity` capacity
|
|
pub fn with_capacity(capacity: usize) -> Self {
|
|
Self {
|
|
buffer: BytesMut::with_capacity(capacity),
|
|
frozen: VecDeque::default(),
|
|
frozen_len: 0,
|
|
}
|
|
}
|
|
|
|
/// Clear the internal queue and buffer, resetting length to zero
|
|
///
|
|
/// if the internal buffer capacity exceeds 64KB or new_capacity, whichever is greater, it will
|
|
/// be freed and a new buffer of capacity `new_capacity` will be allocated
|
|
pub fn clear(&mut self, new_capacity: usize) {
|
|
std::mem::take(&mut self.frozen);
|
|
self.frozen_len = 0;
|
|
self.buffer.clear();
|
|
|
|
if self.buffer.capacity() > new_capacity.max(MAX_CAPACITY) {
|
|
self.buffer = BytesMut::with_capacity(new_capacity);
|
|
}
|
|
}
|
|
|
|
/// Return a mutable reference to the underlying buffer. This should only be used when dealing
|
|
/// with small allocations (e.g. writing headers)
|
|
pub fn buffer_mut(&mut self) -> &mut BytesMut {
|
|
&mut self.buffer
|
|
}
|
|
|
|
/// Return the total length of the bytes stored in BigBytes
|
|
pub fn total_len(&mut self) -> usize {
|
|
self.frozen_len + self.buffer.len()
|
|
}
|
|
|
|
/// Return whether there are no bytes present in the BigBytes
|
|
pub fn is_empty(&self) -> bool {
|
|
self.frozen_len == 0 && self.buffer.is_empty()
|
|
}
|
|
|
|
/// Add the `bytes` to the internal structure. If `bytes` exceeds 64KB, it is pushed into a
|
|
/// queue, otherwise, it is added to a buffer.
|
|
pub fn put_bytes(&mut self, bytes: Bytes) {
|
|
if !self.buffer.is_empty() {
|
|
let current = self.buffer.split().freeze();
|
|
self.frozen_len += current.len();
|
|
self.frozen.push_back(current);
|
|
}
|
|
|
|
if !bytes.is_empty() {
|
|
self.frozen_len += bytes.len();
|
|
self.frozen.push_back(bytes);
|
|
}
|
|
}
|
|
|
|
/// Returns a slice of the frontmost buffer
|
|
///
|
|
/// While there are bytes present in BigBytes, front_slice is guaranteed not to return an empty
|
|
/// slice.
|
|
pub fn front_slice(&self) -> &[u8] {
|
|
if let Some(front) = self.frozen.front() {
|
|
front
|
|
} else {
|
|
&self.buffer
|
|
}
|
|
}
|
|
|
|
/// Advances the first buffer by `count` bytes. If the first buffer is advanced to completion,
|
|
/// it is popped from the queue
|
|
pub fn advance(&mut self, count: usize) {
|
|
if let Some(front) = self.frozen.front_mut() {
|
|
front.advance(count);
|
|
|
|
if front.is_empty() {
|
|
self.frozen.pop_front();
|
|
}
|
|
|
|
self.frozen_len -= count;
|
|
} else {
|
|
self.buffer.advance(count);
|
|
}
|
|
}
|
|
|
|
/// Pops the front Bytes from the BigBytes, or splits and freezes the internal buffer if no
|
|
/// Bytes are present.
|
|
pub fn pop_front(&mut self) -> Option<Bytes> {
|
|
if let Some(front) = self.frozen.pop_front() {
|
|
self.frozen_len -= front.len();
|
|
Some(front)
|
|
} else if !self.buffer.is_empty() {
|
|
Some(self.buffer.split().freeze())
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// Drain the BigBytes, writing everything into the provided BytesMut
|
|
pub fn write_to(&mut self, dst: &mut BytesMut) {
|
|
dst.reserve(self.total_len());
|
|
|
|
for buf in &self.frozen {
|
|
dst.put_slice(buf);
|
|
}
|
|
|
|
dst.put_slice(&self.buffer.split());
|
|
|
|
self.frozen_len = 0;
|
|
|
|
std::mem::take(&mut self.frozen);
|
|
}
|
|
}
|