Key-value store with nested transaction rollback

The base store is a hash map, so set, get and unset are O(1). The real work is begin, rollback and commit with nested transactions.

Brute force. On begin, copy the whole map onto a stack. rollback pops the stack and restores that copy. commit pops it and throws it away. This is correct, but every begin costs O(n) in the size of the store. That fails the usual requirement that transaction cost scale with the keys a transaction touches, not with the store.

Optimal: an undo log per transaction. Keep one live map (data) that always holds the current visible state, so reads never walk a stack. Each open transaction gets an undo log: a dict from key to that key's value before the transaction first touched it, or a MISSING sentinel if the key did not exist.

• Write (set / unset) inside a transaction: if the key is not yet in the top log, record its current value (or MISSING). Then change data. Only the first write per key per transaction is logged, so the log has one entry per touched key, however many times you write it.
• rollback: pop the top log and, for each entry, put the old value back (or delete the key if it was MISSING).
• commit: pop the top log. If there is a parent transaction, merge the child's entries into the parent log, but only for keys the parent has not logged. The parent's older entry already describes the state before the parent began, which is what the parent's rollback needs. If there is no parent, the changes are final and the log is dropped.

Invariants

1. data always equals the state a reader should see right now.
2. For every open transaction T, applying the logs from the top down to and including T's log turns data back into the state at the moment T began.
3. Each log has at most one entry per key, so its size is the number of distinct keys that transaction (and its committed children) changed.

Invariant 2 is why the merge keeps the parent's entry when both logs hold one: the parent's value is older.

commit and rollback act on the innermost transaction and return False when none is open. Interviewers sometimes want an error or a printed NO TRANSACTION instead; ask which.

from typing import Dict, List, Optional


class _Missing:
    """Sentinel: the key did not exist before the transaction touched it."""
    __slots__ = ()


MISSING = _Missing()


class KeyValueStore:
    def __init__(self) -> None:
        self._data: Dict[str, str] = {}
        # One undo log per open transaction, innermost last.
        self._logs: List[Dict[str, object]] = []

    # ---- basic operations -------------------------------------------
    def get(self, key: str) -> Optional[str]:
        return self._data.get(key)

    def set(self, key: str, value: str) -> None:
        self._remember(key)
        self._data[key] = value

    def unset(self, key: str) -> None:
        if key not in self._data:
            return  # no-op; nothing to undo
        self._remember(key)
        del self._data[key]

    # ---- transactions -----------------------------------------------
    def begin(self) -> None:
        self._logs.append({})

    def rollback(self) -> bool:
        if not self._logs:
            return False
        log = self._logs.pop()
        for key, old in log.items():
            if old is MISSING:
                self._data.pop(key, None)
            else:
                self._data[key] = old  # type: ignore[assignment]
        return True

    def commit(self) -> bool:
        if not self._logs:
            return False
        log = self._logs.pop()
        if self._logs:
            parent = self._logs[-1]
            for key, old in log.items():
                # The parent's entry, if any, is older: keep it.
                parent.setdefault(key, old)
        return True

    def in_transaction(self) -> bool:
        return bool(self._logs)

    # ---- helpers ----------------------------------------------------
    def _remember(self, key: str) -> None:
        if not self._logs:
            return
        log = self._logs[-1]
        if key not in log:
            log[key] = self._data.get(key, MISSING)


if __name__ == "__main__":
    kv = KeyValueStore()
    kv.set("a", "1")
    assert kv.get("a") == "1"
    kv.unset("a")
    assert kv.get("a") is None

    # single transaction rollback
    kv.set("a", "1")
    kv.begin()
    kv.set("a", "2")
    kv.set("b", "9")
    assert kv.get("a") == "2" and kv.get("b") == "9"
    assert kv.rollback() is True
    assert kv.get("a") == "1" and kv.get("b") is None

    # nested: inner commit folds into outer, outer rollback undoes both
    kv.begin()
    kv.set("a", "10")
    kv.begin()
    kv.set("a", "20")
    kv.unset("a")
    kv.set("c", "3")
    assert kv.commit() is True
    assert kv.get("a") is None and kv.get("c") == "3"
    assert kv.rollback() is True
    assert kv.get("a") == "1" and kv.get("c") is None

    # nested: inner rollback keeps outer changes
    kv.begin()
    kv.set("x", "1")
    kv.begin()
    kv.set("x", "2")
    kv.rollback()
    assert kv.get("x") == "1"
    kv.commit()
    assert kv.get("x") == "1" and not kv.in_transaction()

    # no transaction open
    assert kv.commit() is False and kv.rollback() is False
    print("ok")

