1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
|
//! Filesystem + JSON implementation of StateMgr.
use crate::{Error, LockStatus, Result, StateMgr};
use serde::{de::DeserializeOwned, Serialize};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
#[cfg(target_family = "unix")]
use std::os::unix::fs::DirBuilderExt;
/// Implementation of StateMgr that stores state as JSON files on disk.
///
/// # Locking
///
/// This manager uses a lock file to determine whether it's allowed to
/// write to the disk. Only one process should write to the disk at
/// a time, though any number may read from the disk.
///
/// By default, every `FsStateMgr` starts out unlocked, and only able
/// to read. Use [`FsStateMgr::try_lock()`] to lock it.
///
/// # Limitations
///
/// 1) This manager only accepts objects that can be serialized as
/// JSON documents. Some types (like maps with non-string keys) can't
/// be serialized as JSON.
///
/// 2) This manager normalizes keys to an fs-safe format before saving
/// data with them. This keeps you from accidentally creating or
/// reading files elsewhere in the filesystem, but it doesn't prevent
/// collisions when two keys collapse to the same fs-safe filename.
/// Therefore, you should probably only use ascii keys that are
/// fs-safe on all systems.
///
/// NEVER use user-controlled or remote-controlled data for your keys.
#[derive(Clone, Debug)]
pub struct FsStateMgr {
/// Inner reference-counted object.
inner: Arc<FsStateMgrInner>,
}
/// Inner reference-counted object, used by `FsStateMgr`.
#[derive(Debug)]
struct FsStateMgrInner {
/// Directory in which we store state files.
statepath: PathBuf,
/// Lockfile to achieve exclusive access to state files.
lockfile: Mutex<fslock::LockFile>,
}
impl FsStateMgr {
/// Construct a new `FsStateMgr` to store data in `path`.
///
/// This function will try to create `path` if it does not already
/// exist.
pub fn from_path<P: AsRef<Path>>(path: P) -> Result<Self> {
let path = path.as_ref();
let statepath = path.join("state");
let lockpath = path.join("state.lock");
{
let mut builder = std::fs::DirBuilder::new();
#[cfg(target_family = "unix")]
builder.mode(0o700);
builder.recursive(true).create(&statepath)?;
}
let lockfile = Mutex::new(fslock::LockFile::open(&lockpath)?);
Ok(FsStateMgr {
inner: Arc::new(FsStateMgrInner {
statepath,
lockfile,
}),
})
}
/// Return a filename to use for storing data with `key`.
///
/// See "Limitations" section on [`FsStateMgr`] for caveats.
fn filename(&self, key: &str) -> PathBuf {
self.inner
.statepath
.join(sanitize_filename::sanitize(key) + ".json")
}
}
impl StateMgr for FsStateMgr {
fn can_store(&self) -> bool {
let lockfile = self
.inner
.lockfile
.lock()
.expect("Poisoned lock on state lockfile");
lockfile.owns_lock()
}
fn try_lock(&self) -> Result<LockStatus> {
let mut lockfile = self
.inner
.lockfile
.lock()
.expect("Poisoned lock on state lockfile");
if lockfile.owns_lock() {
Ok(LockStatus::AlreadyHeld)
} else if lockfile.try_lock()? {
Ok(LockStatus::NewlyAcquired)
} else {
Ok(LockStatus::NoLock)
}
}
fn load<D>(&self, key: &str) -> Result<Option<D>>
where
D: DeserializeOwned,
{
let fname = self.filename(key);
let string = match std::fs::read_to_string(fname) {
Ok(s) => s,
Err(e) => {
if e.kind() == std::io::ErrorKind::NotFound {
return Ok(None);
} else {
return Err(e.into());
}
}
};
Ok(Some(serde_json::from_str(&string)?))
}
fn store<S>(&self, key: &str, val: &S) -> Result<()>
where
S: Serialize,
{
if !self.can_store() {
return Err(Error::NoLock);
}
let fname = self.filename(key);
let output = serde_json::to_string_pretty(val)?;
let fname_tmp = fname.with_extension("tmp");
std::fs::write(&fname_tmp, (&output).as_bytes())?;
std::fs::rename(fname_tmp, fname)?;
Ok(())
}
}
#[cfg(test)]
mod test {
#![allow(clippy::unwrap_used)]
use super::*;
use std::collections::HashMap;
#[test]
fn simple() -> Result<()> {
let dir = tempfile::TempDir::new().unwrap();
let store = FsStateMgr::from_path(dir.path())?;
assert_eq!(store.try_lock()?, LockStatus::NewlyAcquired);
let stuff: HashMap<_, _> = vec![("hello".to_string(), "world".to_string())]
.into_iter()
.collect();
store.store("xyz", &stuff)?;
let stuff2: Option<HashMap<String, String>> = store.load("xyz")?;
let nothing: Option<HashMap<String, String>> = store.load("abc")?;
assert_eq!(Some(stuff), stuff2);
assert!(nothing.is_none());
drop(store); // Do this to release the fs lock.
let store = FsStateMgr::from_path(dir.path())?;
let stuff3: Option<HashMap<String, String>> = store.load("xyz")?;
assert_eq!(stuff2, stuff3);
let stuff4: HashMap<_, _> = vec![("greetings".to_string(), "humans".to_string())]
.into_iter()
.collect();
assert!(matches!(store.store("xyz", &stuff4), Err(Error::NoLock)));
assert_eq!(store.try_lock()?, LockStatus::NewlyAcquired);
store.store("xyz", &stuff4)?;
let stuff5: Option<HashMap<String, String>> = store.load("xyz")?;
assert_eq!(Some(stuff4), stuff5);
Ok(())
}
}
|
