polars_io/file_cache/
file_lock.rs

1use std::fs::{File, OpenOptions};
2use std::path::Path;
3
4use fs4::fs_std::FileExt;
5
6/// Note: this creates the file if it does not exist when acquiring locks.
7pub(super) struct FileLock<T: AsRef<Path>>(T);
8pub(super) struct FileLockSharedGuard(File);
9pub(super) struct FileLockExclusiveGuard(File);
10
11/// Trait to specify a file is lock-guarded without needing a particular type of
12/// guard (i.e. shared/exclusive).
13pub(super) trait FileLockAnyGuard:
14    std::ops::Deref<Target = File> + std::ops::DerefMut<Target = File>
15{
16    const IS_EXCLUSIVE: bool;
17}
18impl FileLockAnyGuard for FileLockSharedGuard {
19    const IS_EXCLUSIVE: bool = false;
20}
21impl FileLockAnyGuard for FileLockExclusiveGuard {
22    const IS_EXCLUSIVE: bool = true;
23}
24
25impl<T: AsRef<Path>> From<T> for FileLock<T> {
26    fn from(path: T) -> Self {
27        Self(path)
28    }
29}
30
31impl<T: AsRef<Path>> FileLock<T> {
32    pub(super) fn acquire_shared(&self) -> Result<FileLockSharedGuard, std::io::Error> {
33        let file = OpenOptions::new()
34            .create(true)
35            .truncate(false)
36            .read(true)
37            .write(true)
38            .open(self.0.as_ref())?;
39        FileExt::lock_shared(&file).map(|_| FileLockSharedGuard(file))
40    }
41
42    pub(super) fn acquire_exclusive(&self) -> Result<FileLockExclusiveGuard, std::io::Error> {
43        let file = OpenOptions::new()
44            .create(true)
45            .truncate(false)
46            .read(true)
47            .write(true)
48            .open(self.0.as_ref())?;
49        file.lock_exclusive().map(|_| FileLockExclusiveGuard(file))
50    }
51}
52
53impl std::ops::Deref for FileLockSharedGuard {
54    type Target = File;
55
56    fn deref(&self) -> &Self::Target {
57        &self.0
58    }
59}
60
61impl std::ops::DerefMut for FileLockSharedGuard {
62    fn deref_mut(&mut self) -> &mut Self::Target {
63        &mut self.0
64    }
65}
66
67impl Drop for FileLockSharedGuard {
68    fn drop(&mut self) {
69        FileExt::unlock(&self.0).unwrap();
70    }
71}
72
73impl std::ops::Deref for FileLockExclusiveGuard {
74    type Target = File;
75
76    fn deref(&self) -> &Self::Target {
77        &self.0
78    }
79}
80
81impl std::ops::DerefMut for FileLockExclusiveGuard {
82    fn deref_mut(&mut self) -> &mut Self::Target {
83        &mut self.0
84    }
85}
86
87impl Drop for FileLockExclusiveGuard {
88    fn drop(&mut self) {
89        FileExt::unlock(&self.0).unwrap();
90    }
91}