Add file-based persistence to the KV store

Author: dinowernli

Cargo.lock

@@ -27,6 +27,7 @@ tracing-subscriber = {version = "0.3.1", features = ["env-filter"]}
 tower = "0.4.13"
 pin-project = "1.1"
 querystring = "1.1"
+async-tempfile = "0.7.0"
 
 [build-dependencies]
 tonic-build = "0.11"

Cargo.toml

@@ -15,5 +15,6 @@ pub mod grpc {
 pub(in crate::keyvalue) mod http;
 #[path = "generated/keyvalue_proto.rs"]
 pub(in crate::keyvalue) mod keyvalue_proto;
+pub(in crate::keyvalue) mod persistence;
 pub(in crate::keyvalue) mod service;
 pub(in crate::keyvalue) mod store;

src/keyvalue/mod.rs

@@ -0,0 +1,113 @@
+use crate::keyvalue::keyvalue_proto::Entry;
+use async_std::fs;
+use async_std::fs::{File, OpenOptions};
+use async_std::io::{BufReader, ReadExt};
+use async_std::path::Path;
+use async_trait::async_trait;
+use bytes::Bytes;
+use futures::AsyncWriteExt;
+use prost::Message;
+use tracing::info;
+
+#[async_trait]
+pub trait Persistence {
+    async fn add(&mut self, entry: &Entry);
+}
+
+pub struct FilePersistence {
+    directory: String,
+
+    log_path: String,
+    log: File,
+}
+
+impl FilePersistence {
+    pub async fn create(directory: String) -> FilePersistence {
+        let meta = fs::metadata(&directory).await.expect("exists");
+        assert!(meta.is_dir());
+        let wal_path = Path::new(directory.as_str()).join("wal");
+        info!("WAL path: {}", &wal_path.display());
+        let wal_file = OpenOptions::new()
+            .write(true)
+            .append(true)
+            .create_new(true)
+            .open(wal_path.clone())
+            .await
+            .expect("open wal file");
+        FilePersistence {
+            directory,
+            log: wal_file,
+            log_path: wal_path.to_str().unwrap().to_string(),
+        }
+    }
+
+    pub async fn read_wal(&self) -> u64 {
+        let file = OpenOptions::new()
+            .read(true)
+            .open(self.log_path.clone())
+            .await
+            .expect("open wal file");
+
+        let mut reader = BufReader::new(file);
+        let mut buf = Vec::new();
+        reader.read_to_end(&mut buf).await.expect("read");
+
+        let mut entries = Vec::new();
+        let mut slice = &buf[..];
+
+        while !slice.is_empty() {
+            match Entry::decode_length_delimited(&mut slice) {
+                Ok(entry) => entries.push(entry),
+                Err(err) => {
+                    eprintln!("Failed to decode Entry: {}", err);
+                    break;
+                }
+            }
+        }
+        entries.len() as u64
+    }
+}
+
+#[async_trait]
+impl Persistence for FilePersistence {
+    async fn add(&mut self, entry: &Entry) {
+        let mut buf = vec![];
+        entry.encode_length_delimited(&mut buf).expect("encode");
+        self.log.write(&buf).await.expect("write");
+        self.log.flush().await.expect("flush");
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use async_tempfile::TempDir;
+    #[tokio::test]
+    async fn test_something() {
+        let tmp_dir = TempDir::new().await.unwrap();
+        let tmp_dir_path = tmp_dir.dir_path().to_str().unwrap().to_string();
+        let mut persistence = FilePersistence::create(tmp_dir_path).await;
+
+        let k1 = Bytes::from("key1");
+        let v1 = Bytes::from("value1");
+        let k2 = Bytes::from("key2");
+        let v2 = Bytes::from("value2");
+
+        persistence
+            .add(&Entry {
+                key: k1.to_vec(),
+                value: v1.to_vec(),
+            })
+            .await;
+        assert_eq!(persistence.read_wal().await, 1);
+
+        persistence
+            .add(&Entry {
+                key: k2.to_vec(),
+                value: v2.to_vec(),
+            })
+            .await;
+        assert_eq!(persistence.read_wal().await, 2);
+    }
+}

src/keyvalue/persistence.rs

@@ -160,6 +160,7 @@ mod tests {
                 .lock()
                 .await
                 .apply(&Bytes::from(copy))
+                .await
                 .expect("bad payload");
             Ok(EntryId { term: 0, index: 0 })
         }
@@ -175,14 +176,15 @@ mod tests {
 
     #[tokio::test]
     async fn test_get() {
-        let service = create_service();
+        let service = create_service().await;
         let store = service.store.clone();
         let server = TestRpcServer::run(KeyValueServer::new(service)).await;
 
         store
             .lock()
             .await
-            .set(Bytes::from("foo"), Bytes::from("bar"));
+            .set(Bytes::from("foo"), Bytes::from("bar"))
+            .await;
 
         let request = GetRequest {
             key: "foo".as_bytes().to_vec(),
@@ -203,7 +205,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_put() {
-        let service = create_service();
+        let service = create_service().await;
         let store = service.store.clone();
         let server = TestRpcServer::run(KeyValueServer::new(service)).await;
 
@@ -227,8 +229,8 @@ mod tests {
     }
 
     // Returns an instance of the service struct we're testing
-    fn create_service() -> KeyValueService {
-        let store = Arc::new(Mutex::new(MapStore::new()));
+    async fn create_service() -> KeyValueService {
+        let store = Arc::new(Mutex::new(MapStore::create_in_memory().await));
         KeyValueService {
             name: "test".into(),
             store: store.clone(),