Skip to main content

tycho_storage/fs/
mod.rs

1#![allow(clippy::disallowed_methods)]
2#![allow(clippy::disallowed_types)] // it's wrapper around Files so
3
4use std::borrow::Cow;
5use std::ffi::OsStr;
6use std::fs::{File, Metadata, OpenOptions};
7use std::os::fd::AsRawFd;
8use std::path::{Path, PathBuf};
9use std::sync::Arc;
10use std::time::Duration;
11
12use anyhow::{Context, Result};
13use tycho_util::fs::MappedFileMut;
14
15const BASE_DIR: &str = "temp";
16
17#[derive(Clone)]
18pub struct TempFileStorage {
19    storage_dir: Dir,
20}
21
22impl TempFileStorage {
23    const MAX_FILE_TTL: Duration = Duration::from_secs(86400); // 1 day
24
25    pub fn new(files_dir: &Dir) -> Result<Self> {
26        Ok(Self {
27            storage_dir: files_dir.create_subdir(BASE_DIR)?,
28        })
29    }
30
31    pub fn dir(&self) -> &Dir {
32        &self.storage_dir
33    }
34
35    pub async fn remove_outdated_files(&self) -> Result<()> {
36        let now = std::time::SystemTime::now();
37
38        let this = self.clone();
39        tokio::task::spawn_blocking(move || {
40            this.retain_files(|path, metadata| match metadata.modified() {
41                Ok(modified) => {
42                    let since_modified = now.duration_since(modified).unwrap_or(Duration::ZERO);
43                    since_modified <= Self::MAX_FILE_TTL
44                }
45                Err(e) => {
46                    tracing::warn!(
47                        path = %path.display(),
48                        "failed to check file metadata: {e:?}"
49                    );
50                    false
51                }
52            })
53        })
54        .await?
55    }
56
57    pub fn retain_files<F>(&self, mut f: F) -> Result<()>
58    where
59        F: FnMut(&Path, &Metadata) -> bool,
60    {
61        let entries = self.storage_dir.entries()?;
62        for e in entries {
63            let e = e?;
64
65            let path = e.path();
66            let Ok(metadata) = std::fs::metadata(&path) else {
67                tracing::warn!(
68                    path = %path.display(),
69                    "failed to check downloaded file metadata: {e:?}"
70                );
71                continue;
72            };
73
74            let is_file = metadata.is_file();
75            let keep = is_file && f(&path, &metadata);
76            tracing::debug!(keep, path = %path.display(), "found downloaded file");
77
78            if keep {
79                continue;
80            }
81
82            let e = if is_file {
83                std::fs::remove_file(&path)
84            } else {
85                std::fs::remove_dir_all(&path)
86            };
87            if let Err(e) = e {
88                tracing::warn!(path = %path.display(), "failed to remove downloads entry: {e:?}");
89            }
90        }
91
92        Ok(())
93    }
94
95    pub fn file<P: AsRef<Path>>(&self, rel_path: P) -> FileBuilder {
96        self.storage_dir.file(&rel_path)
97    }
98
99    pub fn unnamed_file(&self) -> UnnamedFileBuilder {
100        self.storage_dir.unnamed_file()
101    }
102}
103
104#[derive(Clone)]
105pub struct Dir(Arc<DirInner>);
106
107impl Dir {
108    /// Creates a new `Dir` instance.
109    /// If the `root` directory does not exist, it will be created.
110    pub fn new<P>(root: P) -> Result<Self>
111    where
112        P: AsRef<Path>,
113    {
114        std::fs::create_dir_all(root.as_ref())
115            .with_context(|| format!("failed to create {}", root.as_ref().display()))?;
116        Ok(Self(Arc::new(DirInner {
117            base_dir: root.as_ref().to_path_buf(),
118        })))
119    }
120
121    /// Creates a new `Dir` without creating the root directory tree
122    pub fn new_readonly<P>(root: P) -> Self
123    where
124        P: AsRef<Path>,
125    {
126        Self(Arc::new(DirInner {
127            base_dir: root.as_ref().to_path_buf(),
128        }))
129    }
130
131    pub fn path(&self) -> &Path {
132        &self.0.base_dir
133    }
134
135    pub fn create_if_not_exists(&self) -> std::io::Result<()> {
136        std::fs::create_dir_all(&self.0.base_dir)
137    }
138
139    pub fn create_dir_all<P: AsRef<Path>>(&self, rel_path: P) -> std::io::Result<()> {
140        std::fs::create_dir_all(self.0.base_dir.join(rel_path))
141    }
142
143    pub fn remove_file<P: AsRef<Path>>(&self, rel_path: P) -> std::io::Result<()> {
144        std::fs::remove_file(self.0.base_dir.join(rel_path))
145    }
146
147    pub fn file<P: AsRef<Path>>(&self, rel_path: P) -> FileBuilder {
148        FileBuilder {
149            path: self.0.base_dir.join(rel_path.as_ref()),
150            options: std::fs::OpenOptions::new(),
151            prealloc: None,
152        }
153    }
154
155    pub fn unnamed_file(&self) -> UnnamedFileBuilder {
156        UnnamedFileBuilder {
157            base_dir: self.0.base_dir.clone(),
158            prealloc: None,
159        }
160    }
161
162    /// Creates `Dir` instance for a subdirectory of the current one.
163    /// **Note**: The subdirectory will not be created if it does not exist.
164    /// Use `create_subdir` to create it.
165    pub fn subdir_readonly<P: AsRef<Path>>(&self, rel_path: P) -> Self {
166        Self(Arc::new(DirInner {
167            base_dir: self.0.base_dir.join(rel_path),
168        }))
169    }
170
171    /// Creates `Dir` instance for a subdirectory of the current one.
172    /// The subdirectory will be created if it does not exist.
173    pub fn create_subdir<P: AsRef<Path>>(&self, rel_path: P) -> Result<Self> {
174        Self::new(self.0.base_dir.join(rel_path))
175    }
176
177    pub fn file_exists<P: AsRef<Path>>(&self, rel_path: P) -> bool {
178        self.path().join(rel_path).is_file()
179    }
180
181    pub fn entries(&self) -> std::io::Result<std::fs::ReadDir> {
182        std::fs::read_dir(&self.0.base_dir)
183    }
184}
185
186struct DirInner {
187    base_dir: PathBuf,
188}
189
190#[derive(Clone)]
191pub struct FileBuilder {
192    path: PathBuf,
193    options: OpenOptions,
194    prealloc: Option<usize>,
195}
196
197impl FileBuilder {
198    pub fn with_extension<S: AsRef<OsStr>>(&self, extension: S) -> Self {
199        Self {
200            path: self.path.with_extension(extension),
201            options: self.options.clone(),
202            prealloc: self.prealloc,
203        }
204    }
205
206    pub fn open(&self) -> Result<File> {
207        let file = self
208            .options
209            .open(&self.path)
210            .with_context(|| format!("failed to open {}", self.path.display()))?;
211        if let Some(prealloc) = self.prealloc {
212            alloc_file(&file, prealloc)?;
213        }
214        Ok(file)
215    }
216
217    pub fn rename<P: AsRef<Path>>(&self, new_path: P) -> std::io::Result<()> {
218        let new_path = match self.path.parent() {
219            Some(parent) => Cow::Owned(parent.join(new_path)),
220            None => Cow::Borrowed(new_path.as_ref()),
221        };
222        std::fs::rename(&self.path, new_path)
223    }
224
225    pub fn exists(&self) -> bool {
226        std::fs::metadata(&self.path)
227            .ok()
228            .map(|m| m.is_file())
229            .unwrap_or_default()
230    }
231
232    pub fn append(&mut self, append: bool) -> &mut Self {
233        self.options.append(append);
234        self
235    }
236
237    pub fn create(&mut self, create: bool) -> &mut Self {
238        self.options.create(create);
239        self
240    }
241
242    pub fn create_new(&mut self, create_new: bool) -> &mut Self {
243        self.options.create_new(create_new);
244        self
245    }
246
247    pub fn read(&mut self, read: bool) -> &mut Self {
248        self.options.read(read);
249        self
250    }
251
252    pub fn truncate(&mut self, truncate: bool) -> &mut Self {
253        self.options.truncate(truncate);
254        self
255    }
256
257    pub fn write(&mut self, write: bool) -> &mut Self {
258        self.options.write(write);
259        self
260    }
261
262    pub fn prealloc(&mut self, prealloc: usize) -> &mut Self {
263        self.prealloc = Some(prealloc);
264        self
265    }
266
267    pub fn path(&self) -> &Path {
268        &self.path
269    }
270}
271
272pub struct UnnamedFileBuilder {
273    base_dir: PathBuf,
274    prealloc: Option<usize>,
275}
276
277impl UnnamedFileBuilder {
278    pub fn open(self) -> Result<File> {
279        let file = tempfile::tempfile_in(&self.base_dir)?;
280        if let Some(prealloc) = self.prealloc {
281            file.set_len(prealloc as u64)?;
282        }
283
284        Ok(file)
285    }
286
287    pub fn open_as_mapped_mut(&self) -> Result<MappedFileMut> {
288        let file = tempfile::tempfile_in(&self.base_dir).with_context(|| {
289            format!("failed to create a tempfile in {}", self.base_dir.display())
290        })?;
291
292        if let Some(prealloc) = self.prealloc {
293            #[cfg(target_os = "linux")]
294            alloc_file(&file, prealloc)?;
295
296            file.set_len(prealloc as u64)?;
297        } else {
298            anyhow::bail!("prealloc is required for mapping unnamed files");
299        }
300
301        MappedFileMut::from_existing_file(file).map_err(Into::into)
302    }
303
304    pub fn prealloc(&mut self, prealloc: usize) -> &mut Self {
305        self.prealloc = Some(prealloc);
306        self
307    }
308}
309
310#[cfg(not(target_os = "macos"))]
311fn alloc_file(file: &File, len: usize) -> std::io::Result<()> {
312    let res = unsafe { libc::posix_fallocate(file.as_raw_fd(), 0, len as i64) };
313    if res == 0 {
314        Ok(())
315    } else {
316        Err(std::io::Error::last_os_error())
317    }
318}
319
320#[cfg(target_os = "macos")]
321pub fn alloc_file(file: &File, len: usize) -> std::io::Result<()> {
322    let res = unsafe { libc::ftruncate(file.as_raw_fd(), len as i64) };
323    if res < 0 {
324        Err(std::io::Error::last_os_error())
325    } else {
326        Ok(())
327    }
328}