journal_sdk/
persistor.rs

1use crate::config::Config;
2use once_cell::sync::Lazy;
3use rand::prelude::SliceRandom;
4use rand::RngCore;
5use sha2::{Digest, Sha256};
6use std::collections::HashMap;
7use std::sync::{Arc, Mutex};
8
9use rocksdb::{ColumnFamilyDescriptor, Options, DB};
10
11pub static PERSISTOR: Lazy<Box<dyn Persistor + Send + Sync>> = Lazy::new(|| {
12    let config = Config::new();
13    match config.database.as_str() {
14        "" => Box::new(MemoryPersistor::new()),
15        path => Box::new(DatabasePersistor::new(path)),
16    }
17});
18
19/// Size, in bytes, of the global hash algorithm (currently SHA-256)
20pub const SIZE: usize = 32;
21
22/// Byte array describing a hash pointer (currently SHA-256)
23pub type Word = [u8; SIZE];
24
25#[allow(dead_code)]
26#[derive(Debug)]
27pub struct PersistorAccessError(pub String);
28
29pub trait Persistor {
30    fn root_list(&self) -> Vec<Word>;
31    fn root_new(&self, handle: Word, root: Word) -> Result<Word, PersistorAccessError>;
32    fn root_temp(&self, root: Word) -> Result<Word, PersistorAccessError>;
33    fn root_get(&self, handle: Word) -> Result<Word, PersistorAccessError>;
34    fn root_set(&self, handle: Word, old: Word, new: Word) -> Result<Word, PersistorAccessError>;
35    fn root_delete(&self, handle: Word) -> Result<(), PersistorAccessError>;
36    fn branch_set(
37        &self,
38        left: Word,
39        right: Word,
40        digest: Word,
41    ) -> Result<Word, PersistorAccessError>;
42    fn branch_get(&self, branch: Word) -> Result<(Word, Word, Word), PersistorAccessError>;
43    fn leaf_set(&self, content: Vec<u8>) -> Result<Word, PersistorAccessError>;
44    fn leaf_get(&self, leaf: Word) -> Result<Vec<u8>, PersistorAccessError>;
45    fn stump_set(&self, digest: Word) -> Result<Word, PersistorAccessError>;
46    fn stump_get(&self, stump: Word) -> Result<Word, PersistorAccessError>;
47}
48
49#[derive(Clone)]
50pub struct MemoryPersistor {
51    roots: Arc<Mutex<HashMap<Word, (Word, bool)>>>,
52    branches: Arc<Mutex<HashMap<Word, (Word, Word, Word)>>>,
53    leaves: Arc<Mutex<HashMap<Word, Vec<u8>>>>,
54    stumps: Arc<Mutex<HashMap<Word, Word>>>,
55    references: Arc<Mutex<HashMap<Word, usize>>>,
56}
57
58impl MemoryPersistor {
59    pub fn new() -> Self {
60        Self {
61            roots: Arc::new(Mutex::new(HashMap::new())),
62            branches: Arc::new(Mutex::new(HashMap::new())),
63            leaves: Arc::new(Mutex::new(HashMap::new())),
64            stumps: Arc::new(Mutex::new(HashMap::new())),
65            references: Arc::new(Mutex::new(HashMap::new())),
66        }
67    }
68
69    fn reference_increment(&self, node: Word) {
70        let mut references = self.references.lock().expect("Failed to lock references");
71        match references.get(&node) {
72            Some(count) => {
73                let count_ = *count;
74                references.insert(node, count_ + 1);
75            }
76            None => {
77                references.insert(node, 1);
78            }
79        };
80    }
81
82    fn reference_decrement(&self, node: Word) {
83        let mut references = self.references.lock().expect("Failed to lock references");
84        match references.get(&node) {
85            Some(count_old) => {
86                let count_new = *count_old - 1;
87                if count_new > 0 {
88                    references.insert(node, count_new);
89                } else {
90                    references.remove(&node);
91                    let mut branches = self.branches.lock().expect("Failed to lock branches");
92                    if let Some((left, right, _)) = branches.get(&node) {
93                        let left_ = *left;
94                        let right_ = *right;
95                        branches.remove(&node);
96                        drop(references);
97                        drop(branches);
98                        self.reference_decrement(left_);
99                        self.reference_decrement(right_);
100                    } else {
101                        let mut leaves = self.leaves.lock().expect("Failed to lock leaves");
102                        let mut stumps = self.stumps.lock().expect("Failed to lock stumps");
103                        if let Some(_) = leaves.get(&node) {
104                            leaves.remove(&node);
105                        } else if let Some(_) = stumps.get(&node) {
106                            stumps.remove(&node);
107                        }
108                    }
109                }
110            }
111            None => {}
112        };
113    }
114}
115
116impl Persistor for MemoryPersistor {
117    fn root_list(&self) -> Vec<Word> {
118        let mut keys: Vec<Word> = self
119            .roots
120            .lock()
121            .expect("Failed to get locked roots")
122            .iter()
123            .filter(|&(_, &(_, is_persistent))| is_persistent)
124            .map(|(key, _)| key)
125            .cloned()
126            .collect();
127        keys.sort();
128        keys
129    }
130
131    fn root_new(&self, handle: Word, root: Word) -> Result<Word, PersistorAccessError> {
132        let mut roots = self.roots.lock().expect("Failed to lock roots map");
133        match roots.get(&handle) {
134            Some(_) => Err(PersistorAccessError(format!(
135                "Handle {:?} already exists",
136                handle
137            ))),
138            None => {
139                self.reference_increment(root);
140                roots.insert(handle, (root, true));
141                Ok(handle)
142            }
143        }
144    }
145
146    fn root_temp(&self, root: Word) -> Result<Word, PersistorAccessError> {
147        let mut roots = self.roots.lock().expect("Failed to lock roots map");
148        let mut handle: Word = [0 as u8; 32];
149        rand::thread_rng().fill_bytes(&mut handle);
150        match roots.get(&handle) {
151            Some(_) => Err(PersistorAccessError(format!(
152                "Handle {:?} already exists",
153                handle
154            ))),
155            None => {
156                self.reference_increment(root);
157                roots.insert(handle, (root, false));
158                Ok(handle)
159            }
160        }
161    }
162
163    fn root_get(&self, handle: Word) -> Result<Word, PersistorAccessError> {
164        match self
165            .roots
166            .lock()
167            .expect("Failed to lock roots map")
168            .get(&handle)
169        {
170            Some((root, _)) => Ok(*root),
171            None => Err(PersistorAccessError(format!(
172                "Handle {:?} not found",
173                handle
174            ))),
175        }
176    }
177
178    fn root_set(&self, handle: Word, old: Word, new: Word) -> Result<Word, PersistorAccessError> {
179        let mut roots = self.roots.lock().expect("Failed to lock roots map");
180        match roots.get(&handle) {
181            Some((root, true)) if *root == old => {
182                self.reference_increment(new);
183                self.reference_decrement(old);
184                roots.insert(handle, (new, true));
185                Ok(handle)
186            }
187            Some((_, false)) => Err(PersistorAccessError(format!(
188                "Handle {:?} is temporary",
189                handle
190            ))),
191            Some((_, true)) => Err(PersistorAccessError(format!(
192                "Handle {:?} changed since compare",
193                handle
194            ))),
195            None => Err(PersistorAccessError(format!(
196                "Handle {:?} not found",
197                handle
198            ))),
199        }
200    }
201
202    fn root_delete(&self, handle: Word) -> Result<(), PersistorAccessError> {
203        let mut roots = self.roots.lock().expect("Failed to lock roots map");
204        match roots.get(&handle) {
205            Some((old, _)) => {
206                self.reference_decrement(*old);
207                roots.remove(&handle);
208                Ok(())
209            }
210            None => Err(PersistorAccessError(format!(
211                "Handle {:?} not found",
212                handle
213            ))),
214        }
215    }
216
217    fn branch_set(
218        &self,
219        left: Word,
220        right: Word,
221        digest: Word,
222    ) -> Result<Word, PersistorAccessError> {
223        let mut joined = [0 as u8; SIZE * 3];
224        joined[..SIZE].copy_from_slice(&left);
225        joined[SIZE..SIZE * 2].copy_from_slice(&right);
226        joined[SIZE * 2..].copy_from_slice(&digest);
227
228        let branch = Sha256::digest(joined);
229        self.branches
230            .lock()
231            .expect("Failed to lock branches map")
232            .insert(branch.into(), (left, right, digest));
233        self.reference_increment(left);
234        self.reference_increment(right);
235        Ok(Word::from(branch))
236    }
237
238    fn branch_get(&self, branch: Word) -> Result<(Word, Word, Word), PersistorAccessError> {
239        let branches = self.branches.lock().expect("Failed to lock branches map");
240        match branches.get(&branch) {
241            Some((left, right, digest)) => {
242                let mut joined = [0 as u8; SIZE * 3];
243                joined[..SIZE].copy_from_slice(left);
244                joined[SIZE..SIZE * 2].copy_from_slice(right);
245                joined[SIZE * 2..].copy_from_slice(digest);
246                assert!(Vec::from(branch) == Sha256::digest(joined).to_vec());
247                Ok((*left, *right, *digest))
248            }
249            None => Err(PersistorAccessError(format!(
250                "Branch {:?} not found",
251                branch
252            ))),
253        }
254    }
255
256    fn leaf_set(&self, content: Vec<u8>) -> Result<Word, PersistorAccessError> {
257        let leaf = Word::from(Sha256::digest(Sha256::digest(&content)));
258        self.leaves
259            .lock()
260            .expect("Failed to lock leaves map")
261            .insert(leaf, content);
262        Ok(leaf)
263    }
264
265    fn leaf_get(&self, leaf: Word) -> Result<Vec<u8>, PersistorAccessError> {
266        let leaves = self.leaves.lock().expect("Failed to lock leaves map");
267        match leaves.get(&leaf) {
268            Some(content) => {
269                assert!(Vec::from(leaf) == Sha256::digest(Sha256::digest(content)).to_vec());
270                Ok(content.to_vec())
271            }
272            None => Err(PersistorAccessError(format!("Leaf {:?} not found", leaf))),
273        }
274    }
275
276    fn stump_set(&self, digest: Word) -> Result<Word, PersistorAccessError> {
277        let stump = Sha256::digest(digest);
278        self.stumps
279            .lock()
280            .expect("Failed to lock stump map")
281            .insert(stump.into(), digest);
282        Ok(Word::from(stump))
283    }
284
285    fn stump_get(&self, stump: Word) -> Result<Word, PersistorAccessError> {
286        let stumps = self.stumps.lock().expect("Failed to lock stumps map");
287        match stumps.get(&stump) {
288            Some(digest) => {
289                assert!(Vec::from(stump) == Sha256::digest(Vec::from(digest)).to_vec());
290                Ok(*digest)
291            }
292            None => Err(PersistorAccessError(format!("Stump {:?} not found", stump))),
293        }
294    }
295}
296
297pub struct DatabasePersistor {
298    db: Mutex<DB>,
299}
300
301impl DatabasePersistor {
302    pub fn new(path: &str) -> Self {
303        let mut opts = Options::default();
304        opts.create_if_missing(true);
305        opts.create_missing_column_families(true);
306
307        let cfs = vec![
308            ColumnFamilyDescriptor::new("roots", Options::default()),
309            ColumnFamilyDescriptor::new("branches", Options::default()),
310            ColumnFamilyDescriptor::new("leaves", Options::default()),
311            ColumnFamilyDescriptor::new("stumps", Options::default()),
312            ColumnFamilyDescriptor::new("references", Options::default()),
313        ];
314        let persistor = Self {
315            db: Mutex::new(
316                DB::open_cf_descriptors(&opts, path, cfs).expect("Failed to open database"),
317            ),
318        };
319
320        // TODO: clear this due to memory leakage
321        {
322            let mut handles: Vec<Word> = Vec::new();
323            let db = persistor
324                .db
325                .lock()
326                .expect("Failed to acquire database lock");
327            let mut iter = db.raw_iterator_cf(
328                db.cf_handle("roots")
329                    .expect("Failed to get roots column family"),
330            );
331            iter.seek_to_first();
332            while iter.valid() {
333                if (*iter.value().expect("Failed to get iterator value"))[SIZE] == false as u8 {
334                    handles.push(
335                        (*iter.key().expect("Failed to get iterator key"))
336                            .try_into()
337                            .expect("Failed to convert key to Word"),
338                    );
339                }
340                iter.next();
341            }
342            for handle in handles {
343                db.delete_cf(
344                    db.cf_handle("roots")
345                        .expect("Failed to get roots column family"),
346                    handle,
347                )
348                .expect("Failed to delete value from roots");
349            }
350        }
351
352        persistor
353    }
354
355    fn reference_increment(&self, node: Word) {
356        let db = self.db.lock().expect("Failed to acquire db lock");
357        let references = db
358            .cf_handle("references")
359            .expect("Failed to get references handle");
360        match db.get_cf(references, node) {
361            Ok(Some(count)) => {
362                db.put_cf(
363                    references,
364                    node,
365                    (usize::from_ne_bytes(count.try_into().expect("Invalid count bytes")) + 1)
366                        .to_ne_bytes(),
367                )
368                .expect("Failed to increment reference count");
369            }
370            Ok(None) => db
371                .put_cf(references, node, (1 as usize).to_ne_bytes())
372                .expect("Failed to set initial reference count"),
373            Err(e) => {
374                panic! {"{}", e}
375            }
376        };
377    }
378
379    fn reference_decrement(&self, node: Word) {
380        let db = self.db.lock().expect("Failed to acquire db lock");
381        let branches = db
382            .cf_handle("branches")
383            .expect("Failed to get branches handle");
384        let leaves = db.cf_handle("leaves").expect("Failed to get leaves handle");
385        let stumps = db.cf_handle("stumps").expect("Failed to get stumps handle");
386        let references = db
387            .cf_handle("references")
388            .expect("Failed to get references handle");
389        match db
390            .get_cf(references, node)
391            .expect("Failed to get reference count")
392        {
393            Some(count_old) => {
394                let count_new =
395                    usize::from_ne_bytes(count_old.try_into().expect("Invalid count bytes")) - 1;
396                if count_new > 0 {
397                    db.put_cf(references, node, count_new.to_ne_bytes())
398                        .expect("Failed to update reference count");
399                } else {
400                    db.delete_cf(references, node)
401                        .expect("Failed to delete reference");
402                    if let Some(value) = db.get_cf(branches, node).expect("Failed to get branch") {
403                        let left = &value[..SIZE].try_into().expect("Invalid left node bytes");
404                        let right = &value[SIZE..SIZE * 2].try_into().expect("Invalid right node bytes");
405                        db.delete_cf(branches, node)
406                            .expect("Failed to delete branch");
407                        drop(db);
408                        self.reference_decrement(*left);
409                        self.reference_decrement(*right);
410                    } else {
411                        if let Some(_) = db.get_cf(leaves, node).expect("Failed to get leaf") {
412                            db.delete_cf(leaves, node).expect("Failed to delete leaf");
413                        } else if let Some(_) =
414                            db.get_cf(stumps, node).expect("Failed to get stump")
415                        {
416                            db.delete_cf(stumps, node).expect("Failed to delete stump");
417                        }
418                    }
419                }
420            }
421            None => {}
422        };
423    }
424}
425
426impl Persistor for DatabasePersistor {
427    fn root_list(&self) -> Vec<Word> {
428        let mut handles: Vec<Word> = Vec::new();
429        let db = self.db.lock().expect("Failed to acquire db lock");
430        let roots = db
431            .cf_handle("roots")
432            .expect("Failed to get roots column family");
433        let mut iter = db.raw_iterator_cf(roots);
434        iter.seek_to_first();
435        while iter.valid() {
436            if (*iter.value().expect("Failed to get iterator value"))[SIZE] != false as u8 {
437                handles.push(
438                    (*iter.key().expect("Failed to get iterator key"))
439                        .try_into()
440                        .expect("Failed to convert key to Word"),
441                );
442            }
443            iter.next();
444        }
445
446        handles.shuffle(&mut rand::thread_rng());
447        handles
448    }
449
450    fn root_new(&self, handle: Word, root: Word) -> Result<Word, PersistorAccessError> {
451        let mut root_marked = [0 as u8; SIZE + 1];
452        root_marked[..SIZE].copy_from_slice(&root);
453        root_marked[SIZE] = true as u8;
454
455        let db = self.db.lock().expect("Failed to acquire db lock");
456        let roots = db
457            .cf_handle("roots")
458            .expect("Failed to get roots column family");
459        match db.get_cf(roots, handle) {
460            Ok(Some(_)) => Err(PersistorAccessError(format!(
461                "Handle {:?} already exists",
462                handle
463            ))),
464            Ok(None) => {
465                db.put_cf(roots, handle, root_marked)
466                    .expect("Failed to put root in db");
467                drop(db);
468                self.reference_increment(root);
469                Ok(handle)
470            }
471            Err(e) => Err(PersistorAccessError(format!("{}", e))),
472        }
473    }
474
475    fn root_temp(&self, root: Word) -> Result<Word, PersistorAccessError> {
476        let mut root_marked = [0 as u8; SIZE + 1];
477        root_marked[..SIZE].copy_from_slice(&root);
478        root_marked[SIZE] = false as u8;
479
480        let mut handle: Word = [0 as u8; 32];
481        rand::thread_rng().fill_bytes(&mut handle);
482        let db = self.db.lock().expect("Failed to acquire db lock");
483        let roots = db
484            .cf_handle("roots")
485            .expect("Failed to get roots column family");
486        match db.get_cf(roots, handle) {
487            Ok(Some(_)) => Err(PersistorAccessError(format!(
488                "Handle {:?} already exists",
489                handle
490            ))),
491            Ok(None) => {
492                db.put_cf(roots, handle, root_marked)
493                    .expect("Failed to put root in db");
494                drop(db);
495                self.reference_increment(root);
496                Ok(handle)
497            }
498            Err(e) => Err(PersistorAccessError(format!("{}", e))),
499        }
500    }
501
502    fn root_get(&self, handle: Word) -> Result<Word, PersistorAccessError> {
503        let db = self.db.lock().expect("Failed to acquire db lock");
504        let roots = db.cf_handle("roots").expect("Failed to get roots handle");
505        match db.get_cf(roots, handle) {
506            Ok(Some(root_marked)) => Ok(((*root_marked)[..SIZE])
507                .try_into()
508                .expect("Invalid root size")),
509            Ok(None) => Err(PersistorAccessError(format!(
510                "Handle {:?} not found",
511                handle
512            ))),
513            Err(e) => Err(PersistorAccessError(format!("{}", e))),
514        }
515    }
516
517    fn root_set(&self, handle: Word, old: Word, new: Word) -> Result<Word, PersistorAccessError> {
518        let db = self.db.lock().expect("Failed to acquire db lock");
519        let roots = db.cf_handle("roots").expect("Failed to get roots handle");
520        match db.get_cf(roots, handle) {
521            Ok(Some(root_marked)) => match root_marked[SIZE] != false as u8 {
522                true => match root_marked[..SIZE] == old.to_vec() {
523                    true => {
524                        let mut new_marked = [0 as u8; SIZE + 1];
525                        new_marked[..SIZE].copy_from_slice(&new);
526                        new_marked[SIZE] = true as u8;
527                        db.put_cf(roots, handle, new_marked)
528                            .expect("Failed to put root");
529                        drop(db);
530                        self.reference_increment(new);
531                        self.reference_decrement(old);
532                        Ok(handle)
533                    }
534                    false => Err(PersistorAccessError(format!(
535                        "Handle {:?} changed since compare",
536                        handle
537                    ))),
538                },
539                false => Err(PersistorAccessError(format!(
540                    "Handle {:?} is temporary",
541                    handle
542                ))),
543            },
544            Ok(None) => Err(PersistorAccessError(format!(
545                "Handle {:?} not found",
546                handle
547            ))),
548            Err(e) => Err(PersistorAccessError(format!("{}", e))),
549        }
550    }
551
552    fn root_delete(&self, handle: Word) -> Result<(), PersistorAccessError> {
553        let db = self.db.lock().expect("Failed to acquire db lock");
554        let roots = db.cf_handle("roots").expect("Failed to get roots handle");
555        match db.get_cf(roots, handle) {
556            Ok(Some(root_marked)) => {
557                db.delete_cf(roots, handle).expect("Failed to delete root");
558                drop(db);
559                self.reference_decrement(
560                    root_marked[..SIZE].try_into().expect("Invalid root size"),
561                );
562                Ok(())
563            }
564            Ok(None) => Err(PersistorAccessError(format!(
565                "Handle {:?} not found",
566                handle
567            ))),
568            Err(e) => Err(PersistorAccessError(format!("{}", e))),
569        }
570    }
571
572    fn branch_set(
573        &self,
574        left: Word,
575        right: Word,
576        digest: Word,
577    ) -> Result<Word, PersistorAccessError> {
578        let mut joined = [0 as u8; SIZE * 3];
579        joined[..SIZE].copy_from_slice(&left);
580        joined[SIZE..SIZE * 2].copy_from_slice(&right);
581        joined[SIZE * 2..].copy_from_slice(&digest);
582
583        let branch = Sha256::digest(joined);
584
585        let db = self.db.lock().expect("Failed to acquire db lock");
586        let branches = db
587            .cf_handle("branches")
588            .expect("Failed to get branches handle");
589        db.put_cf(branches, branch, joined)
590            .expect("Failed to put branch");
591        drop(db);
592        self.reference_increment(left);
593        self.reference_increment(right);
594
595        Ok(Word::from(branch))
596    }
597
598    fn branch_get(&self, branch: Word) -> Result<(Word, Word, Word), PersistorAccessError> {
599        let db = self.db.lock().expect("Failed to acquire db lock");
600        let branches = db
601            .cf_handle("branches")
602            .expect("Failed to get branches handle");
603        match db.get_cf(branches, branch) {
604            Ok(Some(value)) => {
605                assert!(Vec::from(branch) == Sha256::digest(value.clone()).to_vec());
606                let left = &value[..SIZE].try_into().expect("Invalid left branch size");
607                let right = &value[SIZE..SIZE * 2]
608                    .try_into()
609                    .expect("Invalid right branch size");
610                let digest = &value[SIZE * 2..]
611                    .try_into()
612                    .expect("Invalid digest branch size");
613                Ok((*left, *right, *digest))
614            }
615            Ok(None) => Err(PersistorAccessError(format!(
616                "Branch {:?} not found",
617                branch
618            ))),
619            Err(e) => Err(PersistorAccessError(format!("{}", e))),
620        }
621    }
622
623    fn leaf_set(&self, content: Vec<u8>) -> Result<Word, PersistorAccessError> {
624        let leaf = Word::from(Sha256::digest(Sha256::digest(&content)));
625        let db = self.db.lock().expect("Failed to acquire db lock");
626        let leaves = db.cf_handle("leaves").expect("Failed to get leaves handle");
627        db.put_cf(leaves, leaf, content.clone())
628            .expect("Failed to put leaf");
629        Ok(leaf)
630    }
631
632    fn leaf_get(&self, leaf: Word) -> Result<Vec<u8>, PersistorAccessError> {
633        let db = self.db.lock().expect("Failed to acquire db lock");
634        let leaves = db.cf_handle("leaves").expect("Failed to get leaves handle");
635        match db.get_cf(leaves, leaf) {
636            Ok(Some(content)) => {
637                assert!(leaf == *Sha256::digest(Sha256::digest(content.clone())));
638                Ok(content.to_vec())
639            }
640            Ok(None) => Err(PersistorAccessError(format!("Leaf {:?} not found", leaf))),
641            Err(e) => Err(PersistorAccessError(format!("{}", e))),
642        }
643    }
644
645    fn stump_set(&self, digest: Word) -> Result<Word, PersistorAccessError> {
646        let stump = Word::from(Sha256::digest(Vec::from(digest)));
647        let db = self.db.lock().expect("Failed to acquire db lock");
648        let stumps = db.cf_handle("stumps").expect("Failed to get stumps handle");
649        db.put_cf(stumps, stump, digest)
650            .expect("Failed to put stump");
651        Ok(stump)
652    }
653
654    fn stump_get(&self, stump: Word) -> Result<Word, PersistorAccessError> {
655        let db = self.db.lock().expect("Failed to acquire db lock");
656        let stumps = db.cf_handle("stumps").expect("Failed to get stumps handle");
657        match db.get_cf(stumps, stump) {
658            Ok(Some(digest)) => {
659                assert!(stump == *Sha256::digest(digest.clone()));
660                Ok((&(*digest)[..SIZE])
661                    .try_into()
662                    .expect("Invalid left node bytes"))
663            }
664            Ok(None) => Err(PersistorAccessError(format!(
665                "Stumps {:?} not found",
666                stump
667            ))),
668            Err(e) => Err(PersistorAccessError(format!("{}", e))),
669        }
670    }
671}
672
673#[cfg(test)]
674mod tests {
675    use super::{DatabasePersistor, MemoryPersistor, Persistor, Word, SIZE};
676    use rocksdb::{IteratorMode, DB};
677    use std::fs;
678    use std::sync::Mutex;
679
680    fn test_persistence(persistor: Box<dyn Persistor>) {
681        let zeros: Word = [0 as u8; SIZE];
682
683        assert!(
684            persistor
685                .root_delete(
686                    persistor
687                        .root_temp(zeros)
688                        .expect("Failed to create temp root")
689                )
690                .expect("Failed to delete root")
691                == ()
692        );
693
694        assert!(
695            persistor
696                .root_get(
697                    persistor
698                        .root_set(
699                            persistor
700                                .root_new(zeros, zeros)
701                                .expect("Failed to create new root"),
702                            zeros,
703                            zeros,
704                        )
705                        .expect("Failed to set root"),
706                )
707                .expect("Failed to get root")
708                == zeros
709        );
710
711        assert!(
712            persistor
713                .branch_get(
714                    persistor
715                        .branch_set(zeros, zeros, zeros)
716                        .expect("Failed to set branch"),
717                )
718                .expect("Failed to get branch")
719                == (zeros, zeros, zeros)
720        );
721
722        assert!(
723            persistor
724                .leaf_get(persistor.leaf_set(vec!(0)).expect("Failed to set leaf"),)
725                .expect("Failed to get leaf")
726                == vec!(0)
727        );
728    }
729
730    #[test]
731    fn test_memory_persistence() {
732        test_persistence(Box::new(MemoryPersistor::new()));
733    }
734
735    #[test]
736    fn test_database_persistence() {
737        let db = ".test-database-persistence";
738        let _ = fs::remove_dir_all(db);
739        test_persistence(Box::new(DatabasePersistor::new(db)));
740        let _ = fs::remove_dir_all(db);
741    }
742
743    #[test]
744    fn test_memory_garbage() {
745        let persistor = MemoryPersistor::new();
746        let zeros: Word = [0 as u8; SIZE];
747        let handle: Word = [0 as u8; SIZE];
748
749        let leaf_0 = persistor.leaf_set(vec![0]).expect("Failed to set leaf 0");
750        let leaf_1 = persistor.leaf_set(vec![1]).expect("Failed to set leaf 1");
751        let leaf_2 = persistor.leaf_set(vec![2]).expect("Failed to set leaf 2");
752
753        let stump_0 = persistor
754            .stump_set([0 as u8; SIZE])
755            .expect("Failed to set stump 0");
756
757        let branch_a = persistor
758            .branch_set(leaf_0, leaf_1, zeros)
759            .expect("Failed to set branch A");
760        let branch_b = persistor
761            .branch_set(branch_a, leaf_2, zeros)
762            .expect("Failed to set branch B");
763        let branch_c = persistor
764            .branch_set(branch_b, stump_0, zeros)
765            .expect("Failed to set branch B");
766
767        persistor
768            .root_new(handle, branch_c)
769            .expect("Failed to create new root");
770
771        assert!(persistor.roots.lock().expect("Failed to lock roots").len() == 1);
772        assert!(
773            persistor
774                .branches
775                .lock()
776                .expect("Failed to lock branches")
777                .len()
778                == 3
779        );
780        assert!(
781            persistor
782                .leaves
783                .lock()
784                .expect("Failed to lock leaves")
785                .len()
786                == 3
787        );
788        assert!(
789            persistor
790                .stumps
791                .lock()
792                .expect("Failed to lock leaves")
793                .len()
794                == 1
795        );
796        assert!(
797            persistor
798                .references
799                .lock()
800                .expect("Failed to lock references")
801                .len()
802                == 7
803        );
804
805        let leaf_3 = persistor.leaf_set(vec![3]).expect("Failed to set leaf 3");
806        let branch_d = persistor
807            .branch_set(leaf_2, leaf_3, zeros)
808            .expect("Failed to set branch D");
809        persistor
810            .root_set(handle, branch_c, branch_d)
811            .expect("Failed to set root");
812
813        assert!(persistor.roots.lock().expect("Failed to lock roots").len() == 1);
814        assert!(
815            persistor
816                .branches
817                .lock()
818                .expect("Failed to lock branches")
819                .len()
820                == 1
821        );
822        assert!(
823            persistor
824                .leaves
825                .lock()
826                .expect("Failed to lock leaves")
827                .len()
828                == 2
829        );
830        assert!(
831            persistor
832                .stumps
833                .lock()
834                .expect("Failed to lock stumps")
835                .len()
836                == 0
837        );
838        assert!(
839            persistor
840                .references
841                .lock()
842                .expect("Failed to lock references")
843                .len()
844                == 3
845        );
846    }
847
848    #[test]
849    fn test_database_garbage() {
850        let db = ".test-database-garbage";
851        let _ = fs::remove_dir_all(db);
852        let persistor = DatabasePersistor::new(db);
853        let zeros: Word = [0 as u8; SIZE];
854        let handle: Word = [0 as u8; SIZE];
855        let leaf_0 = persistor.leaf_set(vec![0]).expect("Failed to set leaf 0");
856        let leaf_1 = persistor.leaf_set(vec![1]).expect("Failed to set leaf 1");
857        let leaf_2 = persistor.leaf_set(vec![2]).expect("Failed to set leaf 2");
858
859        let stump_0 = persistor
860            .stump_set([0 as u8; SIZE])
861            .expect("Failed to set stump 0");
862
863        let branch_a = persistor
864            .branch_set(leaf_0, leaf_1, zeros)
865            .expect("Failed to set branch A");
866        let branch_b = persistor
867            .branch_set(branch_a, leaf_2, zeros)
868            .expect("Failed to set branch B");
869        let branch_c = persistor
870            .branch_set(branch_b, stump_0, zeros)
871            .expect("Failed to set branch C");
872
873        persistor
874            .root_new(handle, branch_c)
875            .expect("Failed to create new root");
876
877        let cf_count = |mdb: &Mutex<DB>, cf| {
878            let db_ = mdb.lock().expect("Failed to lock database");
879            db_.iterator_cf(
880                db_.cf_handle(cf).expect("Failed to get CF handle"),
881                IteratorMode::Start,
882            )
883            .count()
884        };
885
886        {
887            assert!(cf_count(&persistor.db, "roots") == 1);
888            assert!(cf_count(&persistor.db, "branches") == 3);
889            assert!(cf_count(&persistor.db, "leaves") == 3);
890            assert!(cf_count(&persistor.db, "stumps") == 1);
891            assert!(cf_count(&persistor.db, "references") == 7);
892        }
893
894        let leaf_3 = persistor.leaf_set(vec![3]).expect("Failed to set leaf 3");
895        let branch_d = persistor
896            .branch_set(leaf_2, leaf_3, zeros)
897            .expect("Failed to set branch D");
898        persistor
899            .root_set(handle, branch_c, branch_d)
900            .expect("Failed to set root");
901
902        {
903            assert!(cf_count(&persistor.db, "roots") == 1);
904            assert!(cf_count(&persistor.db, "branches") == 1);
905            assert!(cf_count(&persistor.db, "leaves") == 2);
906            assert!(cf_count(&persistor.db, "stumps") == 0);
907            assert!(cf_count(&persistor.db, "references") == 3);
908        }
909
910        let _ = fs::remove_dir_all(db);
911    }
912}
journal_sdk/persistor.rs

journal_sdk/
persistor.rs
