Problem 126 · Practical builds

In-Memory Key-Value Store: TTL, Transactions, Snapshots

MediumPractical multi-stage build

01Problem

Build an in-memory key-value store in four stages, each extending the last. The final version supports get/set/delete, per-key TTL against an injected clock, nested BEGIN/ROLLBACK/COMMIT, and snapshot/restore.

Stages — finish each before reading the next

1Core store

Implement set(key, value), get(key, default) and delete(key) -> bool (True if the key existed).

2TTL

set accepts an optional ttl in seconds. Expired keys read as missing; ttl(key) returns remaining seconds, None for no expiry, and raises KeyError for a missing key. Time comes from an injected clock so tests are deterministic.

3Nested transactions

Add begin(), rollback() and commit(). Transactions nest: an inner commit merges into its parent only; rollback discards the innermost layer. Deletes inside a transaction must shadow committed values. Commit/rollback with no open transaction raises.

4Snapshot / restore

snapshot() returns a deep copy of live committed state; restore(snap) replaces the store, dropping entries that expired since. Both refuse to run while a transaction is open.

02Key insight

Model each transaction as a write-overlay dict on a stack, with a sentinel for deletes. Reads walk the stack top-down then fall through to committed data, commit is a merge into the parent layer, rollback is a pop. Store absolute expiry times, not TTLs, so overlays and snapshots stay correct as the clock moves.

03Solution

solution.pypython88 lines

import copy
from typing import Any, Callable, Optional

_DELETED = object()  # marker a transaction layer uses to shadow a key as deleted
Entry = tuple[Any, Optional[float]]  # (value, absolute expiry or None)


class KVStore:
    """get/set/delete + TTL + nested transactions + snapshot/restore."""

    def __init__(self, clock: Callable[[], float]) -> None:
        self._clock = clock
        self._data: dict[str, Entry] = {}
        self._tx: list[dict[str, Any]] = []  # innermost layer last

    # -- internals -------------------------------------------------------
    def _lookup(self, key: str) -> Optional[Entry]:
        for layer in reversed(self._tx):
            if key in layer:
                entry = layer[key]
                return None if entry is _DELETED else entry
        return self._data.get(key)

    def _alive(self, entry: Optional[Entry]) -> bool:
        return entry is not None and (entry[1] is None or self._clock() < entry[1])

    def _write(self, key: str, entry: Any) -> None:
        if self._tx:
            self._tx[-1][key] = entry
        elif entry is _DELETED:
            self._data.pop(key, None)
        else:
            self._data[key] = entry

    # -- stage 1 + 2: CRUD with TTL ----------------------------------------
    def set(self, key: str, value: Any, ttl: Optional[float] = None) -> None:
        if ttl is not None and ttl <= 0:
            raise ValueError("ttl must be positive")
        expires = None if ttl is None else self._clock() + ttl
        self._write(key, (value, expires))

    def get(self, key: str, default: Any = None) -> Any:
        entry = self._lookup(key)
        if not self._alive(entry):
            if entry is not None and not self._tx:
                self._data.pop(key, None)  # lazy expiry outside transactions
            return default
        return entry[0]

    def delete(self, key: str) -> bool:
        existed = self._alive(self._lookup(key))
        self._write(key, _DELETED)
        return existed

    def ttl(self, key: str) -> Optional[float]:
        entry = self._lookup(key)
        if not self._alive(entry):
            raise KeyError(key)
        return None if entry[1] is None else entry[1] - self._clock()

    # -- stage 3: nested transactions ------------------------------------
    def begin(self) -> None:
        self._tx.append({})

    def rollback(self) -> None:
        if not self._tx:
            raise RuntimeError("no open transaction")
        self._tx.pop()

    def commit(self) -> None:
        """Merge the innermost layer into its parent (or into the store)."""
        if not self._tx:
            raise RuntimeError("no open transaction")
        layer = self._tx.pop()
        for key, entry in layer.items():
            self._write(key, entry)

    # -- stage 4: snapshot / restore -------------------------------------
    def snapshot(self) -> dict[str, Entry]:
        if self._tx:
            raise RuntimeError("snapshot requires no open transaction")
        return {k: copy.deepcopy(e) for k, e in self._data.items() if self._alive(e)}

    def restore(self, snap: dict[str, Entry]) -> None:
        if self._tx:
            raise RuntimeError("restore requires no open transaction")
        now = self._clock()
        self._data = {k: copy.deepcopy(e) for k, e in snap.items() if e[1] is None or now < e[1]}

Try it first.

Timeget/set/delete O(d) for transaction depth d; commit O(keys in layer); snapshot O(n)

SpaceO(n + writes in open layers)

tests.py · 13 passing

tests.pypython51 lines

t = [0.0]
kv = KVStore(clock=lambda: t[0])
kv.set("a", 1)
assert kv.get("a") == 1 and kv.get("missing") is None
assert kv.delete("a") is True and kv.delete("a") is False
# TTL
kv.set("s", "tok", ttl=10)
t[0] = 9.0
assert kv.get("s") == "tok" and kv.ttl("s") == 1.0
t[0] = 10.0
assert kv.get("s") is None
try:
    kv.ttl("s"); assert False
except KeyError:
    pass
# nested transactions
kv.set("x", 1)
kv.begin(); kv.set("x", 2)
kv.begin(); kv.delete("x")
assert kv.get("x") is None
kv.rollback()
assert kv.get("x") == 2
kv.begin(); kv.set("y", 9); kv.commit()   # inner commit merges into outer only
assert kv.get("y") == 9
kv.rollback()                              # outer rollback discards both
assert kv.get("x") == 1 and kv.get("y") is None
kv.begin(); kv.set("x", 5); kv.commit()
assert kv.get("x") == 5
try:
    kv.commit(); assert False
except RuntimeError:
    pass
# TTL inside a transaction is evaluated at read time
kv.begin(); kv.set("e", 1, ttl=1); t[0] = 11.5
assert kv.get("e") is None
kv.commit()
# snapshot / restore, deep-copied and expiry-aware
kv.set("list", [1, 2]); kv.set("short", 1, ttl=2)
snap = kv.snapshot()
kv.get("list").append(3); kv.delete("x")
assert kv.get("list") == [1, 2, 3]
kv.restore(snap)
assert kv.get("list") == [1, 2] and kv.get("x") == 5 and kv.get("short") == 1
t[0] = 20.0
kv.restore(snap)
assert kv.get("short") is None and kv.get("x") == 5
kv.begin()
try:
    kv.snapshot(); assert False
except RuntimeError:
    pass

04Staff-level follow-up

How would you make this durable and concurrent? Append-only log of committed writes plus periodic snapshots (Redis AOF + RDB); for concurrency use MVCC (each transaction reads a version, conflicts detected at commit) rather than one global lock; expire actively with a sampled sweep, not only lazily on read.