palace 0.1.0

A tool for mounting datasets into memory for fast loading in deep learning tasks.
Documentation
use crate::utils::{check_mount_point, get_directory_size, PalaceError};
use anyhow::{Context, Result};
use indicatif::{ProgressBar, ProgressStyle};
use std::fs;
use std::io::ErrorKind;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::thread;
use std::time::Duration;

pub struct UnmountOptions {
    pub target: PathBuf,
    pub force: bool,
}

pub fn unmount_dataset(options: &UnmountOptions) -> Result<()> {
    check_mount_point(&options.target)?;

    if !options.target.exists() {
        return Err(PalaceError::InvalidMountPoint(options.target.to_string_lossy().into_owned()).into());
    }

    let size = get_directory_size(&options.target)?;
    println!("Preparing to unmount dataset at: {}", options.target.display());
    println!("This will free approximately {} bytes of memory.", size);

    if !options.force {
        if !confirm_unmount() {
            println!("Unmount operation cancelled.");
            return Ok(());
        }
    }

    unmount_and_clean(&options.target)
}

fn confirm_unmount() -> bool {
    println!("Are you sure you want to unmount this dataset? This operation cannot be undone.");
    println!("Type 'yes' to confirm or any other input to cancel:");

    let mut input = String::new();
    std::io::stdin().read_line(&mut input).expect("Failed to read line");
    input.trim().to_lowercase() == "yes"
}

fn unmount_and_clean(target: &Path) -> Result<()> {
    let pb = ProgressBar::new_spinner();
    pb.set_style(ProgressStyle::default_spinner()
        .template("{spinner:.green} {msg}")
        .unwrap());

    // 首先尝试卸载
    // Unmount the tmpfs filesystem if it's mounted
    if !is_tmpfs_mounted(target)? {
        return Err(anyhow::anyhow!("This target was not mounted in tmpfs."))
    }
    // if is_tmpfs_mounted(target)? {
    pb.set_message("Unmounting tmpfs...");
    unmount_tmpfs(target)?;
    // }

    // 等待一小段时间,确保卸载完成
    thread::sleep(Duration::from_secs(2));

    // 检查是否还有进程在使用该目录
    pb.set_message("Checking for active processes...");
    ensure_no_active_processes(target)?;

    // 删除目录内容
    pb.set_message("Removing dataset contents...");
    remove_directory_contents(target)?;

    // 尝试删除空目录
    pb.set_message("Removing empty directory...");
    for _ in 0..5 { // 尝试最多5次
        match std::fs::remove_dir(target) {
            Ok(_) => {
                pb.finish_with_message("Dataset successfully unmounted and cleaned.");
                return Ok(());
            }
            Err(e) if e.kind() == ErrorKind::WouldBlock || e.kind() == ErrorKind::Other => {
                pb.set_message("Directory still busy, retrying...");
                thread::sleep(Duration::from_secs(1));
            }
            Err(e) => return Err(e.into()),
        }
    }

    Err(anyhow::anyhow!("Failed to remove directory after multiple attempts"))
}


fn is_tmpfs_mounted(target: &Path) -> Result<bool> {
    let output = std::process::Command::new("mount")
        .output()
        .context("Failed to execute 'mount' command")?;
    let mount_info = String::from_utf8_lossy(&output.stdout);
    
    let mut target_str = target.display().to_string();
    
    if target_str.ends_with('/') {
        target_str.pop();
    }
    
    println!("Checking if tmpfs is not mounted on {}", target_str);
    
    Ok(mount_info.contains(&format!("tmpfs on {}", target_str)))
}

fn unmount_tmpfs(target: &Path) -> Result<()> {
    let output = Command::new("sudo")
        .args(&["umount", "-f", target.to_str().unwrap()])
        .output()
        .context("Failed to execute umount command")?;

    if !output.status.success() {
        let error = String::from_utf8_lossy(&output.stderr);
        return Err(anyhow::anyhow!("Failed to unmount tmpfs: {}", error));
    }

    Ok(())
}

fn ensure_no_active_processes(target: &Path) -> Result<()> {
    let output = Command::new("lsof")
        .arg(target.to_str().unwrap())
        .output()
        .context("Failed to execute lsof command")?;

    if !output.status.success() && !output.stdout.is_empty() {
        let processes = String::from_utf8_lossy(&output.stdout);
        return Err(anyhow::anyhow!("Active processes found: {}", processes));
    }

    Ok(())
}

fn remove_directory_contents(dir: &Path) -> Result<()> {
    for entry in fs::read_dir(dir)? {
        let entry = entry?;
        let path = entry.path();
        if path.is_dir() {
            fs::remove_dir_all(path)?;
        } else {
            fs::remove_file(path)?;
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs::File;
    use tempfile::tempdir;

    #[test]
    fn test_unmount_dataset() -> Result<()> {
        let temp_dir = tempdir()?;
        let mount_point = temp_dir.path().join("dataset");
        fs::create_dir(&mount_point)?;
        File::create(mount_point.join("test_file.txt"))?;

        let options = UnmountOptions {
            target: mount_point.clone(),
            force: true, // Skip confirmation for testing
        };

        unmount_dataset(&options)?;

        assert!(!mount_point.exists(), "Mount point should not exist after unmounting");
        Ok(())
    }
}