Building an In-Memory K-V Store in Rust

Over the weekend I decided to learn a little more Rust with PingCAP's PNA course. The first lesson involves writing a tiny key-value store that lives in memory - which means no need to bother with complicated I/O calls (yet!). The final application is supposed to support these operations:

• kvs set <KEY> <VALUE>: Set the value of a string key to a string
• kvs get <KEY>: Get the string value of a given string key
• kvs rm <KEY>: Remove a given key.

Looking at the specification and the accompanying tests, I figured it makes sense to simply write a wrapper around a HashMap.

Wrapping Rust HashMaps

Rust supports regular Computer Science hashmaps through the std::collections::HashMap struct. Also, it uses quadratic probing to resolve hash collisions. I sort of learned about that the week prior so that bit was interesting to me. While there are a bunch of cool things you could do with HashMaps in Rust, I ended up needing just three methods:

• insert to add a new item
• get to retrieve an item (if it exists), and
• remove to delete an item from the map.

Here's what my final library code (in src/lib.rs) looked like:

use std::collections::HashMap;

pub struct KvStore {
    items: HashMap<String, String>,
}

impl Default for KvStore {
    fn default() -> Self {
        Self::new()
    }
}

impl KvStore {
    pub fn new() -> Self {
        KvStore {
            items: HashMap::new(),
        }
    }

    pub fn set(&mut self, key: String, value: String) {
        self.items.insert(key, value);
    }

    pub fn get(&self, key: String) -> Option<String> {
        let val = self.items.get(key.as_str());
        match val {
            Some(v) => Some(v.to_string()),
            _ => None,
        }
    }

    pub fn remove(&mut self, key: String) {
        self.items.remove(key.as_str());
    }
}

Clippy was complaining about a missing Default implementation. Turns out it makes sense for types that can be initialized through a new method to implement the Default trait. That way, callers can write statements like:

let kvs: KvStore = Default::default();

and Rust will initialize it with default values that make sense (to you).

Building the App binary

In Rust lingo, the code above is only a library crate i.e., it is only a dependency that can be used by other projects. To run it, I'd need to use the library crate from a binary crate.

Binary crates are code that can be compiled into executable programs, and they can contain any number of library crates. They must also have a main function that serves as the entry point of a Rust binary.

Since the final project is an executable CLI app, the lesson recommends using either the clap or structopts crates to parse command-line arguments. I tried using both, and clap was such a joy to work with. StructOpts uses macros to generate the equivalent clap code for you, so I ended up using that. Here's my final src/bin/kvs.rs code:

use kvs::KvStore;
use structopt::StructOpt;

#[derive(StructOpt, Debug)]
#[structopt(name="KayVS")]
enum KVS {
    Get {
        key: String,
    },
    Set {
        key: String,
        value: String,
    },
    RM {
        key: String
    }
}

fn main() {
    let mut store = KvStore::new();
    match KVS::from_args() {
        KVS::Get {key} => {
            let val = store.get(key).expect("no matching value found");
            println!("{}", val);
            //panic!("unimplemented");
        },
        KVS::Set {key, value} => {
            store.set(key, value);
            //panic!("unimplemented");
        },
        KVS::RM {key} => {
            store.remove(key);
            //panic!("unimplemented");
        }
    }
}

The library crates’ tests (the KvStore struct) were all passing but trying to get a value from the binary by running the below command fails.

$ cargo run -- get key1

That was because the store only lived in memory as long as the app is running. It's okay though because the next lesson involves persistence and writing the data to disk, and I kind of look forward to seeing how that goes.