Transactional in-memory key-value database

"Transactional" can mean two different problems. Ask which one.

1. One client, nested BEGIN / COMMIT / ROLLBACK. A stack of undo logs solves it (see the nested-rollback question in the bank).
2. Several transactions open at once, each with a handle. This is the version this answer covers, because it is the one with real design content. Pin down:
   • Does a transaction see its own writes? (Yes.)
   • Does it see other transactions' uncommitted writes? (No.)
   • What if two transactions touch the same key? Pick an isolation level. A good default: every transaction reads from a snapshot taken at begin, and commit fails with a conflict if anything it read or wrote was changed by a transaction that committed after its begin. That is optimistic concurrency control, and it makes the schedule serializable.
   • Is commit all or nothing? (Yes: a failed commit applies nothing.)

Plain get / set / delete outside a transaction behave as one-operation transactions.

Level 1: private write sets. Each transaction keeps a dict of its pending writes (with a tombstone for deletes). Reads check the write set first, then the committed store. commit copies the write set into the store; rollback drops it. Uncommitted writes are invisible to everyone else because they never touch the store.

Level 2: concurrent transactions. Two more pieces:

• Snapshots with versions (MVCC). A global commit counter clock. Each key keeps a list of (commit_ts, value) versions in increasing order. A transaction records start = clock at begin and reads the newest version with commit_ts <= start (binary search). Its view never changes, even while others commit.
• Validation at commit. The transaction tracks its read set. At commit, for every key it read or wrote, check the key's newest commit_ts. If any is greater than start, someone committed a change it did not see, so raise Conflict and apply nothing. Otherwise increment clock and append the new versions with that timestamp.

Checking the read set (not just the write set) stops write skew and lost updates: two transactions that each read balance and write balance + 10 cannot both commit. Checking only write sets gives snapshot isolation, which is cheaper but allows write skew. Say which one you built.

Garbage collection. A version can go once no open transaction can read it: keep, per key, the newest version at or below the oldest active start, and everything after it.

import bisect
import itertools
from typing import Any, Dict, List, Optional, Set, Tuple

_DELETED = object()  # tombstone in write sets and version lists


class Conflict(Exception):
    pass


class TransactionClosed(Exception):
    pass


class _Txn:
    __slots__ = ("start", "writes", "reads")

    def __init__(self, start: int) -> None:
        self.start = start
        self.writes: Dict[str, Any] = {}
        self.reads: Set[str] = set()


class Database:
    def __init__(self) -> None:
        self._versions: Dict[str, List[Tuple[int, Any]]] = {}
        self._clock = 0
        self._txns: Dict[int, _Txn] = {}
        self._ids = itertools.count(1)

    # ---- transactions ---------------------------------------------------
    def begin(self) -> int:
        tx_id = next(self._ids)
        self._txns[tx_id] = _Txn(self._clock)
        return tx_id

    def get(self, key: str, tx: Optional[int] = None) -> Optional[Any]:
        if tx is None:
            value = self._read(key, self._clock)
        else:
            t = self._txn(tx)
            if key in t.writes:
                value = t.writes[key]
            else:
                t.reads.add(key)
                value = self._read(key, t.start)
        return None if value is _DELETED else value

    def set(self, key: str, value: Any, tx: Optional[int] = None) -> None:
        self._write(key, value, tx)

    def delete(self, key: str, tx: Optional[int] = None) -> None:
        self._write(key, _DELETED, tx)

    def commit(self, tx: int) -> None:
        t = self._txns.pop(tx, None)
        if t is None:
            raise TransactionClosed(tx)
        for key in t.reads | t.writes.keys():
            if self._latest_ts(key) > t.start:
                self._gc()
                raise Conflict(key)
        if t.writes:
            self._clock += 1
            for key, value in t.writes.items():
                self._versions.setdefault(key, []).append((self._clock, value))
        self._gc()

    def rollback(self, tx: int) -> None:
        if self._txns.pop(tx, None) is None:
            raise TransactionClosed(tx)
        self._gc()

    # ---- helpers --------------------------------------------------------
    def _txn(self, tx: int) -> _Txn:
        t = self._txns.get(tx)
        if t is None:
            raise TransactionClosed(tx)
        return t

    def _write(self, key: str, value: Any, tx: Optional[int]) -> None:
        if tx is None:  # autocommit: a one-operation transaction
            tx = self.begin()
            self._txns[tx].writes[key] = value
            self.commit(tx)
        else:
            self._txn(tx).writes[key] = value

    def _read(self, key: str, as_of: int) -> Any:
        versions = self._versions.get(key)
        if not versions:
            return _DELETED
        i = bisect.bisect_left(versions, (as_of + 1,))  # never compares values
        return versions[i - 1][1] if i else _DELETED

    def _latest_ts(self, key: str) -> int:
        versions = self._versions.get(key)
        return versions[-1][0] if versions else 0

    def _gc(self) -> None:
        oldest = min((t.start for t in self._txns.values()), default=self._clock)
        for key in list(self._versions):
            versions = self._versions[key]
            i = bisect.bisect_left(versions, (oldest + 1,))
            if i > 1:
                del versions[: i - 1]           # keep the one visible at `oldest`
            if len(versions) == 1 and versions[0][1] is _DELETED and versions[0][0] <= oldest:
                del self._versions[key]

    def version_count(self) -> int:
        return sum(len(v) for v in self._versions.values())

Searching for the 1-tuple (ts + 1,) finds the first version newer than ts without ever comparing stored values. The full-scan _gc keeps the code short; in production, track which keys gained versions and only prune those.

Let V be the number of versions kept for one key (bounded by how many commits happen while the oldest transaction is open).

Memory is the live data plus old versions that some open transaction can still see, plus each open transaction's read and write sets.

1. Fewer aborts on hot keys. Under contention OCC aborts and retries a lot. Switch hot keys to pessimistic locks (two-phase locking with a wait-for graph or timeouts to break deadlocks), or let the caller retry with backoff.
2. Range reads. "Read all keys with prefix user:" needs phantom protection: also validate that no key was inserted into the range. Keep the keys sorted (a skip list or sorted container) and record scanned ranges in the read set.
3. Durability. Write each commit's write set to a write-ahead log and fsync before acknowledging. Recovery replays the log; periodic snapshots bound replay time.
4. Threads. Here the methods are not thread-safe. The simplest correct version takes one lock around validation and apply (the commit critical section) and none for reads of committed versions, since versions are append-only.
5. Long transactions. One forgotten transaction pins every version written since it began. Add a maximum transaction age and abort old ones.

Mid-level (L4)

• Gets handles, read-your-writes and isolation from uncommitted writes right with a per-transaction write set and a tombstone for deletes.
• Commit is atomic and rollback leaves no trace; operations on a finished transaction fail clearly.
• Tests the common paths and states the cost of each operation.

Senior (L5)

• Introduces versions and snapshots for concurrent transactions and validates at commit.
• Explains the difference between write-set validation (snapshot isolation) and read-set validation (serializable), with a write-skew example.
• Adds version garbage collection based on the oldest open transaction.

Staff (L6+)

• Compares optimistic and pessimistic control for the workload: abort rate under contention, deadlock handling, retry policy.
• Covers range reads and phantoms, durability with a WAL, and thread safety of the commit critical section.
• Names the operational risks: long-running transactions pinning memory, hot keys, and how to observe them.