Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
pola-rs
GitHub Repository: pola-rs/polars
Path: blob/main/crates/polars-io/src/file_cache/file_lock.rs
6939 views
1
use std::fs::{File, OpenOptions};
2
use std::path::Path;
3
4
use fs4::fs_std::FileExt;
5
6
/// Note: this creates the file if it does not exist when acquiring locks.
7
pub(super) struct FileLock<T: AsRef<Path>>(T);
8
pub(super) struct FileLockSharedGuard(File);
9
pub(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).
13
pub(super) trait FileLockAnyGuard:
14
std::ops::Deref<Target = File> + std::ops::DerefMut<Target = File>
15
{
16
const IS_EXCLUSIVE: bool;
17
}
18
impl FileLockAnyGuard for FileLockSharedGuard {
19
const IS_EXCLUSIVE: bool = false;
20
}
21
impl FileLockAnyGuard for FileLockExclusiveGuard {
22
const IS_EXCLUSIVE: bool = true;
23
}
24
25
impl<T: AsRef<Path>> From<T> for FileLock<T> {
26
fn from(path: T) -> Self {
27
Self(path)
28
}
29
}
30
31
impl<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
53
impl std::ops::Deref for FileLockSharedGuard {
54
type Target = File;
55
56
fn deref(&self) -> &Self::Target {
57
&self.0
58
}
59
}
60
61
impl std::ops::DerefMut for FileLockSharedGuard {
62
fn deref_mut(&mut self) -> &mut Self::Target {
63
&mut self.0
64
}
65
}
66
67
impl Drop for FileLockSharedGuard {
68
fn drop(&mut self) {
69
FileExt::unlock(&self.0).unwrap();
70
}
71
}
72
73
impl std::ops::Deref for FileLockExclusiveGuard {
74
type Target = File;
75
76
fn deref(&self) -> &Self::Target {
77
&self.0
78
}
79
}
80
81
impl std::ops::DerefMut for FileLockExclusiveGuard {
82
fn deref_mut(&mut self) -> &mut Self::Target {
83
&mut self.0
84
}
85
}
86
87
impl Drop for FileLockExclusiveGuard {
88
fn drop(&mut self) {
89
FileExt::unlock(&self.0).unwrap();
90
}
91
}
92
93