Key-value database with concurrent transaction ids

Build an in-memory key-value database that many clients use at once. Every read and write happens inside a transaction, and several transactions can be open at the same time. A transaction sees its own writes right away; nobody else sees them until it commits. One process, one machine, no networking; callers use the API directly.

Interface

Unless you agree on a different interface with the interviewer, assume:

begin() -> int — opens a transaction and returns its id.
get(txid, key) -> value | None — None when the key has no value.
put(txid, key, value) -> None — value=None deletes the key.
commit(txid) -> None — makes the transaction's writes visible to transactions that begin later.
rollback(txid) -> None — throws the transaction's writes away.

The spec arrives in three parts; finish and test each before the next.

Part 1. One transaction at a time: reads see the transaction's own writes, and a commit makes them visible to the next transaction.

Part 2. Many open transactions. Uncommitted writes stay private, rollback discards them, and None deletes. A transaction reads from a snapshot taken when it began: a commit by someone else after that is invisible to it.

Part 3. Two open transactions write the same key. The first to commit wins; the second commit raises ConflictError and ends that transaction. Any call with an id that is unknown, committed or rolled back raises TransactionError. Make the database safe to call from several threads.

Example

a = begin()
put(a, "seats", "40")
commit(a)

b = begin()
c = begin()
get(b, "seats")          -> "40"
put(b, "seats", "39")     # b sells a seat
get(c, "seats")          -> "40"   # b hasn't committed
put(c, "seats", "38")     # c sells two
commit(b)
get(c, "seats")          -> "38"   # its own write
commit(c)                -> ConflictError (b changed "seats" first)

d = begin()
get(d, "seats")          -> "39"
put(d, "seats", None)
commit(d)
get(begin(), "seats")    -> None

class TransactionError(Exception): ...


class ConflictError(TransactionError): ...


class KVDatabase:
    def begin(self) -> int: ...
    def get(self, txid: int, key: str): ...
    def put(self, txid: int, key: str, value) -> None: ...
    def commit(self, txid: int) -> None: ...
    def rollback(self, txid: int) -> None: ...

Clarifying questions

Is put(txid, key, None) a delete? (Yes.)
After another transaction commits, does a still-open transaction see the new value (read committed), or keep reading the snapshot from its begin (repeatable reads)? Say which you pick and why.
If two open transactions write the same key and both commit, should the second silently overwrite the first? (No: that loses an update.)
What happens to a transaction whose commit fails? Can it be retried with the same id?
Do read-only transactions ever conflict?
How many keys and how many open transactions at once?

What a strong answer covers

