ab_farmer_components/file_ext.rs
1//! File extension trait
2
3use std::fs::{File, OpenOptions};
4use std::io::Result;
5
6/// Extension convenience trait that allows setting some file opening options in cross-platform way
7pub trait OpenOptionsExt {
8 /// Advise OS/file system that file will use random access and read-ahead behavior is
9 /// undesirable, only has impact on Windows, for other operating systems see [`FileExt`]
10 fn advise_random_access(&mut self) -> &mut Self;
11
12 /// Advise OS/file system that file will use sequential access and read-ahead behavior is
13 /// desirable, only has impact on Windows, for other operating systems see [`FileExt`]
14 fn advise_sequential_access(&mut self) -> &mut Self;
15
16 /// Use Direct I/O on Linux and disable buffering on Windows.
17 ///
18 /// NOTE: There are major alignment requirements described here:
19 /// <https://learn.microsoft.com/en-us/windows/win32/fileio/file-buffering#alignment-and-file-access-requirements>
20 /// <https://man7.org/linux/man-pages/man2/open.2.html>
21 fn use_direct_io(&mut self) -> &mut Self;
22}
23
24impl OpenOptionsExt for OpenOptions {
25 fn advise_random_access(&mut self) -> &mut Self {
26 cfg_select! {
27 windows => {{
28 use std::os::windows::fs::OpenOptionsExt;
29 // `FILE_FLAG_WRITE_THROUGH` below is a bit of a hack, especially in
30 // `advise_random_access`, but it helps with memory usage and feels like should be
31 // default. Since `.custom_flags()` overrides previous value, we need to set bitwise
32 // OR of two flags rather that two flags separately.
33 self.custom_flags(
34 windows::Win32::Storage::FileSystem::FILE_FLAG_RANDOM_ACCESS.0
35 | windows::Win32::Storage::FileSystem::FILE_FLAG_WRITE_THROUGH.0,
36 )
37 }}
38 _ => {
39 // Not supported
40 self
41 }
42 }
43 }
44
45 fn advise_sequential_access(&mut self) -> &mut Self {
46 cfg_select! {
47 windows => {{
48 use std::os::windows::fs::OpenOptionsExt;
49 self.custom_flags(windows::Win32::Storage::FileSystem::FILE_FLAG_SEQUENTIAL_SCAN.0)
50 }}
51 _ => {
52 // Not supported
53 self
54 }
55 }
56 }
57
58 fn use_direct_io(&mut self) -> &mut Self {
59 cfg_select! {
60 windows => {{
61 use std::os::windows::fs::OpenOptionsExt;
62 self.custom_flags(
63 windows::Win32::Storage::FileSystem::FILE_FLAG_WRITE_THROUGH.0
64 | windows::Win32::Storage::FileSystem::FILE_FLAG_NO_BUFFERING.0,
65 )
66 }}
67 target_os = "linux" => {{
68 use std::os::unix::fs::OpenOptionsExt;
69 self.custom_flags(libc::O_DIRECT)
70 }}
71 _ => {
72 // Not supported
73 self
74 }
75 }
76 }
77}
78
79/// Extension convenience trait that allows pre-allocating files, suggesting random access pattern
80/// and doing cross-platform exact reads/writes
81pub trait FileExt {
82 /// Get file size
83 fn size(&self) -> Result<u64>;
84
85 /// Make sure file has specified number of bytes allocated for it
86 fn preallocate(&self, len: u64) -> Result<()>;
87
88 /// Advise OS/file system that file will use random access and read-ahead behavior is
89 /// undesirable, on Windows this can only be set when file is opened, see [`OpenOptionsExt`]
90 fn advise_random_access(&self) -> Result<()>;
91
92 /// Advise OS/file system that file will use sequential access and read-ahead behavior is
93 /// desirable, on Windows this can only be set when file is opened, see [`OpenOptionsExt`]
94 fn advise_sequential_access(&self) -> Result<()>;
95
96 /// Disable cache on macOS
97 fn disable_cache(&self) -> Result<()>;
98
99 /// Read exact number of bytes at a specific offset
100 fn read_exact_at(&self, buf: &mut [u8], offset: u64) -> Result<()>;
101
102 /// Write all provided bytes at a specific offset
103 fn write_all_at(&self, buf: &[u8], offset: u64) -> Result<()>;
104}
105
106impl FileExt for File {
107 fn size(&self) -> Result<u64> {
108 Ok(self.metadata()?.len())
109 }
110
111 fn preallocate(&self, len: u64) -> Result<()> {
112 fs2::FileExt::allocate(self, len)
113 }
114
115 fn advise_random_access(&self) -> Result<()> {
116 cfg_select! {
117 target_os = "linux" => {{
118 use std::os::unix::io::AsRawFd;
119 // SAFETY: Correct low-level FFI file
120 let err = unsafe { libc::posix_fadvise(self.as_raw_fd(), 0, 0, libc::POSIX_FADV_RANDOM) };
121 if err != 0 {
122 Err(std::io::Error::from_raw_os_error(err))
123 } else {
124 Ok(())
125 }
126 }}
127 target_os = "macos" => {{
128 use std::os::unix::io::AsRawFd;
129 // SAFETY: Correct low-level FFI file
130 if unsafe { libc::fcntl(self.as_raw_fd(), libc::F_RDAHEAD, 0) } != 0 {
131 Err(std::io::Error::last_os_error())
132 } else {
133 Ok(())
134 }
135 }}
136 _ => {
137 // Not supported
138 Ok(())
139 }
140 }
141 }
142
143 fn advise_sequential_access(&self) -> Result<()> {
144 cfg_select! {
145 target_os = "linux" => {{
146 use std::os::unix::io::AsRawFd;
147 // SAFETY: Correct low-level FFI file
148 let err =
149 unsafe { libc::posix_fadvise(self.as_raw_fd(), 0, 0, libc::POSIX_FADV_SEQUENTIAL) };
150 if err != 0 {
151 Err(std::io::Error::from_raw_os_error(err))
152 } else {
153 Ok(())
154 }
155 }}
156 target_os = "macos" => {{
157 use std::os::unix::io::AsRawFd;
158 // SAFETY: Correct low-level FFI file
159 if unsafe { libc::fcntl(self.as_raw_fd(), libc::F_RDAHEAD, 1) } != 0 {
160 Err(std::io::Error::last_os_error())
161 } else {
162 Ok(())
163 }
164 }}
165 _ => {
166 // Not supported
167 Ok(())
168 }
169 }
170 }
171
172 fn disable_cache(&self) -> Result<()> {
173 cfg_select! {
174 target_os = "macos" => {{
175 use std::os::unix::io::AsRawFd;
176 // SAFETY: Correct low-level FFI file
177 if unsafe { libc::fcntl(self.as_raw_fd(), libc::F_NOCACHE, 1) } != 0 {
178 Err(std::io::Error::last_os_error())
179 } else {
180 Ok(())
181 }
182 }}
183 _ => {
184 // Not supported
185 Ok(())
186 }
187 }
188 }
189
190 fn read_exact_at(&self, buf: &mut [u8], offset: u64) -> Result<()> {
191 cfg_select! {
192 unix => {
193 std::os::unix::fs::FileExt::read_exact_at(self, buf, offset)
194 }
195 windows => {{
196 let mut buf = buf;
197 let mut offset = offset;
198
199 while !buf.is_empty() {
200 match std::os::windows::fs::FileExt::seek_read(self, buf, offset) {
201 Ok(0) => {
202 break;
203 }
204 Ok(n) => {
205 buf = &mut buf[n..];
206 offset += n as u64;
207 }
208 Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => {
209 // Try again
210 }
211 Err(e) => {
212 return Err(e);
213 }
214 }
215 }
216
217 if !buf.is_empty() {
218 Err(std::io::Error::new(
219 std::io::ErrorKind::UnexpectedEof,
220 "failed to fill whole buffer",
221 ))
222 } else {
223 Ok(())
224 }
225 }}
226 _ => {
227 compile_error!("Unsupported platform (consider contributing)");
228 }
229 }
230 }
231
232 fn write_all_at(&self, buf: &[u8], offset: u64) -> Result<()> {
233 cfg_select! {
234 unix => {
235 std::os::unix::fs::FileExt::write_all_at(self, buf, offset)
236 }
237 windows => {{
238 let mut buf = buf;
239 let mut offset = offset;
240
241 while !buf.is_empty() {
242 match std::os::windows::fs::FileExt::seek_write(self, buf, offset) {
243 Ok(0) => {
244 return Err(std::io::Error::new(
245 std::io::ErrorKind::WriteZero,
246 "failed to write whole buffer",
247 ));
248 }
249 Ok(n) => {
250 buf = &buf[n..];
251 offset += n as u64;
252 }
253 Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => {
254 // Try again
255 }
256 Err(e) => {
257 return Err(e);
258 }
259 }
260 }
261
262 Ok(())
263 }}
264 _ => {
265 compile_error!("Unsupported platform (consider contributing)");
266 }
267 }
268 }
269}