sys_util: add WriteZeroesAt trait

Add a variant of WriteZeroes that allows the caller to specify the
offset explicitly instead of using the file's cursor.  This gets rid of
one of the last bits of shared state between disk file users, which will
help in implementing multi-queue support.

Additionally, modify the WriteZeroes trait to use a generic
implementation based on WriteZeroesAt + Seek when possible.

BUG=chromium:858815
TEST=Boot Termina in crosvm

Change-Id: If710159771aeeb55f4f7746dd4354b6c042144e8
Signed-off-by: Daniel Verkamp <dverkamp@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosvm/+/1913519

2 files changed, 50 insertions, 9 deletions

diff --git a/sys_util/src/lib.rs b/sys_util/src/lib.rs
index 331dbac..31b1662 100644
--- a/sys_util/src/lib.rs
+++ b/sys_util/src/lib.rs
@@ -70,7 +70,7 @@ pub use crate::guest_memory::Error as GuestMemoryError;
 pub use crate::mmap::Error as MmapError;
 pub use crate::seek_hole::SeekHole;
 pub use crate::signalfd::Error as SignalFdError;
-pub use crate::write_zeroes::{PunchHole, WriteZeroes};
+pub use crate::write_zeroes::{PunchHole, WriteZeroes, WriteZeroesAt};
 
 use std::ffi::CStr;
 use std::fs::{remove_file, File};
diff --git a/sys_util/src/write_zeroes.rs b/sys_util/src/write_zeroes.rs
index 0e733c7..7e28f53 100644
--- a/sys_util/src/write_zeroes.rs
+++ b/sys_util/src/write_zeroes.rs
@@ -4,7 +4,8 @@
 
 use std::cmp::min;
 use std::fs::File;
-use std::io::{self, Error, ErrorKind, Seek, SeekFrom, Write};
+use std::io::{self, Error, ErrorKind, Seek, SeekFrom};
+use std::os::unix::fs::FileExt;
 
 use crate::fallocate;
 use crate::FallocateMode;
@@ -51,13 +52,43 @@ pub trait WriteZeroes {
     }
 }
 
-impl<T: PunchHole + Seek + Write> WriteZeroes for T {
-    fn write_zeroes(&mut self, length: usize) -> io::Result<usize> {
+/// A trait for writing zeroes to an arbitrary position in a file.
+pub trait WriteZeroesAt {
+    /// Write up to `length` bytes of zeroes starting at `offset`, returning how many bytes were
+    /// written.
+    fn write_zeroes_at(&mut self, offset: u64, length: usize) -> io::Result<usize>;
+
+    /// Write zeroes starting at `offset` until `length` bytes have been written.
+    ///
+    /// This method will continuously call `write_zeroes_at` until the requested
+    /// `length` is satisfied or an error is encountered.
+    fn write_zeroes_all_at(&mut self, mut offset: u64, mut length: usize) -> io::Result<()> {
+        while length > 0 {
+            match self.write_zeroes_at(offset, length) {
+                Ok(0) => return Err(Error::from(ErrorKind::WriteZero)),
+                Ok(bytes_written) => {
+                    length = length
+                        .checked_sub(bytes_written)
+                        .ok_or(Error::from(ErrorKind::Other))?;
+                    offset = offset
+                        .checked_add(bytes_written as u64)
+                        .ok_or(Error::from(ErrorKind::Other))?;
+                }
+                Err(e) => {
+                    if e.kind() != ErrorKind::Interrupted {
+                        return Err(e);
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
+impl WriteZeroesAt for File {
+    fn write_zeroes_at(&mut self, offset: u64, length: usize) -> io::Result<usize> {
         // Try to punch a hole first.
-        let offset = self.seek(SeekFrom::Current(0))?;
         if let Ok(()) = self.punch_hole(offset, length as u64) {
-            // Advance the seek cursor as if we had done a real write().
-            self.seek(SeekFrom::Current(length as i64))?;
             return Ok(length);
         }
 
@@ -71,17 +102,27 @@ impl<T: PunchHole + Seek + Write> WriteZeroes for T {
         while nwritten < length {
             let remaining = length - nwritten;
             let write_size = min(remaining, buf_size);
-            nwritten += self.write(&buf[0..write_size])?;
+            nwritten += self.write_at(&buf[0..write_size], offset + nwritten as u64)?;
         }
         Ok(length)
     }
 }
 
+impl<T: WriteZeroesAt + Seek> WriteZeroes for T {
+    fn write_zeroes(&mut self, length: usize) -> io::Result<usize> {
+        let offset = self.seek(SeekFrom::Current(0))?;
+        let nwritten = self.write_zeroes_at(offset, length)?;
+        // Advance the seek cursor as if we had done a real write().
+        self.seek(SeekFrom::Current(nwritten as i64))?;
+        Ok(length)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use std::fs::OpenOptions;
-    use std::io::{Read, Seek, SeekFrom};
+    use std::io::{Read, Seek, SeekFrom, Write};
     use tempfile::TempDir;
 
     #[test]