From e0e7058d14ca63ae506544f6418fb35ec19be4ef Mon Sep 17 00:00:00 2001
From: weishu <twsxtd@gmail.com>
Date: Mon, 19 Feb 2024 15:27:27 +0800
Subject: [PATCH] ksud: reclaim sparse space when install/uninstall modules.
 close #1367

---
 userspace/ksud/src/cli.rs    |  8 ++++
 userspace/ksud/src/module.rs |  8 ++++
 userspace/ksud/src/utils.rs  | 85 ++++++++++++++++++++++++++++++++++++
 3 files changed, 101 insertions(+)

diff --git a/userspace/ksud/src/cli.rs b/userspace/ksud/src/cli.rs
index 1f63583a..5f0b0c4c 100644
--- a/userspace/ksud/src/cli.rs
+++ b/userspace/ksud/src/cli.rs
@@ -119,6 +119,13 @@ enum Debug {
         /// destination file
         dst: String,
     },
+
+    /// Punch hole file
+    PunchHole {
+        /// file path
+        file: String,
+    },
+
     /// For testing
     Test,
 }
@@ -292,6 +299,7 @@ pub fn run() -> Result<()> {
                 utils::copy_sparse_file(src, dst)?;
                 Ok(())
             }
+            Debug::PunchHole { file } => utils::punch_hole(file),
             Debug::Test => todo!(),
         },
 
diff --git a/userspace/ksud/src/module.rs b/userspace/ksud/src/module.rs
index e71e2113..11b4f7aa 100644
--- a/userspace/ksud/src/module.rs
+++ b/userspace/ksud/src/module.rs
@@ -444,6 +444,10 @@ fn _install_module(zip: &str) -> Result<()> {
 
     exec_install_script(zip)?;
 
+    if let Err(e) = utils::punch_hole(tmp_module_img) {
+        warn!("Failed to punch hole: {}", e);
+    }
+
     info!("rename {tmp_module_img} to {}", defs::MODULE_UPDATE_IMG);
     // all done, rename the tmp image to modules_update.img
     if std::fs::rename(tmp_module_img, defs::MODULE_UPDATE_IMG).is_err() {
@@ -507,6 +511,10 @@ where
     // call the operation func
     let result = func(id, update_dir);
 
+    if let Err(e) = utils::punch_hole(modules_update_tmp_img) {
+        warn!("Failed to punch hole: {}", e);
+    }
+
     if let Err(e) = std::fs::rename(modules_update_tmp_img, defs::MODULE_UPDATE_IMG) {
         warn!("Rename image failed: {e}, try copy it.");
         utils::copy_sparse_file(modules_update_tmp_img, defs::MODULE_UPDATE_IMG)
diff --git a/userspace/ksud/src/utils.rs b/userspace/ksud/src/utils.rs
index dc782475..f5172406 100644
--- a/userspace/ksud/src/utils.rs
+++ b/userspace/ksud/src/utils.rs
@@ -231,3 +231,88 @@ pub fn copy_sparse_file<P: AsRef<Path>, Q: AsRef<Path>>(src: P, dst: Q) -> Resul
 
     Ok(())
 }
+
+pub fn punch_hole(src: impl AsRef<Path>) -> Result<()> {
+    let mut src_file = OpenOptions::new().write(true).read(true).open(src)?;
+
+    let st = rustix::fs::fstat(&src_file)?;
+
+    let bufsz = st.st_blksize;
+    let mut buf = vec![0u8; bufsz as usize];
+
+    let mut ct = 0;
+    let mut hole_sz = 0;
+    let mut hole_start = 0;
+
+    let segments = src_file.scan_chunks()?;
+    for segment in segments {
+        if segment.segment_type != SegmentType::Data {
+            continue;
+        }
+
+        let mut off = segment.start;
+        let end = segment.end + 1;
+
+        while off < end {
+            let mut rsz = rustix::io::pread(&src_file, &mut buf, off)? as u64;
+            if rsz > 0 && rsz > end - off {
+                // exceed the end of the boundary
+                rsz = end - off;
+            }
+
+            if rsz == 0 {
+                break;
+            }
+
+            if buf.iter().all(|&x| x == 0) {
+                // the whole buf is zero, mark it as a hole
+                if hole_sz == 0 {
+                    hole_start = off;
+                }
+                // for continuous zero, we can merge them into a bigger hole
+                hole_sz += rsz;
+            } else if hole_sz > 0 {
+                if let Err(e) = rustix::fs::fallocate(
+                    &src_file,
+                    rustix::fs::FallocateFlags::PUNCH_HOLE | rustix::fs::FallocateFlags::KEEP_SIZE,
+                    hole_start,
+                    hole_sz,
+                ) {
+                    log::warn!("Failed to punch hole: {:?}", e);
+                }
+
+                ct += hole_sz;
+
+                hole_sz = 0;
+                hole_start = 0;
+            }
+
+            off += rsz;
+        }
+
+        // if the last segment is a hole, we need to punch it
+        if hole_sz > 0 {
+            let mut alloc_sz = hole_sz;
+            if off >= end {
+                alloc_sz += st.st_blksize as u64;
+            }
+            if let Err(e) = rustix::fs::fallocate(
+                &src_file,
+                rustix::fs::FallocateFlags::PUNCH_HOLE | rustix::fs::FallocateFlags::KEEP_SIZE,
+                hole_start,
+                alloc_sz,
+            ) {
+                log::warn!("Failed to punch hole: {:?}", e);
+            }
+
+            ct += hole_sz;
+        }
+    }
+
+    log::info!(
+        "Punched {} of hole",
+        humansize::format_size(ct, humansize::DECIMAL)
+    );
+
+    Ok(())
+}