tyler/proxmox-backup
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

src/backup/chunk_stream.rs: optimize FixedChunkStream (use BytesMut)

Browse Source
This commit is contained in:
Dietmar Maurer 2019-05-22 09:05:35 +02:00
parent 3be3f3dcc4
commit 169c0e060f
1 changed files with 27 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

57
src/backup/chunk_stream.rs
View File

@ -73,27 +73,40 @@ impl <S: Stream<Item=Vec<u8>, Error=Error>> Stream for ChunkStream<S> {
} }
} }
use bytes::BytesMut;
/// Split input stream into fixed sized chunks /// Split input stream into fixed sized chunks
pub struct FixedChunkStream<S: Stream<Item=Vec<u8>, Error=Error>> { pub struct FixedChunkStream<S> {
input: S, input: S,
chunk_size: usize, chunk_size: usize,
buffer: Option<Vec<u8>>, buffer: BytesMut,
} }
impl <S: Stream<Item=Vec<u8>, Error=Error>> FixedChunkStream<S> { impl <S> FixedChunkStream<S> {
pub fn new(input: S, chunk_size: usize) -> Self { pub fn new(input: S, chunk_size: usize) -> Self {
Self { input, chunk_size, buffer: None } Self { input, chunk_size, buffer: BytesMut::new() }
} }
} }
impl <S: Stream<Item=Vec<u8>, Error=Error>> Stream for FixedChunkStream<S> { impl <S> Stream for FixedChunkStream<S>
where S: Stream,
S::Item: AsRef<[u8]>,
{
type Item = Vec<u8>; type Item = BytesMut;
type Error = Error; type Error = S::Error;
fn poll(&mut self) -> Poll<Option<Vec<u8>>, Error> { fn poll(&mut self) -> Poll<Option<BytesMut>, S::Error> {
loop { loop {
if self.buffer.len() == self.chunk_size {
return Ok(Async::Ready(Some(self.buffer.take())));
} else if self.buffer.len() > self.chunk_size {
let result = self.buffer.split_to(self.chunk_size);
return Ok(Async::Ready(Some(result)));
}
match self.input.poll() { match self.input.poll() {
Err(err) => { Err(err) => {
return Err(err); return Err(err);
@ -103,30 +116,14 @@ impl <S: Stream<Item=Vec<u8>, Error=Error>> Stream for FixedChunkStream<S> {
} }
Ok(Async::Ready(None)) => { Ok(Async::Ready(None)) => {
// last chunk can have any size // last chunk can have any size
return Ok(Async::Ready(self.buffer.take())); if self.buffer.len() > 0 {
return Ok(Async::Ready(Some(self.buffer.take())));
} else {
return Ok(Async::Ready(None));
}
} }
Ok(Async::Ready(Some(data))) => { Ok(Async::Ready(Some(data))) => {
let buffer = self.buffer.get_or_insert_with(|| Vec::with_capacity(1024*1024)); self.buffer.extend_from_slice(data.as_ref());
let need = self.chunk_size - buffer.len();
if need > data.len() {
buffer.extend(data);
// continue poll
} else if need == data.len() {
buffer.extend(data);
return Ok(Async::Ready(self.buffer.take()));
} else if need < data.len() {
let (left, right) = data.split_at(need);
buffer.extend(left);
let result = self.buffer.take();
self.buffer = Some(Vec::from(right));
return Ok(Async::Ready(result));
} else {
unreachable!();
}
} }
} }
} }