use super::disk_adaptor::{DirectDiskAdaptor, DiskAdaptor};
use crate::error::{Aria2Error, FatalError, Result};
use crate::filesystem::disk_space::check_disk_space;
use std::path::Path;
#[cfg_attr(target_os = "linux", allow(dead_code))]
static SECURE_FALLOC_WARN_ONCE: std::sync::Once = std::sync::Once::new();
#[derive(Debug, Clone, Copy, PartialEq, Default)]
pub enum AllocationStrategy {
#[default]
None,
Prealloc,
Falloc,
Trunc,
Mmap,
}
impl AllocationStrategy {
#[allow(clippy::should_implement_trait)]
pub fn from_str(s: &str) -> Self {
match s {
"prealloc" => AllocationStrategy::Prealloc,
"falloc" => AllocationStrategy::Falloc,
"trunc" => AllocationStrategy::Trunc,
"mmap" => AllocationStrategy::Mmap,
_ => AllocationStrategy::None,
}
}
}
pub async fn allocate_file<D: DiskAdaptor>(
adaptor: &mut D,
_path: &Path,
length: u64,
strategy: AllocationStrategy,
secure: bool,
) -> Result<()> {
match strategy {
AllocationStrategy::None => Ok(()),
AllocationStrategy::Prealloc => preallocate(adaptor, length).await,
AllocationStrategy::Falloc => fallocate(adaptor, length, secure).await,
AllocationStrategy::Trunc => truncate(adaptor, length).await,
AllocationStrategy::Mmap => fallocate(adaptor, length, secure).await,
}
}
pub async fn preallocate_file(
path: &Path,
length: u64,
strategy: &str,
secure: bool,
) -> Result<()> {
preallocate_file_with_progress(path, length, strategy, None::<&fn(u64, u64)>, secure).await
}
pub async fn preallocate_file_with_progress<F>(
path: &Path,
length: u64,
strategy: &str,
on_progress: Option<&F>,
secure: bool,
) -> Result<()>
where
F: Fn(u64, u64) + Send + Sync,
{
let alloc_strategy = AllocationStrategy::from_str(strategy);
if length == 0 || alloc_strategy == AllocationStrategy::None {
return Ok(());
}
if let Err(_e) = check_disk_space(path, length) {
return Err(Aria2Error::Fatal(FatalError::DiskSpaceExhausted));
}
if let Some(parent) = path.parent() {
let parent: &Path = parent;
if !parent.exists() {
tokio::fs::create_dir_all(parent)
.await
.map_err(|e: std::io::Error| Aria2Error::Io(e.to_string()))?;
}
}
const PROGRESS_THRESHOLD: u64 = 100 * 1024 * 1024;
if let Some(cb) = on_progress
&& length >= PROGRESS_THRESHOLD
{
cb(0, length);
}
let mut adaptor = DirectDiskAdaptor::new();
adaptor.open(path).await?;
allocate_file(&mut adaptor, path, length, alloc_strategy, secure).await?;
if let Some(cb) = on_progress
&& length >= PROGRESS_THRESHOLD
{
cb(length, length);
}
adaptor.close().await
}
async fn preallocate<D: DiskAdaptor>(adaptor: &mut D, length: u64) -> Result<()> {
adaptor.truncate(length).await
}
async fn async_zero_fill<D: DiskAdaptor>(adaptor: &mut D, length: u64) -> Result<()> {
const CHUNK_SIZE: usize = 1024 * 1024; let zero_chunk = vec![0u8; CHUNK_SIZE];
let mut remaining = length;
let mut offset: u64 = 0;
while remaining > 0 {
let write_len = remaining.min(CHUNK_SIZE as u64) as usize;
adaptor.write(offset, &zero_chunk[..write_len]).await?;
offset += write_len as u64;
remaining -= write_len as u64;
tokio::task::yield_now().await;
}
Ok(())
}
#[cfg_attr(target_os = "linux", allow(unused_variables))]
async fn fallocate<D: DiskAdaptor>(adaptor: &mut D, length: u64, secure: bool) -> Result<()> {
#[cfg(target_os = "linux")]
{
if let Some(fd) = adaptor.unix_raw_fd() {
let ret = unsafe {
libc::fallocate(
fd,
0 as libc::c_int,
0,
length as libc::off_t,
)
};
if ret == 0 {
return Ok(());
}
let errno = unsafe { *libc::__errno_location() };
if errno == libc::EOPNOTSUPP {
tracing::warn!(
length,
"fallocate(2) not supported by filesystem; \
falling back to async zero-fill"
);
adaptor.truncate(length).await?;
return async_zero_fill(adaptor, length).await;
}
Err(Aria2Error::Io(
std::io::Error::from_raw_os_error(errno).to_string(),
))
} else {
adaptor.truncate(length).await
}
}
#[cfg(all(unix, not(target_os = "linux")))]
{
#[cfg(target_os = "macos")]
{
match adaptor.unix_raw_fd() {
Some(fd) => {
adaptor.truncate(length).await?;
let fstore = libc::fstore_t {
fst_flags: libc::F_ALLOCATEALL as libc::c_uint,
fst_posmode: libc::F_PEOFPOSMODE,
fst_offset: 0,
fst_length: length as libc::off_t,
fst_bytesalloc: 0,
};
let ret = unsafe { libc::fcntl(fd, libc::F_PREALLOCATE, &fstore) };
if ret != 0 {
tracing::warn!(
length,
"F_PREALLOCATE failed on macOS, file remains sparse"
);
return Ok(());
}
if secure {
async_zero_fill(adaptor, length).await
} else {
SECURE_FALLOC_WARN_ONCE.call_once(|| {
tracing::warn!(
"F_PREALLOCATE succeeded but does not zero-fill; \
residual disk data may be exposed. \
Enable --secure-falloc=true to zero-fill (at a \
performance cost). This warning is logged once."
);
});
Ok(())
}
}
None => adaptor.truncate(length).await,
}
}
#[cfg(not(target_os = "macos"))]
{
adaptor.truncate(length).await
}
}
#[cfg(not(unix))]
{
adaptor.truncate(length).await?;
let valid_data_succeeded: bool = if let Some(handle) = adaptor.windows_raw_handle() {
let ok = unsafe {
windows_sys::Win32::Storage::FileSystem::SetFileValidData(handle, length as i64)
};
if ok == 0 {
let err = unsafe { windows_sys::Win32::Foundation::GetLastError() };
if err == windows_sys::Win32::Foundation::ERROR_PRIVILEGE_NOT_HELD {
tracing::warn!(
length,
"SetFileValidData requires SE_MANAGE_VOLUME_PRIVILEGE; \
falling back to sparse file. This causes ~2x I/O on \
first write because each block must be zeroed by the \
writer instead of being pre-allocated."
);
} else {
tracing::warn!(length, err, "SetFileValidData failed; file remains sparse");
}
false
} else {
true
}
} else {
false
};
if valid_data_succeeded {
if secure {
async_zero_fill(adaptor, length).await
} else {
SECURE_FALLOC_WARN_ONCE.call_once(|| {
tracing::warn!(
"SetFileValidData succeeded but does not zero-fill; \
residual disk data may be exposed. \
Enable --secure-falloc=true to zero-fill (at a \
performance cost). This warning is logged once."
);
});
Ok(())
}
} else {
Ok(())
}
}
}
async fn truncate<D: DiskAdaptor>(adaptor: &mut D, length: u64) -> Result<()> {
adaptor.truncate(length).await
}
pub async fn get_available_space(path: &Path) -> Result<u64> {
let parent = path.parent().unwrap_or_else(|| Path::new("."));
#[cfg(target_os = "linux")]
{
let _metadata = tokio::fs::metadata(parent)
.await
.map_err(|e| Aria2Error::Io(e.to_string()))?;
let statvfs_result = unsafe {
let mut stat: libc::statvfs64 = std::mem::zeroed();
let ret = libc::statvfs64(
parent.to_str().unwrap_or(".").as_ptr() as *const i8,
&mut stat,
);
(ret, stat)
};
if statvfs_result.0 == 0 {
let stat = statvfs_result.1;
Ok(stat.f_bavail as u64 * stat.f_frsize as u64)
} else {
Err(Aria2Error::Io("Failed to get disk space".to_string()))
}
}
#[cfg(all(unix, not(target_os = "linux")))]
{
let _metadata = tokio::fs::metadata(parent)
.await
.map_err(|e| Aria2Error::Io(e.to_string()))?;
let statvfs_result = unsafe {
let mut stat: libc::statvfs = std::mem::zeroed();
let ret = libc::statvfs(
parent.to_str().unwrap_or(".").as_ptr() as *const i8,
&mut stat,
);
(ret, stat)
};
if statvfs_result.0 == 0 {
let stat = statvfs_result.1;
Ok(stat.f_bavail as u64 * stat.f_frsize as u64)
} else {
Err(Aria2Error::Io("Failed to get disk space".to_string()))
}
}
#[cfg(windows)]
{
let metadata = tokio::fs::metadata(parent)
.await
.map_err(|e| Aria2Error::Io(e.to_string()))?;
let free = metadata.len();
if free > 0 { Ok(free) } else { Ok(u64::MAX / 2) }
}
#[cfg(all(not(unix), not(windows)))]
{
Ok(u64::MAX)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_allocation_strategy_from_str() {
assert_eq!(
AllocationStrategy::from_str("none"),
AllocationStrategy::None
);
assert_eq!(
AllocationStrategy::from_str("prealloc"),
AllocationStrategy::Prealloc
);
assert_eq!(
AllocationStrategy::from_str("falloc"),
AllocationStrategy::Falloc
);
assert_eq!(
AllocationStrategy::from_str("trunc"),
AllocationStrategy::Trunc
);
assert_eq!(
AllocationStrategy::from_str("mmap"),
AllocationStrategy::Mmap
);
assert_eq!(
AllocationStrategy::from_str("invalid"),
AllocationStrategy::None
);
assert_eq!(AllocationStrategy::from_str(""), AllocationStrategy::None);
}
#[tokio::test]
async fn test_preallocate_file_none() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test_none.bin");
preallocate_file(&path, 1024, "none", false).await.unwrap();
assert!(!path.exists());
}
#[tokio::test]
async fn test_preallocate_file_trunc() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test_trunc.bin");
preallocate_file(&path, 4096, "trunc", false).await.unwrap();
let metadata = tokio::fs::metadata(&path).await.unwrap();
assert_eq!(metadata.len(), 4096);
}
#[tokio::test]
async fn test_preallocate_file_prealloc() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test_prealloc.bin");
preallocate_file(&path, 1024 * 1024, "prealloc", false)
.await
.unwrap();
let metadata = tokio::fs::metadata(&path).await.unwrap();
assert_eq!(metadata.len(), 1024 * 1024);
}
#[tokio::test]
async fn test_preallocate_zero_length() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test_zero.bin");
preallocate_file(&path, 0, "trunc", false).await.unwrap();
assert!(!path.exists());
}
#[tokio::test]
async fn test_preallocate_creates_parent_dirs() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("sub1").join("sub2").join("test_nested.bin");
preallocate_file(&path, 100, "trunc", false).await.unwrap();
assert!(path.exists());
let metadata = tokio::fs::metadata(&path).await.unwrap();
assert_eq!(metadata.len(), 100);
}
#[tokio::test]
async fn test_get_available_space_returns_value() {
let dir = tempfile::tempdir().unwrap();
let space = get_available_space(dir.path()).await;
assert!(space.is_ok());
let val = space.unwrap();
assert!(val > 0);
}
#[tokio::test]
async fn test_preallocate_overwrite_existing() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test_overwrite.bin");
tokio::fs::write(&path, b"original data").await.unwrap();
preallocate_file(&path, 2048, "trunc", false).await.unwrap();
let metadata = tokio::fs::metadata(&path).await.unwrap();
assert_eq!(metadata.len(), 2048);
}
#[tokio::test]
async fn test_allocate_file_cross_platform_prealloc() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test_alloc_prealloc.bin");
tokio::fs::write(&path, b"hello").await.unwrap();
preallocate_file(&path, 10 * 1024 * 1024, "prealloc", false)
.await
.unwrap();
let metadata = tokio::fs::metadata(&path).await.unwrap();
assert_eq!(metadata.len(), 10 * 1024 * 1024);
}
#[tokio::test]
async fn test_allocate_file_cross_platform_falloc() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test_alloc_falloc.bin");
tokio::fs::write(&path, b"initial data").await.unwrap();
preallocate_file(&path, 5 * 1024 * 1024, "falloc", false)
.await
.unwrap();
let metadata = tokio::fs::metadata(&path).await.unwrap();
assert_eq!(metadata.len(), 5 * 1024 * 1024);
}
#[tokio::test]
async fn test_fallocate_creates_sparse_file_of_correct_size() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test_falloc_sparse.bin");
preallocate_file(&path, 1024 * 1024, "falloc", false)
.await
.unwrap();
let metadata = tokio::fs::metadata(&path).await.unwrap();
assert_eq!(metadata.len(), 1024 * 1024);
}
#[tokio::test]
async fn test_allocate_file_cross_platform_trunc() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test_alloc_trunc.bin");
tokio::fs::write(&path, b"some initial content here")
.await
.unwrap();
preallocate_file(&path, 1024 * 1024, "trunc", false)
.await
.unwrap();
let metadata = tokio::fs::metadata(&path).await.unwrap();
assert_eq!(metadata.len(), 1024 * 1024);
}
#[tokio::test]
async fn test_allocate_file_cross_platform_none() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test_alloc_none.bin");
preallocate_file(&path, 1024 * 1024, "none", false)
.await
.unwrap();
assert!(!path.exists());
}
#[tokio::test]
async fn test_allocate_large_file() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test_large_alloc.bin");
preallocate_file(&path, 50 * 1024 * 1024, "falloc", false)
.await
.unwrap();
let metadata = tokio::fs::metadata(&path).await.unwrap();
assert_eq!(metadata.len(), 50 * 1024 * 1024);
use tokio::io::{AsyncSeekExt, AsyncWriteExt};
let mut file = tokio::fs::OpenOptions::new()
.write(true)
.open(&path)
.await
.unwrap();
file.seek(std::io::SeekFrom::Start(49 * 1024 * 1024))
.await
.unwrap();
file.write_all(b"end marker").await.unwrap();
file.flush().await.unwrap();
drop(file);
let metadata = tokio::fs::metadata(&path).await.unwrap();
assert_eq!(metadata.len(), 50 * 1024 * 1024);
}
#[tokio::test]
async fn test_all_strategies_same_result() {
let dir = tempfile::tempdir().unwrap();
let test_size: u64 = 1024 * 100;
let strategies = ["prealloc", "falloc", "trunc"];
for (i, strategy) in strategies.iter().enumerate() {
let path = dir.path().join(format!("test_strategy_{}.bin", i));
preallocate_file(&path, test_size, strategy, false)
.await
.unwrap();
let metadata = tokio::fs::metadata(&path).await.unwrap();
assert_eq!(
metadata.len(),
test_size,
"Strategy {} produced wrong size",
strategy
);
}
}
#[tokio::test]
async fn test_preallocate_with_progress_callback() {
use std::sync::{Arc, Mutex};
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test_progress.bin");
let progress_calls: Arc<Mutex<Vec<(u64, u64)>>> = Arc::new(Mutex::new(Vec::new()));
let pc = progress_calls.clone();
preallocate_file_with_progress(
&path,
150 * 1024 * 1024, "prealloc",
Some(&|allocated, total| {
pc.lock().unwrap().push((allocated, total));
}),
false,
)
.await
.unwrap();
{
let calls = progress_calls.lock().unwrap();
assert!(
calls.len() >= 2,
"expected at least 2 progress calls, got {}",
calls.len()
);
assert_eq!(calls.first().unwrap().0, 0); assert_eq!(
calls.last().unwrap().0,
150 * 1024 * 1024 );
assert_eq!(calls.last().unwrap().1, 150 * 1024 * 1024);
}
let metadata = tokio::fs::metadata(&path).await.unwrap();
assert_eq!(metadata.len(), 150 * 1024 * 1024);
}
#[tokio::test]
async fn test_preallocate_small_file_no_progress_callback() {
use std::sync::{Arc, Mutex};
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test_small_progress.bin");
let progress_calls: Arc<Mutex<Vec<(u64, u64)>>> = Arc::new(Mutex::new(Vec::new()));
let pc = progress_calls.clone();
preallocate_file_with_progress(
&path,
1024, "trunc",
Some(&|allocated, total| {
pc.lock().unwrap().push((allocated, total));
}),
false,
)
.await
.unwrap();
let calls = progress_calls.lock().unwrap();
assert!(
calls.is_empty(),
"small file should not trigger progress, got {} calls",
calls.len()
);
}
#[test]
fn test_default_allocation_is_falloc() {
use crate::constants;
assert_eq!(constants::DEFAULT_FILE_ALLOCATION, "falloc");
assert_eq!(
AllocationStrategy::from_str(constants::DEFAULT_FILE_ALLOCATION),
AllocationStrategy::Falloc
);
}
#[tokio::test]
async fn test_async_zero_fill() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("test_zero_fill.bin");
let mut adaptor = DirectDiskAdaptor::new();
adaptor.open(&path).await.unwrap();
adaptor.truncate(5 * 1024 * 1024).await.unwrap(); async_zero_fill(&mut adaptor, 5 * 1024 * 1024)
.await
.unwrap();
adaptor.close().await.unwrap();
let metadata = tokio::fs::metadata(&path).await.unwrap();
assert_eq!(metadata.len(), 5 * 1024 * 1024);
let content = tokio::fs::read(&path).await.unwrap();
assert!(content.iter().all(|&b| b == 0), "File should be all zeros");
}
#[test]
fn test_secure_falloc_default() {
use crate::request::request_group::DownloadOptions;
let opts = DownloadOptions::default();
assert!(!opts.secure_falloc, "secure_falloc should default to false");
}
}