Skip to main content

automerge_persistent_fs/
lib.rs

1use std::{
2    collections::HashMap,
3    fs,
4    os::unix::prelude::OsStrExt,
5    path::{Path, PathBuf},
6};
7
8use automerge::ActorId;
9use automerge_persistent::{Persister, StoredSizes};
10#[cfg(feature = "async")]
11use futures::{Future, FutureExt, TryStreamExt};
12use hex::FromHexError;
13
14#[derive(Debug)]
15pub struct FsPersister {
16    changes_path: PathBuf,
17    doc_path: PathBuf,
18    sync_states_path: PathBuf,
19    cache: FsPersisterCache,
20    sizes: StoredSizes,
21}
22
23#[derive(Debug)]
24pub struct FsPersisterCache {
25    changes: HashMap<(ActorId, u64), Vec<u8>>,
26    document: Option<Vec<u8>>,
27    sync_states: HashMap<Vec<u8>, Vec<u8>>,
28}
29
30impl FsPersisterCache {
31    fn flush_changes(&mut self, changes_path: PathBuf) -> Result<usize, std::io::Error> {
32        let mut flushed = 0;
33        for ((a, s), c) in self.changes.drain() {
34            fs::write(make_changes_path(&changes_path, &a, s), &c)?;
35            flushed += c.len();
36        }
37        Ok(flushed)
38    }
39
40    fn flush_document(&mut self, doc_path: PathBuf) -> Result<usize, std::io::Error> {
41        let mut flushed = 0;
42        if let Some(data) = self.document.take() {
43            fs::write(&doc_path, &data)?;
44            flushed = data.len();
45        }
46        Ok(flushed)
47    }
48
49    fn flush_sync_states(&mut self, sync_states_path: PathBuf) -> Result<usize, std::io::Error> {
50        let mut flushed = 0;
51        for (peer_id, sync_state) in self.sync_states.drain() {
52            fs::write(make_peer_path(&sync_states_path, &peer_id), &sync_state)?;
53            flushed += sync_state.len();
54        }
55        Ok(flushed)
56    }
57
58    #[cfg(feature = "async")]
59    async fn flush_changes_async(
60        &mut self,
61        changes_path: PathBuf,
62    ) -> Result<usize, std::io::Error> {
63        let futs = futures::stream::FuturesUnordered::new();
64        for ((a, s), c) in self.changes.drain() {
65            let len = c.len();
66            futs.push(
67                tokio::fs::write(make_changes_path(&changes_path, &a, s), c).map(move |_| Ok(len)),
68            );
69        }
70        let res: Result<Vec<usize>, std::io::Error> = futs.try_collect().await;
71        Ok(res?.iter().sum())
72    }
73
74    #[cfg(feature = "async")]
75    async fn flush_document_async(&mut self, doc_path: PathBuf) -> Result<usize, std::io::Error> {
76        let mut flushed = 0;
77        if let Some(data) = self.document.take() {
78            tokio::fs::write(&doc_path, &data).await?;
79            flushed = data.len();
80        }
81        Ok(flushed)
82    }
83
84    #[cfg(feature = "async")]
85    async fn flush_sync_states_async(
86        &mut self,
87        sync_states_path: PathBuf,
88    ) -> Result<usize, std::io::Error> {
89        let futs = futures::stream::FuturesUnordered::new();
90        for (peer_id, sync_state) in self.sync_states.drain() {
91            let len = sync_state.len();
92            futs.push(
93                tokio::fs::write(make_peer_path(&sync_states_path, &peer_id), sync_state)
94                    .map(move |_| Ok(len)),
95            );
96        }
97        let res: Result<Vec<usize>, std::io::Error> = futs.try_collect().await;
98        Ok(res?.iter().sum())
99    }
100
101    #[cfg(feature = "async")]
102    pub async fn flush_async(
103        &mut self,
104        doc_path: PathBuf,
105        changes_path: PathBuf,
106        sync_states_path: PathBuf,
107    ) -> Result<usize, std::io::Error> {
108        let mut flushed = 0;
109        flushed += self.flush_document_async(doc_path).await?;
110        flushed += self.flush_changes_async(changes_path).await?;
111        flushed += self.flush_sync_states_async(sync_states_path).await?;
112        Ok(flushed)
113    }
114
115    pub fn flush(
116        &mut self,
117        doc_path: PathBuf,
118        changes_path: PathBuf,
119        sync_states_path: PathBuf,
120    ) -> Result<usize, std::io::Error> {
121        let mut flushed = 0;
122        flushed += self.flush_document(doc_path)?;
123        flushed += self.flush_changes(changes_path)?;
124        flushed += self.flush_sync_states(sync_states_path)?;
125        Ok(flushed)
126    }
127
128    fn drain_clone(&mut self) -> Self {
129        Self {
130            changes: self.changes.drain().collect(),
131            document: self.document.take(),
132            sync_states: self.sync_states.drain().collect(),
133        }
134    }
135}
136
137/// Possible errors from persisting.
138#[derive(Debug, thiserror::Error)]
139pub enum FsPersisterError {
140    #[error(transparent)]
141    Io(#[from] std::io::Error),
142    #[error(transparent)]
143    Hex(#[from] FromHexError),
144}
145
146const CHANGES_DIR: &str = "changes";
147const DOC_FILE: &str = "doc";
148const SYNC_DIR: &str = "sync";
149
150impl FsPersister {
151    pub fn new<R: AsRef<Path>, P: AsRef<Path>>(
152        root: R,
153        prefix: P,
154    ) -> Result<Self, FsPersisterError> {
155        let root_path = root.as_ref().join(&prefix);
156        fs::create_dir_all(&root_path)?;
157
158        let changes_path = root_path.join(CHANGES_DIR);
159        if fs::metadata(&changes_path).is_err() {
160            fs::create_dir(&changes_path)?;
161        }
162
163        let doc_path = root_path.join(DOC_FILE);
164
165        let sync_states_path = root_path.join(SYNC_DIR);
166        if fs::metadata(&sync_states_path).is_err() {
167            fs::create_dir(&sync_states_path)?;
168        }
169
170        let mut s = Self {
171            changes_path,
172            doc_path,
173            sync_states_path,
174            cache: FsPersisterCache {
175                changes: HashMap::new(),
176                document: None,
177                sync_states: HashMap::new(),
178            },
179            sizes: StoredSizes::default(),
180        };
181
182        s.sizes.changes = s.get_changes()?.iter().map(|v| v.len() as u64).sum();
183        s.sizes.document = s.get_document()?.unwrap_or_default().len() as u64;
184        s.sizes.sync_states = s
185            .get_peer_ids()?
186            .iter()
187            .map(|id| {
188                s.get_sync_state(id)
189                    .map(|o| o.unwrap_or_default().len() as u64)
190            })
191            .collect::<Result<Vec<u64>, _>>()?
192            .iter()
193            .sum();
194
195        Ok(s)
196    }
197
198    #[cfg(feature = "async")]
199    pub fn flush_cache_async(&mut self) -> impl Future<Output = Result<usize, std::io::Error>> {
200        let doc_path = self.doc_path.clone();
201        let changes_path = self.changes_path.clone();
202        let sync_states_path = self.sync_states_path.clone();
203        let mut cache = self.cache.drain_clone();
204        async move {
205            cache
206                .flush_async(doc_path, changes_path, sync_states_path)
207                .await
208        }
209    }
210
211    pub fn load<R: AsRef<Path>, P: AsRef<Path>>(
212        root: R,
213        prefix: P,
214    ) -> Result<Option<Self>, FsPersisterError> {
215        if !root.as_ref().join(&prefix).exists() {
216            return Ok(None);
217        }
218        let doc = Self::new(root, prefix)?;
219        Ok(Some(doc))
220    }
221}
222
223fn make_changes_path<P: AsRef<Path>>(changes_path: P, actor_id: &ActorId, seq: u64) -> PathBuf {
224    changes_path
225        .as_ref()
226        .join(format!("{}-{}", actor_id.to_hex_string(), seq))
227}
228
229fn make_peer_path<P: AsRef<Path>>(sync_states_path: P, peer_id: &[u8]) -> PathBuf {
230    sync_states_path.as_ref().join(hex::encode(peer_id))
231}
232
233impl Persister for FsPersister {
234    type Error = FsPersisterError;
235
236    fn get_changes(&self) -> Result<Vec<Vec<u8>>, Self::Error> {
237        fs::read_dir(&self.changes_path)?
238            .filter_map(|entry| {
239                if let Ok((Ok(file_type), path)) =
240                    entry.map(|entry| (entry.file_type(), entry.path()))
241                {
242                    if file_type.is_file() {
243                        Some(fs::read(path).map_err(FsPersisterError::from))
244                    } else {
245                        None
246                    }
247                } else {
248                    None
249                }
250            })
251            .collect()
252    }
253
254    fn insert_changes(&mut self, changes: Vec<(ActorId, u64, Vec<u8>)>) -> Result<(), Self::Error> {
255        for (a, s, c) in changes {
256            self.sizes.changes += c.len() as u64;
257            if let Some(old) = self.cache.changes.insert((a, s), c) {
258                self.sizes.changes -= old.len() as u64;
259            }
260        }
261        Ok(())
262    }
263
264    fn remove_changes(&mut self, changes: Vec<(&ActorId, u64)>) -> Result<(), Self::Error> {
265        for (a, s) in changes {
266            if let Some(old) = self.cache.changes.remove(&(a.clone(), s)) {
267                // not flushed yet
268                self.sizes.changes -= old.len() as u64;
269                continue;
270            }
271
272            let path = make_changes_path(&self.changes_path, a, s);
273            if let Ok(meta) = fs::metadata(&path) {
274                if meta.is_file() {
275                    fs::remove_file(&path)?;
276                    self.sizes.changes -= meta.len();
277                }
278            }
279        }
280        Ok(())
281    }
282
283    fn get_document(&self) -> Result<Option<Vec<u8>>, Self::Error> {
284        if let Some(ref doc) = self.cache.document {
285            return Ok(Some(doc.clone()));
286        }
287        if fs::metadata(&self.doc_path).is_ok() {
288            return Ok(fs::read(&self.doc_path).map(|v| if v.is_empty() { None } else { Some(v) })?);
289        }
290        Ok(None)
291    }
292
293    fn set_document(&mut self, data: Vec<u8>) -> Result<(), Self::Error> {
294        self.sizes.document = data.len() as u64;
295        self.cache.document = Some(data);
296        Ok(())
297    }
298
299    fn get_sync_state(&self, peer_id: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
300        if let Some(sync_state) = self.cache.sync_states.get(peer_id) {
301            return Ok(Some(sync_state.clone()));
302        }
303        let path = make_peer_path(&self.sync_states_path, peer_id);
304        if fs::metadata(&path).is_ok() {
305            return Ok(fs::read(&path).map(|v| if v.is_empty() { None } else { Some(v) })?);
306        }
307        Ok(None)
308    }
309
310    fn set_sync_state(&mut self, peer_id: Vec<u8>, sync_state: Vec<u8>) -> Result<(), Self::Error> {
311        self.sizes.sync_states += sync_state.len() as u64;
312        if let Some(old) = self.cache.sync_states.insert(peer_id, sync_state) {
313            self.sizes.sync_states -= old.len() as u64;
314        }
315        Ok(())
316    }
317
318    fn remove_sync_states(&mut self, peer_ids: &[&[u8]]) -> Result<(), Self::Error> {
319        for peer_id in peer_ids {
320            if let Some(old) = self.cache.sync_states.remove(*peer_id) {
321                // not flushed yet
322                self.sizes.sync_states -= old.len() as u64;
323                continue;
324            }
325            let path = make_peer_path(&self.sync_states_path, peer_id);
326            if let Ok(meta) = fs::metadata(&path) {
327                if meta.is_file() {
328                    fs::remove_file(&path)?;
329                    self.sizes.sync_states -= meta.len();
330                }
331            }
332        }
333        Ok(())
334    }
335
336    fn get_peer_ids(&self) -> Result<Vec<Vec<u8>>, Self::Error> {
337        fs::read_dir(&self.sync_states_path)?
338            .filter_map(|entry| {
339                if let Ok((Ok(file_type), path)) =
340                    entry.map(|entry| (entry.file_type(), entry.path()))
341                {
342                    if file_type.is_file() {
343                        Some(
344                            hex::decode(path.file_name().unwrap().as_bytes())
345                                .map_err(FsPersisterError::from),
346                        )
347                    } else {
348                        None
349                    }
350                } else {
351                    None
352                }
353            })
354            .collect()
355    }
356
357    fn sizes(&self) -> StoredSizes {
358        self.sizes.clone()
359    }
360
361    fn flush(&mut self) -> Result<usize, Self::Error> {
362        self.cache
363            .drain_clone()
364            .flush(
365                self.doc_path.clone(),
366                self.changes_path.clone(),
367                self.sync_states_path.clone(),
368            )
369            .map_err(FsPersisterError::from)
370    }
371}