Private writes per transaction (a dict of buffered writes, with a marker for deletes so a delete isn't mistaken for "not written"), checked before the shared store on every get.
Snapshot reads: a commit counter as the version, each transaction remembers the version at begin, and each key keeps its committed values with the version that wrote them, so a read is a binary search for the newest value at or below the snapshot.
Commit as the single point where writes become visible, and why writing to the shared store in put breaks isolation.
Write conflicts: at commit, a key this transaction wrote that was committed after its snapshot means someone else got there first; raise and end the transaction (first committer wins).
Clear errors for unknown and finished ids, and one lock around each operation for thread safety, with the reasoning for why that is enough on a single core.
Tests: the example, private writes, rollback, deletes, snapshot reads, disjoint keys and read-only transactions committing, and a threaded retry loop that loses no increments.

Follow-ups

Read committed instead of snapshots: what changes in get, and which anomaly does it allow?
Write skew: two transactions read the same keys and write different ones. Does first-committer-wins catch it? What would?
Old versions pile up: when can a key's history be trimmed?
Nested transactions inside one id (savepoints).
Make it durable with a write-ahead log.

from solution import KVDatabase


def test_a_transaction_reads_its_own_writes():
    db = KVDatabase()
    t = db.begin()
    assert db.get(t, "views") is None
    db.put(t, "views", "10")
    assert db.get(t, "views") == "10"
    db.commit(t)


def test_a_later_transaction_sees_the_commit():
    db = KVDatabase()
    t = db.begin()
    db.put(t, "views", "10")
    db.commit(t)
    u = db.begin()
    assert db.get(u, "views") == "10"


def test_ids_are_distinct():
    db = KVDatabase()
    assert len({db.begin() for _ in range(100)}) == 100

from solution import KVDatabase


def test_uncommitted_writes_stay_private():
    db = KVDatabase()
    a, b = db.begin(), db.begin()
    db.put(a, "color", "red")
    assert db.get(b, "color") is None
    db.put(b, "color", "blue")
    assert db.get(a, "color") == "red"
    assert db.get(b, "color") == "blue"


def test_rollback_discards_everything():
    db = KVDatabase()
    seed = db.begin()
    db.put(seed, "color", "red")
    db.commit(seed)
    t = db.begin()
    db.put(t, "color", "green")
    db.put(t, "size", "L")
    db.rollback(t)
    u = db.begin()
    assert db.get(u, "color") == "red"
    assert db.get(u, "size") is None


def test_none_deletes():
    db = KVDatabase()
    seed = db.begin()
    db.put(seed, "color", "red")
    db.commit(seed)
    t = db.begin()
    db.put(t, "color", None)
    assert db.get(t, "color") is None
    other = db.begin()
    assert db.get(other, "color") == "red"
    db.commit(t)
    assert db.get(db.begin(), "color") is None
    db.put(other, "size", "M")
    db.commit(other)  # other never wrote "color": no conflict


def test_reads_come_from_the_snapshot_taken_at_begin():
    db = KVDatabase()
    early = db.begin()
    writer = db.begin()
    db.put(writer, "color", "red")
    db.commit(writer)
    assert db.get(early, "color") is None   # began before the commit
    late = db.begin()
    assert db.get(late, "color") == "red"

import threading

from solution import ConflictError, KVDatabase, TransactionError


def raises(exc, fn, *args):
    try:
        fn(*args)
    except exc:
        return True
    return False


def test_first_committer_wins():
    db = KVDatabase()
    a, b = db.begin(), db.begin()
    db.put(a, "stock", "9")
    db.put(b, "stock", "8")
    db.commit(a)
    assert raises(ConflictError, db.commit, b)
    assert db.get(db.begin(), "stock") == "9"


def test_a_failed_commit_ends_the_transaction():
    db = KVDatabase()
    a, b = db.begin(), db.begin()
    db.put(a, "stock", "9")
    db.put(b, "stock", "8")
    db.commit(a)
    assert raises(ConflictError, db.commit, b)
    assert raises(TransactionError, db.get, b, "stock")


def test_different_keys_and_read_only_transactions_do_not_conflict():
    db = KVDatabase()
    a, b, reader = db.begin(), db.begin(), db.begin()
    db.put(a, "x", "1")
    db.put(b, "y", "2")
    assert db.get(reader, "x") is None
    db.commit(a)
    db.commit(b)
    db.commit(reader)
    t = db.begin()
    assert (db.get(t, "x"), db.get(t, "y")) == ("1", "2")


def test_a_delete_conflicts_like_a_write():
    db = KVDatabase()
    seed = db.begin()
    db.put(seed, "k", "v")
    db.commit(seed)
    a, b = db.begin(), db.begin()
    db.put(a, "k", None)
    db.put(b, "k", "w")
    db.commit(b)
    assert raises(ConflictError, db.commit, a)


def test_finished_and_unknown_ids_raise():
    db = KVDatabase()
    t = db.begin()
    db.commit(t)
    r = db.begin()
    db.rollback(r)
    for txid in (t, r, 999):
        assert raises(TransactionError, db.get, txid, "k")
        assert raises(TransactionError, db.put, txid, "k", "v")
        assert raises(TransactionError, db.commit, txid)
        assert raises(TransactionError, db.rollback, txid)


def test_concurrent_increments_lose_nothing():
    db = KVDatabase()
    seed = db.begin()
    db.put(seed, "n", "0")
    db.commit(seed)

    def worker():
        for _ in range(200):
            while True:
                t = db.begin()
                db.put(t, "n", str(int(db.get(t, "n")) + 1))
                try:
                    db.commit(t)
                    break
                except ConflictError:
                    continue

    threads = [threading.Thread(target=worker) for _ in range(4)]
    for th in threads:
        th.start()
    for th in threads:
        th.join()
    assert db.get(db.begin(), "n") == "800"