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