Design an In-Memory Key-Value Store

Description

An in-memory key-value store - a small Redis - starts as a hash map and becomes interesting when the interviewer adds keys that expire, atomic counters and transactions that can roll back. It tests time handling (absolute expiry, lazy deletion plus sweeping), exact rollback with an undo log, and atomicity under concurrency. This skill builds each of those on top of one map.

An in-memory key-value store is a HashMap - until the interviewer adds expiring keys, counters and transactions:

Map<String, String> data = new HashMap<>();
Map<String, Long> secondsLeft = new HashMap<>();            // a timer thread decrements these every second

void begin()    { snapshot = new HashMap<>(data); }         // copy everything
void rollback() { data = snapshot; }

long incr(String key) {
    long n = Long.parseLong(data.getOrDefault(key, "0"));
    data.put(key, Long.toString(n + 1));
    return n + 1;
}

data: dict[str, str] = {}
seconds_left: dict[str, int] = {}               # a timer thread decrements these every second

def begin() -> None:
    global snapshot
    snapshot = dict(data)                       # copy everything

def rollback() -> None:
    global data
    data = snapshot

def incr(key: str) -> int:
    n = int(data.get(key, "0"))
    data[key] = str(n + 1)
    return n + 1

The timer thread races with every reader, drifts when it's late, and does work proportional to all keys every second. begin copies the whole store - a million keys to change one. The snapshot forgets the TTL map, so a rollback leaves values and their expiries out of step, and nested transactions can't work at all. And incr loses updates as soon as two threads call it: both read 41, both write 42.

We'll go from simple to more technical:

1. The model: one map of entries, each with an optional absolute expiry time.
2. Get, set, increment and delete - with expiry checked on access.
3. Transactions with an undo log, including nested ones.
4. Expiry and memory, concurrency, and the bugs in between.

The approach

1. Clarify. Which operations - get, set (with or without a TTL), delete, expire, ttl, increment? Values are strings? What does set do to an existing TTL? Transactions: nested? Per client or global? Persistence? Concurrency?
2. Model. An Entry(value, expiresAt) record - expiresAt an absolute Instant, or null for never. One map from key to entry. A stack of undo maps, one per open transaction. An injected Clock.
3. API. Redis-like: set(key, value), set(key, value, ttl), get, delete, expire, ttlMillis (-1 for no expiry, -2 for missing), size, increment, begin, commit, rollback.
4. Core. Lazy expiry: a key is live while now is before its expiresAt, checked on every read. One write path: every change goes through a single write method, which is where an open transaction records the key's previous state.

The pieces in code

private record Entry(String value, Instant expiresAt) {}        // expiresAt == null: never
private final Map<String, Entry> data = new HashMap<>();
private final Deque<Map<String, Optional<Entry>>> undo = new ArrayDeque<>();

private Entry live(String key) {                                // lazy expiry
    Entry e = data.get(key);
    boolean expired = e != null && e.expiresAt() != null && !clock.instant().isBefore(e.expiresAt());
    return e == null || expired ? null : e;
}

private void write(String key, Entry entry) {                   // the only place data changes
    if (!undo.isEmpty()) undo.peek().putIfAbsent(key, Optional.ofNullable(data.get(key)));
    if (entry == null) data.remove(key); else data.put(key, entry);
}

public synchronized void rollback() {
    if (undo.isEmpty()) throw new IllegalStateException("No transaction is open");
    undo.pop().forEach((key, before) -> {
        if (before.isPresent()) data.put(key, before.get());    // value *and* expiry come back
        else data.remove(key);                                  // it didn't exist
    });
}

public synchronized void commit() {
    if (undo.isEmpty()) throw new IllegalStateException("No transaction is open");
    Map<String, Optional<Entry>> inner = undo.pop();
    if (!undo.isEmpty()) inner.forEach(undo.peek()::putIfAbsent);   // the outer keeps its own "before"
}

import threading
from collections.abc import Callable
from dataclasses import dataclass
from datetime import datetime

_ABSENT = object()                              # "the key didn't exist" - never mistaken for a value


@dataclass(frozen=True)
class Entry:
    value: str
    expires_at: datetime | None                 # None: never


