diff options
Diffstat (limited to 'src/cmd/record.rs')
-rw-r--r-- | src/cmd/record.rs | 36 |
1 files changed, 19 insertions, 17 deletions
diff --git a/src/cmd/record.rs b/src/cmd/record.rs index 6fda88f..46e76e1 100644 --- a/src/cmd/record.rs +++ b/src/cmd/record.rs @@ -62,7 +62,7 @@ enum FileState { fut: tokio::fs::file::CreateFuture<String>, }, Open { - file: crate::ttyrec::File, + writer: crate::ttyrec::Writer<tokio::fs::File>, }, } @@ -141,9 +141,11 @@ impl RecordSession { filename: filename.clone(), } })); - let mut file = crate::ttyrec::File::new(file); - file.write_frame(self.buffer.contents())?; - self.file = FileState::Open { file }; + let mut writer = crate::ttyrec::Writer::new(file); + if !self.buffer.contents().is_empty() { + writer.frame(self.buffer.contents())?; + } + self.file = FileState::Open { writer }; Ok(component_future::Async::DidWork) } FileState::Open { .. } => { @@ -173,8 +175,8 @@ impl RecordSession { } tokio_pty_process_stream::Event::Output { data } => { self.record_bytes(&data); - if let FileState::Open { file } = &mut self.file { - file.write_frame(&data)?; + if let FileState::Open { writer } = &mut self.file { + writer.frame(&data)?; } } } @@ -222,23 +224,23 @@ impl RecordSession { } fn poll_write_file(&mut self) -> component_future::Poll<(), Error> { - let file = match &mut self.file { - FileState::Open { file } => file, + let writer = match &mut self.file { + FileState::Open { writer } => writer, _ => { return Ok(component_future::Async::NothingToDo); } }; - match file.poll_write()? { - futures::Async::Ready(()) => Ok(component_future::Async::DidWork), - futures::Async::NotReady => { - // ship all data to the server before actually ending - if self.done && file.is_empty() { - Ok(component_future::Async::Ready(())) - } else { - Ok(component_future::Async::NotReady) - } + if writer.is_empty() { + // finish writing to the file before actually ending + if self.done { + Ok(component_future::Async::Ready(())) + } else { + Ok(component_future::Async::NothingToDo) } + } else { + component_future::try_ready!(writer.poll_write()); + Ok(component_future::Async::DidWork) } } } |