lol_core/storage/
memory.rs

1use super::{Ballot, Entry};
2use crate::Index;
3use anyhow::Result;
4use std::collections::BTreeMap;
5use std::sync::Arc;
6use tokio::sync::{Mutex, RwLock};
7
8pub struct Storage {
9    entries: Arc<RwLock<BTreeMap<u64, super::Entry>>>,
10    ballot: Arc<Mutex<Ballot>>,
11}
12impl Storage {
13    pub fn new() -> Self {
14        Self {
15            entries: Arc::new(RwLock::new(BTreeMap::new())),
16            ballot: Arc::new(Mutex::new(Ballot::new())),
17        }
18    }
19}
20#[async_trait::async_trait]
21impl super::RaftStorage for Storage {
22    async fn get_head_index(&self) -> Result<Index> {
23        let x = self.entries.read().await;
24        let r = match x.iter().next() {
25            Some((k, _)) => *k,
26            None => 0,
27        };
28        Ok(r)
29    }
30    async fn get_last_index(&self) -> Result<Index> {
31        let x = self.entries.read().await;
32        let r = match x.iter().next_back() {
33            Some((k, _)) => *k,
34            None => 0,
35        };
36        Ok(r)
37    }
38    async fn delete_entry(&self, i: Index) -> Result<()> {
39        self.entries.write().await.remove(&i);
40        Ok(())
41    }
42    async fn insert_entry(&self, i: Index, e: Entry) -> Result<()> {
43        self.entries.write().await.insert(i, e);
44        Ok(())
45    }
46    async fn get_entry(&self, i: Index) -> Result<Option<Entry>> {
47        let r = self.entries.read().await.get(&i).cloned();
48        Ok(r)
49    }
50    async fn save_ballot(&self, v: Ballot) -> Result<()> {
51        *self.ballot.lock().await = v;
52        Ok(())
53    }
54    async fn load_ballot(&self) -> Result<Ballot> {
55        let r = self.ballot.lock().await.clone();
56        Ok(r)
57    }
58}
59
60#[cfg(test)]
61mod tests {
62    use super::*;
63    use crate::storage;
64
65    #[tokio::test]
66    async fn test_mem_storage() -> Result<()> {
67        let s = Storage::new();
68        storage::test_storage(s).await?;
69        Ok(())
70    }
71}