class Store:
    def __init__(self, now: Callable[[], datetime]):
        self._now = now
        self._data: dict[str, Entry] = {}
        self._undo: list[dict[str, object]] = []
        self._lock = threading.RLock()

    def _live(self, key: str) -> Entry | None:                  # lazy expiry
        e = self._data.get(key)
        expired = e is not None and e.expires_at is not None and self._now() >= e.expires_at
        return None if e is None or expired else e

    def _write(self, key: str, entry: Entry | None) -> None:    # the only place data changes
        if self._undo:
            self._undo[-1].setdefault(key, self._data.get(key, _ABSENT))
        if entry is None:
            self._data.pop(key, None)
        else:
            self._data[key] = entry

    def rollback(self) -> None:
        with self._lock:
            if not self._undo:
                raise RuntimeError("No transaction is open")
            for key, before in self._undo.pop().items():
                if before is _ABSENT:
                    self._data.pop(key, None)                   # it didn't exist
                else:
                    self._data[key] = before                    # value *and* expiry come back

    def commit(self) -> None:
        with self._lock:
            if not self._undo:
                raise RuntimeError("No transaction is open")
            inner = self._undo.pop()
            if self._undo:
                for key, before in inner.items():
                    self._undo[-1].setdefault(key, before)      # the outer keeps its own "before"

• An absolute expiresAt makes expiry a comparison, not a job: nothing ticks, and a fake clock in a test can jump an hour.
• putIfAbsent in write records only a key's state before its first change in a transaction; recording every change would make rollback restore an intermediate value.
• Optional as the "before" value is deliberate. With null meaning "the key didn't exist", putIfAbsent - which treats a key mapped to null as absent - would let an inner commit overwrite the outer transaction's record, and rolling back the outer transaction would resurrect a value that never existed before it. (Python's setdefault checks the key, not the value, so None would survive the merge - but a sentinel such as _ABSENT = object() still can't be mistaken for a stored value, and if outer.get(key) is None would bring the trap right back.)
• Redis semantics are worth matching because interviewers know them: a plain set clears any TTL, increment keeps it, and ttl reports -1 (no expiry) or -2 (missing).

Expiry and memory

Lazy expiry alone is always correct but leaks memory: a key that's never read again is never removed. Redis therefore combines it with active expiry - many times a second it samples a few keys that have a TTL, deletes the expired ones, and keeps sampling while a large share turn out to be expired. An interview-sized version can run a similar sweep, or clean up during writes; either way, correctness must never depend on the sweep having run.

Concurrency

increment is a read-parse-write: without atomicity two threads read the same number and one update is lost. A ConcurrentHashMap makes each get and put safe, but not the sequence; compute(key, ...) runs the whole update atomically for that key. The simplest correct store makes every method synchronized - every operation is a few map calls, so the lock is held briefly. (In Python: one threading.Lock - an RLock if the methods call each other - taken with with in every public method.)

Transactions here belong to the store, which suits one client. With many clients, one client's rollback must not undo another's writes: real systems keep a transaction per connection. Redis MULTI/EXEC queues a client's commands and runs them together, with no rollback at all; databases isolate transactions with locks or multi-version concurrency control.

Which tool for which job?

Common mistakes

Check yourself

Try to answer before opening each one - recalling beats re-reading.

1. A key was set at 12:00:00 with a 60-second TTL. At 12:00:30, set(key, "new") is called without a TTL. Under Redis semantics, when does the key expire - and what if it had been increment instead?

2. Inside a transaction, a key goes from "a" to "b" to "c". What must the undo log hold for it, and what does rollback restore?

3. The outer transaction creates key y; an inner transaction changes it and commits; then the outer rolls back. What should y be, and how can null get this wrong?

Key takeaways

• Store an absolute expiry per value and check it on read; sweep separately only to reclaim memory.
• Match known semantics: set clears a TTL, increment keeps it; ttl reports -1 and -2.
• Funnel every change through one write method, where transactions record each key's state before their first change.
• Rollback restores exactly those states (TTLs included); commit folds them into the enclosing transaction.
• Mark "absent" explicitly - null plus putIfAbsent silently loses the record.
• Make read-modify-write operations atomic, and keep transactions per client once there are several.