Time Traveling Key-Value Store, This Time in Rust

2022-04-21

Previously…

We’ve talked previously about implementing a time-traveling key-value store. Having worked more with Rust in the interim, I tried my hand at a Rust implementation. Rust’s type system, standard library, and attention to pedantic details make this my favorite version yet.

Implementation

The standard library provides everything we need:

use std::collections::BTreeMap;
use std::time::Instant;

pub struct Ttkv<K, V> {
    started: Instant,
    mapping: BTreeMap<u128, (K, V)>,
}

Recording the std::time::Instant when we created our store allows us to monotonically record insertion timestamps—on most hardware, at least—and we can still check if this assumption is violated.

We build a new store via implementing Default:

impl<K, V> Default for Ttkv<K, V> {
    fn default() -> Self {
        Self {
            started: Instant::now(),
            mapping: BTreeMap::default(),
        }
    }
}

Digging into the implementation (the following snippets are all wrapped in an impl block), it’s straightforward to check for emptiness:

pub fn is_empty(&self) -> bool {
    self.mapping.is_empty()
}

Adding a pair to our store gives us an opportunity to assert that time is monotonic:

pub fn put(&mut self, key: K, value: V, timestamp: Option<u128>) {
    let t = timestamp.unwrap_or_else(|| {
        Instant::now()
            .checked_duration_since(self.started)
            .unwrap_or_else(|| panic!("non-monotonic insertion"))
            .as_nanos()
    });
    self.mapping.insert(t, (key, value));
}

Retrieval from the store is similar to previous implementations:

pub fn get(&self, key: &K, timestamp: Option<u128>) -> Option<&V>
where
    K: PartialEq,
{
    self.mapping
        .range(0..timestamp.unwrap_or(u128::MAX))
        .filter(|(_, (k, _))| k == key)
        .last()
        .map(|(_, (_, v))| v)
}

Finally, collecting the insertion times in order:

pub fn times(&self) -> Vec<u128> {
    self.mapping.keys().cloned().collect()
}

Testing

Rust’s (and cargo’s) approach to testing is a pleasure to use. We can do regular unit-style tests:

#[test]
fn initially_empty() {
    let t = Ttkv::<String, String>::default();
    assert!(t.is_empty());
    assert!(t.times().is_empty());
}

And with the help of proptest we can do property-based testing:

#[test]
fn two_gets_different_keys(a: String, b: String, x: String, y: String) {
    prop_assume!(a != b);
    let mut t = Ttkv::default();
    t.put(a.clone(), x.clone(), None);
    t.put(b.clone(), y.clone(), None);
    prop_assert_eq!(t.times().len(), 2);
    prop_assert_eq!(t.get(&a, None), Some(&x));
    prop_assert_eq!(t.get(&b, None), Some(&y));
}

As a side note, this test originally failed for me in an early draft where the timestamping mechanism wasn’t monotonic.

Fin.

This was a fun exercise, and again is probably my favorite implementation. See the repo for the full code!