src/backup/chunk_stream.rs: optimize FixedChunkStream (use BytesMut)
This commit is contained in:
parent
3be3f3dcc4
commit
169c0e060f
@ -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!();
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
Loading…
Reference in New Issue
Block a user