| use std::{ |
| io::{Error, ErrorKind, Read, Result, Seek, SeekFrom, Write}, |
| vec::Vec, |
| }; |
| |
| struct Op { |
| offset: u64, |
| data: Vec<u8>, |
| } |
| |
| impl Op { |
| fn end(&self) -> u64 { |
| self.offset + self.data.len() as u64 |
| } |
| } |
| |
| // TransactionManager provides the ability to queue up changes and then later commit or revert those |
| // changes. It simplifies the process of reverting all changes. It is not designed to work with very |
| // large numbers of changes. |
| pub struct TransactionManager<T> { |
| inner: T, |
| ops: Vec<Op>, |
| active: bool, |
| } |
| |
| impl<T: Seek + Write> TransactionManager<T> { |
| pub fn new(inner: T) -> Self { |
| TransactionManager { inner, ops: Vec::new(), active: false } |
| } |
| |
| #[allow(dead_code)] |
| pub fn into_inner(self) -> T { |
| self.inner |
| } |
| |
| #[must_use] |
| pub fn begin_transaction(&mut self) -> bool { |
| if self.active { |
| false |
| } else { |
| self.active = true; |
| true |
| } |
| } |
| |
| pub fn commit(&mut self) -> Result<()> { |
| assert!(self.active); |
| self.active = false; |
| for op in self.ops.drain(..) { |
| self.inner.seek(SeekFrom::Start(op.offset))?; |
| let mut buf = op.data.as_slice(); |
| while !buf.is_empty() { |
| let done = self.inner.write(buf)?; |
| if done == 0 { |
| return Err(Error::new(ErrorKind::WriteZero, "Inner write failed")); |
| } |
| buf = &buf[done..]; |
| } |
| } |
| Ok(()) |
| } |
| |
| pub fn revert(&mut self) { |
| assert!(self.active); |
| self.active = false; |
| self.ops.clear(); |
| } |
| |
| pub fn borrow_inner(&self) -> &T { |
| &self.inner |
| } |
| } |
| |
| impl<T: Read + Seek> Read for TransactionManager<T> { |
| fn read(&mut self, mut buf: &mut [u8]) -> Result<usize> { |
| if !self.active { |
| return self.inner.read(buf); |
| } |
| |
| let mut offset = self.inner.seek(SeekFrom::Current(0))?; |
| let mut i = 0; |
| while i < self.ops.len() && self.ops[i].end() <= offset { |
| i += 1; |
| } |
| let mut done = 0; |
| while !buf.is_empty() { |
| let mut to_do; |
| let next_offset = if i < self.ops.len() { self.ops[i].offset } else { u64::MAX }; |
| if next_offset <= offset { |
| // Copy from our copy. |
| to_do = std::cmp::min(self.ops[i].end() - offset, buf.len() as u64) as usize; |
| let data_offset = (offset - next_offset) as usize; |
| buf[..to_do].copy_from_slice(&self.ops[i].data[data_offset..data_offset + to_do]); |
| i += 1; |
| self.inner.seek(SeekFrom::Current(to_do as i64))?; |
| } else { |
| // Fall through to inner. |
| to_do = std::cmp::min(next_offset - offset, buf.len() as u64) as usize; |
| to_do = self.inner.read(&mut buf[..to_do])?; |
| if to_do == 0 { |
| return Ok(done); |
| } |
| } |
| buf = &mut buf[to_do..]; |
| offset += to_do as u64; |
| done += to_do; |
| } |
| return Ok(done); |
| } |
| } |
| |
| impl<T: Seek + Write> Write for TransactionManager<T> { |
| fn write(&mut self, buf: &[u8]) -> Result<usize> { |
| if !self.active { |
| return self.inner.write(buf); |
| } |
| |
| let offset = self.inner.seek(SeekFrom::Current(0))?; |
| let mut i = 0; |
| while i < self.ops.len() && self.ops[i].end() < offset { |
| i += 1; |
| } |
| if i >= self.ops.len() { |
| // End of the list, just append a new operation. |
| self.ops.push(Op { offset, data: buf.to_vec() }); |
| } else if self.ops[i].end() == offset { |
| // Extend an existing entry. |
| self.ops[i].data.extend_from_slice(buf); |
| } else if self.ops[i].offset < offset { |
| // Existing entry overlaps. |
| let data_offset = (offset - self.ops[i].offset) as usize; |
| if self.ops[i].end() >= offset + buf.len() as u64 { |
| // Existing entry already encompasses all of this write. |
| self.ops[i].data[data_offset..data_offset + buf.len()].copy_from_slice(buf); |
| } else { |
| // Partial overlap + extension. |
| let to_do = self.ops[i].data.len() - data_offset; |
| self.ops[i].data[data_offset..].copy_from_slice(&buf[..to_do]); |
| self.ops[i].data.extend_from_slice(&buf[to_do..]); |
| } |
| } else { |
| // No overlap at the beginning. |
| self.ops.insert(i, Op { offset: offset, data: buf.to_vec() }); |
| } |
| // Trim any of the following extents as required. |
| let end = self.ops[i].end(); |
| i += 1; |
| while i < self.ops.len() && self.ops[i].offset < end { |
| if self.ops[i].end() <= end { |
| self.ops.remove(i); |
| } else { |
| // Delete the beginning of the entry. |
| let to_delete = end - self.ops[i].offset; |
| self.ops[i].offset += to_delete; |
| self.ops[i].data.drain(..to_delete as usize); |
| break; |
| } |
| } |
| self.inner.seek(SeekFrom::Current(buf.len() as i64))?; |
| Ok(buf.len()) |
| } |
| |
| fn flush(&mut self) -> Result<()> { |
| self.inner.flush() |
| } |
| } |
| |
| impl<T: std::io::Seek> Seek for TransactionManager<T> { |
| fn seek(&mut self, pos: SeekFrom) -> Result<u64> { |
| self.inner.seek(pos) |
| } |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use { |
| super::TransactionManager, |
| std::io::{Cursor, Read, Seek, SeekFrom, Write}, |
| }; |
| |
| #[test] |
| fn test_read_fall_through_when_no_transaction() { |
| let mut manager = TransactionManager::new(Cursor::new(vec![55; 100])); |
| let mut read_buf = vec![0; 3]; |
| assert_eq!(manager.read(&mut read_buf).expect("read failed"), 3); |
| assert_eq!(&read_buf, &[55, 55, 55]); |
| } |
| |
| #[test] |
| fn test_write_fall_through_when_no_transaction() { |
| let mut manager = TransactionManager::new(Cursor::new(vec![0; 4])); |
| let write_buf = vec![1, 2, 3]; |
| assert_eq!(manager.write(&write_buf).expect("write failed"), 3); |
| assert_eq!(&manager.into_inner().into_inner(), &[1, 2, 3, 0]); |
| } |
| |
| #[test] |
| fn test_read_part_transaction_part_inner() { |
| let mut manager = TransactionManager::new(Cursor::new(vec![55; 100])); |
| assert!(manager.begin_transaction()); |
| assert_eq!(manager.write(&[1, 2, 3]).expect("write failed"), 3); |
| assert_eq!(manager.seek(SeekFrom::Start(10)).expect("seek failed"), 10); |
| assert_eq!(manager.write(&[4, 5, 6]).expect("write failed"), 3); |
| assert_eq!(manager.seek(SeekFrom::Start(7)).expect("seek failed"), 7); |
| let mut read_buf = vec![0; 7]; |
| assert_eq!(manager.read(&mut read_buf).expect("read failed"), 7); |
| assert_eq!(&read_buf, &[55, 55, 55, 4, 5, 6, 55]); |
| assert_eq!(manager.seek(SeekFrom::Current(0)).expect("seek failed"), 14); |
| } |
| |
| #[test] |
| fn test_write_extend_entry() { |
| let mut manager = TransactionManager::new(Cursor::new(vec![55; 100])); |
| assert!(manager.begin_transaction()); |
| assert_eq!(manager.write(&[1, 2, 3]).expect("write failed"), 3); |
| assert_eq!(manager.seek(SeekFrom::Start(10)).expect("seek failed"), 10); |
| assert_eq!(manager.write(&[4, 5, 6]).expect("write failed"), 3); |
| assert_eq!(manager.write(&[7, 8, 9]).expect("write failed"), 3); |
| assert_eq!(manager.seek(SeekFrom::Start(11)).expect("seek failed"), 11); |
| let mut read_buf = vec![0; 4]; |
| assert_eq!(manager.read(&mut read_buf).expect("read failed"), 4); |
| assert_eq!(&read_buf, &[5, 6, 7, 8]); |
| } |
| |
| #[test] |
| fn test_write_existing_entry_encompasses_write() { |
| let mut manager = TransactionManager::new(Cursor::new(vec![55; 100])); |
| assert!(manager.begin_transaction()); |
| assert_eq!(manager.write(&[1, 2, 3]).expect("write failed"), 3); |
| assert_eq!(manager.seek(SeekFrom::Start(10)).expect("seek failed"), 10); |
| assert_eq!(manager.write(&[4, 5, 6, 7, 8, 9]).expect("write failed"), 6); |
| assert_eq!(manager.seek(SeekFrom::Start(12)).expect("seek failed"), 12); |
| assert_eq!(manager.write(&[99, 100]).expect("write failed"), 2); |
| assert_eq!(manager.seek(SeekFrom::Start(11)).expect("seek failed"), 11); |
| let mut read_buf = vec![0; 4]; |
| assert_eq!(manager.read(&mut read_buf).expect("read failed"), 4); |
| assert_eq!(&read_buf, &[5, 99, 100, 8]); |
| } |
| |
| #[test] |
| fn test_write_partial_overlap_and_extension() { |
| let mut manager = TransactionManager::new(Cursor::new(vec![55; 100])); |
| assert!(manager.begin_transaction()); |
| assert_eq!(manager.write(&[1, 2, 3]).expect("write failed"), 3); |
| assert_eq!(manager.seek(SeekFrom::Start(10)).expect("seek failed"), 10); |
| assert_eq!(manager.write(&[4, 5, 6, 7, 8, 9]).expect("write failed"), 6); |
| assert_eq!(manager.seek(SeekFrom::Start(14)).expect("seek failed"), 14); |
| assert_eq!(manager.write(&[99, 100, 101, 102]).expect("write failed"), 4); |
| assert_eq!(manager.seek(SeekFrom::Start(13)).expect("seek failed"), 13); |
| let mut read_buf = vec![0; 6]; |
| assert_eq!(manager.read(&mut read_buf).expect("read failed"), 6); |
| assert_eq!(&read_buf, &[7, 99, 100, 101, 102, 55]); |
| } |
| |
| #[test] |
| fn test_write_no_overlap() { |
| let mut manager = TransactionManager::new(Cursor::new(vec![55; 100])); |
| assert!(manager.begin_transaction()); |
| assert_eq!(manager.write(&[1, 2, 3]).expect("write failed"), 3); |
| assert_eq!(manager.seek(SeekFrom::Start(8)).expect("seek failed"), 8); |
| assert_eq!(manager.write(&[4, 5, 6]).expect("write failed"), 3); |
| assert_eq!(manager.seek(SeekFrom::Start(4)).expect("seek failed"), 4); |
| assert_eq!(manager.write(&[7, 8, 9]).expect("write failed"), 3); |
| assert_eq!(manager.seek(SeekFrom::Start(2)).expect("seek failed"), 2); |
| let mut read_buf = vec![0; 10]; |
| assert_eq!(manager.read(&mut read_buf).expect("read failed"), 10); |
| assert_eq!(&read_buf, &[3, 55, 7, 8, 9, 55, 4, 5, 6, 55]); |
| } |
| |
| #[test] |
| fn test_trim_following_entries() { |
| let mut manager = TransactionManager::new(Cursor::new(vec![55; 100])); |
| assert!(manager.begin_transaction()); |
| assert_eq!(manager.seek(SeekFrom::Start(2)).expect("seek failed"), 2); |
| assert_eq!(manager.write(&[1, 2, 3]).expect("write failed"), 3); |
| assert_eq!(manager.seek(SeekFrom::Current(1)).expect("seek failed"), 6); |
| assert_eq!(manager.write(&[4, 5, 6]).expect("write failed"), 3); |
| assert_eq!(manager.seek(SeekFrom::Start(1)).expect("seek failed"), 1); |
| assert_eq!(manager.write(&[100, 101, 102, 103, 104, 105]).expect("write failed"), 6); |
| assert_eq!(manager.seek(SeekFrom::Start(0)).expect("seek failed"), 0); |
| let mut read_buf = vec![0; 10]; |
| assert_eq!(manager.read(&mut read_buf).expect("read failed"), 10); |
| assert_eq!(&read_buf, &[55, 100, 101, 102, 103, 104, 105, 5, 6, 55]); |
| } |
| |
| #[test] |
| fn test_revert_transaction() { |
| let mut manager = TransactionManager::new(Cursor::new(vec![1, 2, 3, 4, 5, 6])); |
| assert!(manager.begin_transaction()); |
| assert_eq!(manager.write(&[7, 8, 9]).expect("write failed"), 3); |
| manager.revert(); |
| assert_eq!(manager.seek(SeekFrom::Start(1)).expect("seek failed"), 1); |
| let mut read_buf = vec![0; 4]; |
| assert_eq!(manager.read(&mut read_buf).expect("read failed"), 4); |
| assert_eq!(&read_buf, &[2, 3, 4, 5]); |
| } |
| |
| #[test] |
| fn test_commit_transaction() { |
| let mut manager = TransactionManager::new(Cursor::new(vec![55; 10])); |
| assert!(manager.begin_transaction()); |
| assert_eq!(manager.write(&[1, 2]).expect("write failed"), 2); |
| assert_eq!(manager.seek(SeekFrom::Current(1)).expect("seek failed"), 3); |
| assert_eq!(manager.write(&[3, 4]).expect("write failed"), 2); |
| assert_eq!(manager.seek(SeekFrom::Current(1)).expect("seek failed"), 6); |
| assert_eq!(manager.write(&[5, 6]).expect("write failed"), 2); |
| manager.commit().expect("commit failed"); |
| assert_eq!(&manager.into_inner().into_inner(), &[1, 2, 55, 3, 4, 55, 5, 6, 55, 55]); |
| } |
| } |
