file_alloc/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use std::future::Future;
4use std::io::{self, SeekFrom};
5use tokio::fs::File;
6use tokio::io::{AsyncSeekExt, AsyncWriteExt};
7
8mod unix;
9mod windows;
10
11#[cfg(unix)]
12use unix::try_fast_preallocate;
13#[cfg(windows)]
14use windows::try_fast_preallocate;
15
16#[cfg(unix)]
17pub use unix::init_fast_alloc;
18#[cfg(windows)]
19pub use windows::init_fast_alloc;
20
21pub trait FileAlloc {
22    fn allocate(&mut self, size: u64) -> impl Future<Output = io::Result<()>> + Send + Sync + '_;
23}
24
25impl FileAlloc for File {
26    async fn allocate(&mut self, size: u64) -> io::Result<()> {
27        let current_size = self.metadata().await?.len();
28        if current_size >= size || try_fast_preallocate(self, current_size, size).await? {
29            Ok(())
30        } else {
31            async_zero_fill(self, current_size, size).await
32        }
33    }
34}
35
36const CHUNK_SIZE: usize = 1024 * 1024;
37static ZEROS: [u8; CHUNK_SIZE] = [0; CHUNK_SIZE];
38
39async fn async_zero_fill(
40    file: &mut File,
41    mut current_size: u64,
42    target_size: u64,
43) -> io::Result<()> {
44    file.seek(SeekFrom::Start(current_size)).await?;
45    while current_size < target_size {
46        let remaining = target_size - current_size;
47        #[allow(clippy::cast_possible_truncation)]
48        let to_write = CHUNK_SIZE.min(remaining as usize);
49        let n = file.write(&ZEROS[..to_write]).await?;
50        if n == 0 {
51            return Err(io::ErrorKind::WriteZero.into());
52        }
53        current_size += n as u64;
54    }
55    file.flush().await?;
56    Ok(())
57}
58
59#[cfg(not(any(windows, unix)))]
60async fn try_fast_preallocate(_file: &File, _current_size: u64, _size: u64) -> io::Result<bool> {
61    Ok(false)
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67    use std::io::Read;
68    use tempfile::NamedTempFile;
69
70    /// 测试基础分配功能
71    #[tokio::test]
72    async fn test_allocate_basic() -> io::Result<()> {
73        let temp_file = NamedTempFile::new()?;
74        let mut file = File::options()
75            .read(true)
76            .write(true)
77            .open(temp_file.path())
78            .await?;
79
80        let target_size = 5 * 1024 * 1024; // 5MB
81        file.allocate(target_size).await?;
82
83        let metadata = file.metadata().await?;
84        assert_eq!(metadata.len(), target_size);
85
86        Ok(())
87    }
88
89    /// 测试幂等性：分配比当前更小的大小不应改变文件
90    #[tokio::test]
91    async fn test_allocate_idempotency() -> io::Result<()> {
92        let temp_file = NamedTempFile::new()?;
93        let mut file = File::options()
94            .read(true)
95            .write(true)
96            .open(temp_file.path())
97            .await?;
98
99        // 先分配 2MB
100        file.allocate(2 * 1024 * 1024).await?;
101        let size1 = file.metadata().await?.len();
102
103        // 尝试分配 1MB (应该直接返回 Ok)
104        file.allocate(1024 * 1024).await?;
105        let size2 = file.metadata().await?.len();
106
107        assert_eq!(size1, 2 * 1024 * 1024);
108        assert_eq!(size1, size2);
109        Ok(())
110    }
111
112    /// 测试大文件分块分配（触发循环写 0 逻辑）
113    #[tokio::test]
114    async fn test_allocate_large_chunk() -> io::Result<()> {
115        let temp_file = NamedTempFile::new()?;
116        let mut file = File::options()
117            .read(true)
118            .write(true)
119            .open(temp_file.path())
120            .await?;
121
122        // 分配 2.5MB，超过 1MB 的 CHUNK_SIZE
123        #[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
124        let target_size = (2.5 * 1024.0 * 1024.0) as u64;
125        file.allocate(target_size).await?;
126
127        assert_eq!(file.metadata().await?.len(), target_size);
128
129        // 验证文件末尾是否可以写入数据
130        file.seek(SeekFrom::End(0)).await?;
131        file.write_all(b"end").await?;
132        file.flush().await?;
133        assert_eq!(file.metadata().await?.len(), target_size + 3);
134
135        Ok(())
136    }
137
138    /// 验证分配出的空间读取出来全是 0
139    #[tokio::test]
140    #[cfg(not(windows))]
141    async fn test_allocate_zero_verification() -> io::Result<()> {
142        let temp_file = NamedTempFile::new()?;
143        let mut file = File::options()
144            .read(true)
145            .write(true)
146            .open(temp_file.path())
147            .await?;
148
149        let target_size = 100 * 1024; // 100KB
150        file.allocate(target_size).await?;
151
152        // 必须通过 std File 读取来验证内容
153        let mut std_file = std::fs::File::open(temp_file.path())?;
154        let mut buffer = Vec::new();
155        std_file.read_to_end(&mut buffer)?;
156
157        assert_eq!(buffer.len() as u64, target_size);
158        assert!(buffer.iter().all(|&b| b == 0));
159
160        Ok(())
161    }
162
163    /// 测试在已有数据的文件后面追加分配
164    #[tokio::test]
165    async fn test_allocate_append() -> io::Result<()> {
166        let temp_file = NamedTempFile::new()?;
167        let mut file = File::options()
168            .read(true)
169            .write(true)
170            .open(temp_file.path())
171            .await?;
172
173        // 先写入 10 字节数据
174        let initial_data = b"0123456789";
175        file.write_all(initial_data).await?;
176        file.flush().await?;
177
178        // 预分配到 100 字节
179        file.allocate(100).await?;
180
181        let mut std_file = std::fs::File::open(temp_file.path())?;
182        let mut buffer = Vec::new();
183        std_file.read_to_end(&mut buffer)?;
184
185        assert_eq!(buffer.len(), 100);
186        assert_eq!(&buffer[0..10], initial_data); // 原数据应保持不变
187        assert!(buffer[10..].iter().all(|&b| b == 0)); // 后续应全为 0
188
189        Ok(())
190    }
191}
file_alloc/lib.rs

file_alloc/
lib.rs