Let k be the number of distinct keys changed in the innermost transaction.

Total space is O(n + sum of log sizes), where n is the number of live keys. Reads never depend on transaction depth, because data is always current.

• set("a","1"); get("a") returns "1"; set("a","2"); get("a") returns "2" (overwrite).
• get("missing") returns None; unset("missing") is a no-op and does not raise.
• commit() or rollback() with no open transaction returns False and changes nothing.
• set("a","1"); begin(); set("a","2"); rollback(); get("a") returns "1".
• begin(); set("b","5"); rollback(); get("b") returns None: a key created inside a rolled-back transaction disappears.
• set("a","1"); begin(); unset("a"); rollback(); get("a") returns "1": rollback restores deleted keys.
• Writing the same key many times in one transaction (set a=2, set a=3, unset a) then rollback() restores the value from before begin.
• Nested: begin(); set a=10; begin(); set a=20; commit(); get a returns "20"; then rollback(); get a returns the value from before the outer begin.
• Nested: begin(); set x=1; begin(); set x=2; rollback(); get x returns "1"; commit() keeps "1".
• Outermost commit() makes changes permanent: a later rollback() returns False and does not undo them.
• Deep nesting (1,000 begin calls, one write each, then 1,000 rollback calls) restores the original store with each call O(1).

1. count(value): how many keys hold a value? Keep a second map value -> count and update it in set and unset (decrement the old value, increment the new one). Rollback goes through the same write path, so the counts stay right with no extra logging. count is O(1).
2. Commit everything. Some versions make commit close all open transactions. Then just clear the stack: O(depth), since data is already current. Rollback stays innermost-only.
3. Point-in-time reads or snapshots. "Read the value as of transaction t" or "open a read-only snapshot" means switching to MVCC. Store each key's history as a list of (version, value) and binary-search it on read. Writes append, a snapshot is just a version number, and old versions are garbage-collected once no snapshot can see them.
4. Concurrent clients. With one transaction per client, keep a private write set per transaction (the layered-overlay design) instead of writing to shared data, and validate at commit. Optimistic: abort if any key you read changed since you began. Pessimistic: take per-key locks. Discuss isolation levels and what "read your own writes" means here.
5. Durability. Add a write-ahead log. Append BEGIN, each SET/UNSET, and COMMIT records, and fsync on the outermost commit. On restart, replay only fully committed transactions, and take snapshots so you can truncate the log.

Mid-level (L4)

• Correct set / get / unset and a working single-level transaction, even if begin copies the whole map at first.
• Handles the edge cases: rollback with no transaction open, unsetting a missing key, a key created and then rolled back.
• Writes small tests as they go and can state the complexity of their version.

Senior (L5)

• Gets to the undo-log (or overlay-stack) design without prompting and explains why a full copy on begin is too expensive.
• Gets nested commit right: merges into the parent, keeps the parent's older entry, and handles deletes with a sentinel or tombstone.
• States the invariant clearly and adds count(value) cleanly when asked.

Staff (L6+)

• Compares the undo log (fast reads, writes in place) with an overlay stack (cheap discard, reads walk the stack) and picks one for the stated workload.
• Extends the design to snapshots/MVCC, concurrent transactions with isolation levels, and durability with a WAL, and names the costs of each.
• Raises API questions early: what commit means with nesting, errors versus return values, value types, memory limits on long transactions.