Key-value store with TTL and transactions

Design an in-memory key-value store with expiry (TTL), an optional size cap that removes the least recently used key, and nested transactions with commit and rollback. A snapshot (the Memento pattern) is the simplest model answer for rollback.

Design an in-memory key-value store, like a small Redis. It keeps string values by key. A key may expire after a time to live (TTL). The store may have a size cap and then removes the least recently used key. Callers can open transactions, nested if they like, and commit or roll back.

Questions to ask first

  • Which operations do we need? Assume put (with or without a TTL), get, delete and a call that lists the live keys. Values are strings.
  • How does expiry work? Assume the TTL is given in minutes with the put. An expired key behaves as if it does not exist. Time comes from a clock that tests control.
  • Do we remove expired keys with a background thread? Assume no. We check on read (lazy expiry) and when the store needs room. The deep dive compares this with active cleanup.
  • What does the size cap count? Assume live keys only. When a new key does not fit, expired keys go first, then the least recently used key.
  • What does a transaction do? Assume it is a way to undo. Reads inside it see its changes, rollback restores the state at begin, and commit keeps the changes. There is no isolation between threads.
  • Can transactions nest? Yes. Committing an inner one keeps its changes, but rolling back an outer one still undoes them.
  • Is it used by many threads? Assume yes, so one lock guards every call. Remember that get changes the order of use.

Requirements

  • Put a value with no expiry or with a TTL, get it, delete it, and list the live keys.
  • A key is gone from the moment its time is up. It is never returned after that.
  • An optional cap on live keys. When it is full, drop expired keys first, then the least recently used key.
  • Begin, commit and rollback, nested to any depth.
  • Rollback restores values, expiry times and the order of use as they were at begin, including keys that the cap removed in the meantime.
  • Safe for many threads, and time comes from a clock that tests control.

Out of scope

  • Saving data to disk and replication.
  • Values other than strings.
  • Isolation between the transactions of different users.
  • Eviction rules other than LRU (an extension below).

Entities

  • KeyValueStore: put, putWithTtl, get, delete, keys, begin, commit, rollback
  • Entry: value, expiresAt; expiredAt(now). Never changed after it is made
  • EvictionPolicy: chooseVictim(keys in order of use); LruEviction, NoEviction
  • Snapshot: a frozen copy of the whole state, the memento
  • TransactionManager: a stack of snapshots, the caretaker; begin, commit, rollback

Classes in the model design

  • KeyValueStore (class)
  • Entry (class)
  • EvictionPolicy (interface)
  • LruEviction (class)
  • NoEviction (class)
  • Snapshot (class)
  • TransactionManager (class)

The main use case, step by step

  1. At minute 0 the caller runs `putWithTtl("session", "tok", 10)`. The store makes an `Entry` with `expiresAt` 10, removes any old entry for the key and puts the new one at the newest end of its map. That map order is the order of use.
  2. At minute 9 `get("session")` finds the entry, sees that it has not expired, moves it to the newest end and returns "tok". At minute 10 the entry is expired. The store removes it and returns null.
  3. With a cap of 3 and keys a, b and c, the call `put("d", ...)` makes the size 4. `makeRoom` drops expired entries first. Then it asks the `EvictionPolicy`, and `LruEviction` answers with the first key in the map, the oldest use.
  4. `begin()` asks the `TransactionManager` to push a `Snapshot`: a copy of the map with its order. The entries are shared and not copied, because an Entry never changes. The call returns depth 1.
  5. Inside the transaction `put` and `delete` work as always, and reads see their effects. Nothing is logged, because the snapshot already holds the old state.
  6. `rollback()` pops the snapshot, and the store replaces its map with a copy of it. Changed values, deleted keys and keys that the cap removed all come back, with the old order of use.
  7. With two `begin()` calls there are two snapshots. An inner `commit()` pops one and keeps the changes. A later outer `rollback()` restores the first snapshot, so the inner changes are undone as well.

Patterns used

  • Memento: Snapshot is the memento, TransactionManager is the caretaker, and KeyValueStore is the originator that creates and restores snapshots.. Rollback needs the old state back, but the caller must not see how the store keeps it. The store saves its own state in a sealed object, and a stack of them gives nesting for free.
  • Strategy: EvictionPolicy with LruEviction and NoEviction.. Which key leaves first is a rule that changes. Behind an interface, LFU or FIFO is a new class, and KeyValueStore does not change.
  • Immutable value object: Entry never changes after it is made. A put creates a new Entry.. A snapshot can then copy only the map and share the entries. No later change can reach into a saved snapshot by accident.

Key code

Entry and the eviction rule

/** Never changed after it is made, so a snapshot can share entries without copying them. */
public final class Entry {
    private final String value;
    private final long expiresAt;          // the first moment the entry is gone; Long.MAX_VALUE means never

