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

src/backup/chunk_stream.rs: fix behaviour for large input buffers

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

57
src/backup/chunk_stream.rs
View File

@ -9,13 +9,13 @@ pub struct ChunkStream<S: Stream<Item=Vec<u8>, Error=Error>> {
input: S, input: S,
chunker: Chunker, chunker: Chunker,
buffer: Option<Vec<u8>>, buffer: Option<Vec<u8>>,
rest: Option<Vec<u8>>, scan: Option<Vec<u8>>,
} }
impl <S: Stream<Item=Vec<u8>, Error=Error>> ChunkStream<S> { impl <S: Stream<Item=Vec<u8>, Error=Error>> ChunkStream<S> {
pub fn new(input: S) -> Self { pub fn new(input: S) -> Self {
Self { input, chunker: Chunker::new(4 * 1024 * 1024), buffer: None, rest: None } Self { input, chunker: Chunker::new(4 * 1024 * 1024), buffer: None, scan: None}
} }
} }
@ -26,6 +26,27 @@ impl <S: Stream<Item=Vec<u8>, Error=Error>> Stream for ChunkStream<S> {
fn poll(&mut self) -> Poll<Option<Vec<u8>>, Error> { fn poll(&mut self) -> Poll<Option<Vec<u8>>, Error> {
loop { loop {
if let Some(data) = self.scan.take() {
let buffer = self.buffer.get_or_insert_with(|| Vec::with_capacity(1024*1024));
let boundary = self.chunker.scan(&data);
if boundary == 0 {
buffer.extend(data);
// continue poll
} else if boundary == data.len() {
buffer.extend(data);
return Ok(Async::Ready(self.buffer.take()));
} else if boundary < data.len() {
let (left, right) = data.split_at(boundary);
buffer.extend(left);
self.scan = Some(right.to_vec());
return Ok(Async::Ready(self.buffer.take()));
} else {
panic!("got unexpected chunk boundary from chunker");
}
}
match self.input.poll() { match self.input.poll() {
Err(err) => { Err(err) => {
return Err(err); return Err(err);
@ -35,7 +56,7 @@ impl <S: Stream<Item=Vec<u8>, Error=Error>> Stream for ChunkStream<S> {
} }
Ok(Async::Ready(None)) => { Ok(Async::Ready(None)) => {
let mut data = self.buffer.take().or_else(|| Some(vec![])).unwrap(); let mut data = self.buffer.take().or_else(|| Some(vec![])).unwrap();
if let Some(rest) = self.rest.take() { data.extend(rest); } if let Some(rest) = self.scan.take() { data.extend(rest); }
if data.len() > 0 { if data.len() > 0 {
return Ok(Async::Ready(Some(data))); return Ok(Async::Ready(Some(data)));
@ -44,34 +65,8 @@ impl <S: Stream<Item=Vec<u8>, Error=Error>> Stream for ChunkStream<S> {
} }
} }
Ok(Async::Ready(Some(data))) => { Ok(Async::Ready(Some(data))) => {
let scan = self.scan.get_or_insert_with(|| Vec::with_capacity(1024*1024));
let data = if let Some(mut rest) = self.rest.take() { scan.extend(data);
rest.extend(data);
rest
} else {
data
};
let buffer = self.buffer.get_or_insert_with(|| Vec::with_capacity(1024*1024));
let boundary = self.chunker.scan(&data);
if boundary == 0 {
buffer.extend(data);
// continue poll
} else if boundary == data.len() {
buffer.extend(data);
return Ok(Async::Ready(self.buffer.take()));
} else if boundary < data.len() {
let (left, right) = data.split_at(boundary);
buffer.extend(left);
let rest = self.rest.get_or_insert_with(|| Vec::with_capacity(right.len()));
rest.extend(right);
return Ok(Async::Ready(self.buffer.take()));
} else {
panic!("got unexpected chunk boundary from chunker");
}
} }
} }
} }