Add shared KvStore with Arc type, thread safe move to heap

Author: hindenbug

kvs-client-server/Cargo.toml

@@ -9,7 +9,7 @@ edition = "2018"
 [dependencies]
 clap = "2.33.0"
 failure = "0.1.5"
-serde = { version = "1.0.89", features = ["derive"] }
+serde = { version = "1.0.89", features = ["derive", "rc"] }
 serde_json = "1.0.39"
 bincode = "1.2.1"
 structopt = "0.3"

kvs-client-server/src/bin/kvs-server.rs

@@ -79,7 +79,7 @@ fn main() -> Result<()> {
 }
 
 fn start_server_with<E: KvsEngine>(addr: &str, engine: E) -> Result<()> {
-    let mut server = Server::new(addr, engine);
+    let server = Server::new(addr, engine);
     server.serve()?;
     Ok(())
 }

kvs-client-server/src/engines/kvs.rs

@@ -5,12 +5,18 @@ use std::collections::BTreeMap;
 use std::fs::{self, File};
 use std::io::{BufReader, BufWriter, Seek, SeekFrom, Write};
 use std::path::PathBuf;
+use std::sync::{Arc, Mutex};
 
 const COMPACTION_THRESHOLD: u64 = 1024;
 const LOG_FILE_NAME: &str = "current.db";
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct KvStore {
+    inner: Arc<Mutex<InnerKvStore>>,
+}
+
+#[derive(Debug)]
+struct InnerKvStore {
     path: PathBuf,
     log: File,
     map: BTreeMap<String, LogPointer>,
@@ -36,44 +42,53 @@ struct LogPointer {
 }
 
 impl KvsEngine for KvStore {
-    fn set(&mut self, key: String, value: String) -> Result<()> {
-        let offset = self.log.seek(SeekFrom::End(0))?;
+    fn set(&self, key: String, value: String) -> Result<()> {
+        let mut inner = self.inner.lock().unwrap();
+
+        let offset = inner.log.seek(SeekFrom::End(0))?;
         let command = Command {
             cmd: CommandType::Set,
             key: key.clone(),
             value,
         };
         // encoding before writing to log
-        serde_json::to_writer(&mut self.log, &command)?;
-        self.log.flush()?;
-        let current_offset = self.log.seek(SeekFrom::End(0))?;
-        self.map.insert(
+        serde_json::to_writer(&mut inner.log, &command)?;
+        inner.log.flush()?;
+        let current_offset = inner.log.seek(SeekFrom::End(0))?;
+        inner.map.insert(
             key,
             LogPointer {
                 offset,
                 len: current_offset - offset,
             },
         );
         if current_offset > COMPACTION_THRESHOLD {
-            self.compact()?;
+            inner.compact()?;
         }
         Ok(())
     }
 
-    fn get(&mut self, key: String) -> Result<Option<String>> {
-        if let Some(pointer) = self.map.get(&key) {
-            &self.log.seek(SeekFrom::Start(pointer.offset));
-            let mut de = serde_json::Deserializer::from_reader(&self.log);
-            let cmd: Command = serde::de::Deserialize::deserialize(&mut de)?;
-            Ok(Some(cmd.value))
-        } else {
-            Ok(None)
+    fn get(&self, key: String) -> Result<Option<String>> {
+        let mut inner = self.inner.lock().unwrap();
+
+        if !inner.map.contains_key(&key) {
+            return Ok(None);
         }
+
+        let pointer = inner.map.get(&key).unwrap();
+        let pos = SeekFrom::Start(pointer.offset);
+
+        inner.log.seek(pos)?;
+        let mut de = serde_json::Deserializer::from_reader(&inner.log);
+        let cmd: Command = serde::de::Deserialize::deserialize(&mut de)?;
+        Ok(Some(cmd.value))
     }
 
-    fn remove(&mut self, key: String) -> Result<()> {
-        self.log.seek(SeekFrom::End(0))?;
-        if self.map.get(&key).is_none() {
+    fn remove(&self, key: String) -> Result<()> {
+        let mut inner = self.inner.lock().unwrap();
+
+        inner.log.seek(SeekFrom::End(0))?;
+        if inner.map.get(&key).is_none() {
             return Err(KvsError::KeyNotFound);
         }
         let command = Command {
@@ -82,30 +97,35 @@ impl KvsEngine for KvStore {
             value: String::new(),
         };
         // encoding before writing to log
-        serde_json::to_writer(&mut self.log, &command)?;
-        self.log.flush()?;
-        self.map.remove(&key);
+        serde_json::to_writer(&mut inner.log, &command)?;
+        inner.log.flush()?;
+        inner.map.remove(&key);
         Ok(())
     }
 }
 
 impl KvStore {
     pub fn open(path: impl Into<PathBuf>) -> Result<KvStore> {
         let path = path.into();
-        fs::create_dir_all(&path)?;
-        let log = Self::new_log_file(&path)?;
+        fs::create_dir_all(&*path)?;
+        let log = InnerKvStore::new_log_file(&path)?;
 
-        let mut store = KvStore {
-            path: path,
-            log: log,
+        let mut inner = InnerKvStore {
+            path,
+            log,
             map: BTreeMap::new(),
         };
 
         // Load from log files
-        store.load_from_log()?;
-        Ok(store)
+        inner.load_from_log()?;
+
+        Ok(KvStore {
+            inner: Arc::new(Mutex::new(inner)),
+        })
     }
+}
 
+impl InnerKvStore {
     fn load_from_log(&mut self) -> Result<()> {
         let mut reader = BufReader::new(File::open(&self.path.join(LOG_FILE_NAME))?);
         let mut offset = reader.seek(SeekFrom::Start(0))?;
@@ -141,7 +161,7 @@ impl KvStore {
         Ok(())
     }
 
-    fn compact(&mut self) -> Result<()> {
+    fn compact(&self) -> Result<()> {
         let tmp_path: PathBuf = self.path.join("tmp.db");
         let file_path: PathBuf = self.path.join(LOG_FILE_NAME);
         let mut work_file = self.log.try_clone()?;

kvs-client-server/src/engines/mod.rs

@@ -1,23 +1,23 @@
 use crate::Result;
 
 /// Trait for a key value storage engine.
-pub trait KvsEngine {
+pub trait KvsEngine: Clone + Send + 'static {
     /// Sets the value of a string key to a string.
     ///
     /// If the key already exists, the previous value will be overwritten.
-    fn set(&mut self, key: String, value: String) -> Result<()>;
+    fn set(&self, key: String, value: String) -> Result<()>;
 
     /// Gets the string value of a given string key.
     ///
     /// Returns `None` if the given key does not exist.
-    fn get(&mut self, key: String) -> Result<Option<String>>;
+    fn get(&self, key: String) -> Result<Option<String>>;
 
     /// Removes a given key.
     ///
     /// # Errors
     ///
     /// It returns `KvsError::KeyNotFound` if the given key is not found.
-    fn remove(&mut self, key: String) -> Result<()>;
+    fn remove(&self, key: String) -> Result<()>;
 }
 
 mod kvs;

kvs-client-server/src/engines/sled.rs

@@ -1,6 +1,8 @@
 use crate::{KvsEngine, KvsError, Result};
 use sled::Db;
 use std::path::PathBuf;
+
+#[derive(Clone)]
 pub struct SledKvsEngine {
     tree: Db,
 }
@@ -13,13 +15,13 @@ impl SledKvsEngine {
 }
 
 impl KvsEngine for SledKvsEngine {
-    fn set(&mut self, key: String, value: String) -> Result<()> {
+    fn set(&self, key: String, value: String) -> Result<()> {
         self.tree.insert(key, value.as_bytes())?;
         self.tree.flush()?;
         Ok(())
     }
 
-    fn get(&mut self, key: String) -> Result<Option<String>> {
+    fn get(&self, key: String) -> Result<Option<String>> {
         Ok(self
             .tree
             .get(key)?
@@ -28,7 +30,7 @@ impl KvsEngine for SledKvsEngine {
             .transpose()?)
     }
 
-    fn remove(&mut self, key: String) -> Result<()> {
+    fn remove(&self, key: String) -> Result<()> {
         self.tree.remove(key)?.ok_or(KvsError::KeyNotFound)?;
         self.tree.flush()?;
         Ok(())

kvs-client-server/src/server.rs

@@ -22,7 +22,7 @@ impl<E: KvsEngine> Server<E> {
         Server { listener, engine }
     }
 
-    pub fn serve(&mut self) -> Result<()> {
+    pub fn serve(&self) -> Result<()> {
         debug!("Waiting for connections...");
         let listnr = self.listener.try_clone().unwrap();
         for stream in listnr.incoming() {
@@ -38,7 +38,7 @@ impl<E: KvsEngine> Server<E> {
         Ok(())
     }
 
-    fn handle_client(&mut self, stream: TcpStream) -> Result<()> {
+    fn handle_client(&self, stream: TcpStream) -> Result<()> {
         debug!(
             "Connection established from {}, waiting for data...",
             stream.peer_addr()?