    public Entry(String value, long expiresAt) { this.value = value; this.expiresAt = expiresAt; }
    public String value() { return value; }
    public boolean expiredAt(long now) { return now >= expiresAt; }
}

/** Which key goes when the store is full? The map is in order of use, oldest first. */
public interface EvictionPolicy {
    Optional<String> chooseVictim(LinkedHashMap<String, Entry> inUseOrder);
}

public final class LruEviction implements EvictionPolicy {
    @Override public Optional<String> chooseVictim(LinkedHashMap<String, Entry> inUseOrder) {
        return inUseOrder.keySet().stream().findFirst();      // the least recently used key
    }
}

public final class NoEviction implements EvictionPolicy {
    @Override public Optional<String> chooseVictim(LinkedHashMap<String, Entry> inUseOrder) {
        return Optional.empty();
    }
}

Snapshot and TransactionManager: the memento and its caretaker

/** The memento: a frozen copy of the whole state, including the order of use. Only the store can open it. */
public final class Snapshot {
    private final LinkedHashMap<String, Entry> entries;

    Snapshot(Map<String, Entry> state) { this.entries = new LinkedHashMap<>(state); }
    LinkedHashMap<String, Entry> restore() { return new LinkedHashMap<>(entries); }
}

/** The caretaker: it keeps the snapshots in a stack and knows nothing about what is inside. */
public final class TransactionManager {
    private final Deque<Snapshot> stack = new ArrayDeque<>();

    public int begin(Map<String, Entry> state) { stack.push(new Snapshot(state)); return stack.size(); }
    public boolean commit() { return stack.poll() != null; }          // the changes stay, an outer rollback can still undo them
    public Optional<Snapshot> rollback() { return Optional.ofNullable(stack.poll()); }
}

KeyValueStore: expiry, the cap and the transaction calls

public final class KeyValueStore {
    private final int capacity;                              // 0 means no limit
    private final EvictionPolicy policy;
    private final LongSupplier clock;                        // a test passes its own, so TTL needs no sleeping
    private final TransactionManager transactions = new TransactionManager();
    private LinkedHashMap<String, Entry> data = new LinkedHashMap<>();   // oldest use first

    public KeyValueStore(int capacity, EvictionPolicy policy, LongSupplier clock) {
        this.capacity = capacity; this.policy = policy; this.clock = clock;
    }

    public synchronized void put(String key, String value) {
        store(key, new Entry(value, Long.MAX_VALUE));
    }

    public synchronized void putWithTtl(String key, String value, long ttl) {
        if (ttl < 1) throw new IllegalArgumentException("ttl must be at least 1");
        store(key, new Entry(value, clock.getAsLong() + ttl));
    }

    private void store(String key, Entry entry) {
        data.remove(key);                                    // a put is a use: the key goes to the newest end
        data.put(key, entry);
        if (capacity > 0 && data.size() > capacity) makeRoom();
    }

    private void makeRoom() {
        long now = clock.getAsLong();
        data.values().removeIf(e -> e.expiredAt(now));       // expired keys go first
        while (data.size() > capacity) {
            data.remove(policy.chooseVictim(data).orElseThrow());
        }
    }

    /** get changes the order of use, so it is a write: it needs the same lock as put. */
    public synchronized String get(String key) {
        Entry e = data.remove(key);
        if (e == null) return null;
        if (e.expiredAt(clock.getAsLong())) return null;     // expired: leave it removed
        data.put(key, e);
        return e.value();
    }

    public synchronized int begin() { return transactions.begin(data); }
    public synchronized boolean commit() { return transactions.commit(); }

    public synchronized boolean rollback() {
        Optional<Snapshot> saved = transactions.rollback();
        saved.ifPresent(s -> data = s.restore());
        return saved.isPresent();
    }
}

A cheaper rollback: an undo log

/** A cheaper way to roll back: remember only the old entry of each key that changes. */
public final class UndoLog {
    private final Map<String, Optional<Entry>> before = new HashMap<>();   // empty means the key was absent

    /** Call before every change inside a transaction. Only the first change of a key is kept. */
    public void record(String key, Entry oldEntry) {
        before.putIfAbsent(key, Optional.ofNullable(oldEntry));
    }

    public void undo(Map<String, Entry> data) {
        before.forEach((key, old) -> {
            if (old.isPresent()) data.put(key, old.get()); else data.remove(key);
        });
    }
}

Deep dives

How do we undo a transaction?

Weak: Undo by hand. Each method logs its own inverse: a put logs a delete, a delete logs a put. An overwrite loses the old value, and a removal by the size cap is not logged at all. Every new operation must remember to log, and forgetting one breaks rollback with no error.

