Add sst write/read to main lib functions; we have 🚀. Also some metrics.

src/lib.rs

@@ -1,5 +1,9 @@
-use std::{collections::BTreeMap, path::Path};
+use std::{
+    collections::BTreeMap,
+    path::{Path, PathBuf},
+};
 
+use sstable::{new_reader, new_writer};
 use wal::WAL;
 
 mod kvrecord;
@@ -18,32 +22,68 @@ impl From<std::io::Error> for ToyKVError {
     }
 }
 
+#[derive(Default)]
+pub struct ToyKVMetrics {
+    sst_flushes: u64,
+    reads: u64,
+    writes: u64,
+}
+
 pub struct ToyKV<'a> {
     /// d is the folder that the KV store owns.
+    d: PathBuf,
     memtable: BTreeMap<Vec<u8>, Vec<u8>>,
     wal: WAL<'a>,
+    pub metrics: ToyKVMetrics,
 }
 
 pub fn open(d: &Path) -> Result<ToyKV, ToyKVError> {
     // TODO check path, fail if it doesn't exist.
     // TODO initialise if empty (new WAL?)
     let mut wal = wal::new(d);
     let memtable = wal.replay()?;
-    Ok(ToyKV { memtable, wal })
+    Ok(ToyKV {
+        d: d.to_path_buf(),
+        memtable,
+        wal,
+        metrics: Default::default(),
+    })
 }
 
+/// How many writes to a WAL before we SSTable it.
+const WAL_WRITE_THRESHOLD: u32 = 1000;
+
 impl<'a> ToyKV<'a> {
     /// Sets key k to v.
     pub fn set(&mut self, k: Vec<u8>, v: Vec<u8>) -> Result<(), ToyKVError> {
         self.wal.write(&k, &v)?;
         self.memtable.insert(k, v);
+        if self.wal.wal_writes > WAL_WRITE_THRESHOLD {
+            let mut sst = new_writer(self.d.as_path())?;
+            for entry in &self.memtable {
+                sst.write(entry.0, entry.1)?;
+            }
+            sst.finalise()?;
+            self.wal.reset()?;
+            self.memtable.clear();
+            self.metrics.sst_flushes += 1;
+        }
+        self.metrics.writes += 1;
         Ok(())
     }
 
     /// Get the value for k.
-    pub fn get(&self, k: &[u8]) -> Result<Option<Vec<u8>>, ToyKVError> {
-        let r: Vec<u8> = self.memtable[k].clone();
-        Ok(Some(r))
+    pub fn get(&mut self, k: &[u8]) -> Result<Option<Vec<u8>>, ToyKVError> {
+        let r = self.memtable.get(k);
+        let r = match r {
+            Some(r) => Some(r.clone()),
+            None => {
+                let sst = new_reader(self.d.as_path())?;
+                sst.get(k)?
+            }
+        };
+        self.metrics.reads += 1;
+        Ok(r)
     }
 
     /// Perform a graceful shutdown.

src/sstable.rs

@@ -133,6 +133,7 @@ mod tests_sorted_sstable_files {
 
 /// Provides methods to write already-sorted KVRecords to an on-disk file.
 pub(crate) struct SSTableFileWriter {
+    #[allow(dead_code)] // used in tests to check filepath
     p: PathBuf,
     f: File,
 }

tests/endtoend.rs

@@ -21,6 +21,23 @@ fn insert_and_readback() -> Result<(), ToyKVError> {
 
     Ok(())
 }
+#[test]
+fn grace_on_missing_key() -> Result<(), ToyKVError> {
+    let k = "foo".to_string();
+    let v = "the rain in spain falls mainly on the plain".to_string();
+
+    let tmp_dir = tempfile::tempdir().unwrap();
+
+    let mut db = toykv::open(tmp_dir.path())?;
+    match db.set(k.clone().into_bytes(), v.clone().into_bytes()) {
+        Ok(it) => it,
+        Err(err) => return Err(err),
+    };
+    let got = db.get("bar".as_bytes())?;
+    assert_eq!(got, None, "Didn't get None for missing key");
+
+    Ok(())
+}
 
 #[test]
 fn data_survive_restart() -> Result<(), ToyKVError> {
@@ -36,7 +53,7 @@ fn data_survive_restart() -> Result<(), ToyKVError> {
     };
     db.shutdown();
 
-    let db2 = toykv::open(tmp_dir.path())?;
+    let mut db2 = toykv::open(tmp_dir.path())?;
     let got = db2.get(k.as_bytes())?;
     assert_eq!(
         got.unwrap(),