Skip to main content

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        fs4::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 =
121                    unsafe { libc::posix_fadvise(self.as_raw_fd(), 0, 0, libc::POSIX_FADV_RANDOM) };
122                if err != 0 {
123                    Err(std::io::Error::from_raw_os_error(err))
124                } else {
125                    Ok(())
126                }
127            }
128            target_os = "macos" => {
129                use std::os::unix::io::AsRawFd;
130                // SAFETY: Correct low-level FFI file
131                if unsafe { libc::fcntl(self.as_raw_fd(), libc::F_RDAHEAD, 0) } != 0 {
132                    Err(std::io::Error::last_os_error())
133                } else {
134                    Ok(())
135                }
136            }
137            _ => {
138                // Not supported
139                Ok(())
140            }
141        }
142    }
143
144    fn advise_sequential_access(&self) -> Result<()> {
145        cfg_select! {
146            target_os = "linux" => {
147                use std::os::unix::io::AsRawFd;
148                // SAFETY: Correct low-level FFI file
149                let err = unsafe {
150                    libc::posix_fadvise(self.as_raw_fd(), 0, 0, libc::POSIX_FADV_SEQUENTIAL)
151                };
152                if err != 0 {
153                    Err(std::io::Error::from_raw_os_error(err))
154                } else {
155                    Ok(())
156                }
157            }
158            target_os = "macos" => {
159                use std::os::unix::io::AsRawFd;
160                // SAFETY: Correct low-level FFI file
161                if unsafe { libc::fcntl(self.as_raw_fd(), libc::F_RDAHEAD, 1) } != 0 {
162                    Err(std::io::Error::last_os_error())
163                } else {
164                    Ok(())
165                }
166            }
167            _ => {
168                // Not supported
169                Ok(())
170            }
171        }
172    }
173
174    fn disable_cache(&self) -> Result<()> {
175        cfg_select! {
176            target_os = "macos" => {
177                use std::os::unix::io::AsRawFd;
178                // SAFETY: Correct low-level FFI file
179                if unsafe { libc::fcntl(self.as_raw_fd(), libc::F_NOCACHE, 1) } != 0 {
180                    Err(std::io::Error::last_os_error())
181                } else {
182                    Ok(())
183                }
184            }
185            _ => {
186                // Not supported
187                Ok(())
188            }
189        }
190    }
191
192    fn read_exact_at(&self, buf: &mut [u8], offset: u64) -> Result<()> {
193        cfg_select! {
194            unix => std::os::unix::fs::FileExt::read_exact_at(self, buf, offset),
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(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                    Ok(())
219                } else {
220                    Err(std::io::Error::new(
221                        std::io::ErrorKind::UnexpectedEof,
222                        "failed to fill whole buffer",
223                    ))
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 => std::os::unix::fs::FileExt::write_all_at(self, buf, offset),
235            windows => {
236                let mut buf = buf;
237                let mut offset = offset;
238
239                while !buf.is_empty() {
240                    match std::os::windows::fs::FileExt::seek_write(self, buf, offset) {
241                        Ok(0) => {
242                            return Err(std::io::Error::new(
243                                std::io::ErrorKind::WriteZero,
244                                "failed to write whole buffer",
245                            ));
246                        }
247                        Ok(n) => {
248                            buf = &buf[n..];
249                            offset += n as u64;
250                        }
251                        Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {
252                            // Try again
253                        }
254                        Err(e) => {
255                            return Err(e);
256                        }
257                    }
258                }
259
260                Ok(())
261            }
262            _ => {
263                compile_error!("Unsupported platform (consider contributing)");
264            }
265        }
266    }
267}