Good: The Memento pattern. begin saves a copy of the whole map in a Snapshot and pushes it on a stack. rollback pops it and puts the copy back. Every operation, today's and tomorrow's, is covered because the whole state is saved. Entries never change, so the copy holds references and no entry is cloned. Nesting is the stack.

Great: Keep this contract and make begin cheaper when the store is large. An undo log (the last code block) saves only the old entry of each key that changes, so it costs time in proportion to the changes. Or use copy-on-write: begin freezes the current map, writes go to a small overlay, reads check the overlay first, commit merges it and rollback drops it. Order of use needs extra care in both.

Trade-off: The full snapshot is the simplest and obviously correct, and it restores the order of use for free. It costs O(n) time and memory for each open transaction. Choose the undo log or copy-on-write when n is large and transactions are short, and say what you do about the order of use.

How do keys expire?

Weak: A timer thread per key that deletes it when the time is up, or a sweeper that locks the store every second and scans every key. A million keys mean a million timers, or a long pause for every call while the sweeper runs.

Good: Lazy expiry. Store the absolute time expiresAt in the entry and compare it with the clock in get and keys. A key is invisible the moment it expires, with no thread at all. Also drop expired entries when the store needs room. The weak point is memory: a key that nobody reads again stays in the map.

Great: Lazy expiry plus a bounded cleanup. Keep a PriorityQueue of (expiresAt, key). A cleanup step takes only the entries at the head that are due, so its cost follows the number of expired keys and not the size of the store. A heap entry can be stale because the key was overwritten, so compare it with the current entry before removing.

Trade-off: Lazy expiry is the least code and has no background cost, but it can hold dead keys for a long time. The heap bounds the memory and adds a structure and a schedule to run it. In a machine-coding round, do lazy expiry first and name the heap as the next step.

Locking, when get changes the order

Weak: A ConcurrentHashMap for the data and a separate list for the order of use, without a shared lock. A key can leave the list but stay in the map, or the other way round. Two threads can both see room and both add, so the store goes over its cap.

Good: One lock, such as synchronized, around every public method. Check and change are one step, so the cap and the order stay right. Note that get also writes, because it moves the key in the order of use. A read-write lock with get under the read lock would let two readers reorder the map at the same time and corrupt it.

Great: Start with one lock. If a measurement shows contention, split the keys into shards by hash, each with its own lock and its share of the cap. Or make the order approximate: move a key to the newest end only if it has not moved in the last second. A read then often needs no exclusive lock.

Trade-off: Sharding scales, but the cap becomes approximate, and a transaction that snapshots the whole store must lock every shard. Also say that one transaction stack serves all threads of one store here. A real system gives each session its own transaction.

Now add one more feature

  • Now add LFU, which removes the least frequently used key. What changes? Add `LfuEviction`, which keeps a count per key. But the counts are state too, so a snapshot must save them, or rollback leaves wrong counts. Give the policy a way to save and restore its own state.
  • Now support lists and counters. What changes? `Entry` holds a `Value` (a string, a list, a number) and the commands change it. A mutable list inside a shared entry would break the snapshot, so a change must create a new copy of the value.
  • Now survive a restart. What changes? Append each put and delete to a log file before applying it, and replay the file on start. Write a full snapshot to a file now and then, so the log can be cut short.
  • Now give each user an isolated transaction. What changes? Each session keeps its own overlay map. Reads check the overlay and then the store. Commit applies the overlay to the store under the lock, and rollback drops it. Other users never see uncommitted writes.

What each level is expected to show

SDE-1: Model KeyValueStore and Entry, with an absolute expiry time per entry checked on read. Add an LRU cap with a LinkedHashMap or a map plus a list. Explain how rollback could work, for example by copying the map at begin.

SDE-2: Everything at SDE-1, plus an EvictionPolicy strategy, entries that never change, and transactions as a stack of snapshots (Memento) with correct nesting. Drop expired keys before evicting live ones. Use one lock and explain why get is a write. Name the cheaper options for large stores: an undo log, copy-on-write or a heap of expiry times.

Common mistakes

  • Storing the remaining TTL and counting it down, instead of an absolute `expiresAt` compared with the clock. Every tick must then touch every entry.
  • Reading the system clock inside the logic. Tests then need real sleeps, and they are slow and flaky.
  • Copying the map for a snapshot but sharing entries that can change. A later change then reaches into the snapshot, and rollback restores nothing.
  • Evicting a live key while expired keys still use up room in the cap.
  • Putting `get` under a read lock although it reorders the keys.
  • Letting an inner commit throw away the outer snapshot, so an outer rollback can no longer undo the inner changes.
  • Rolling back only the last operation instead of everything since begin.

Lessons behind this answer