Low Level Design: In-Memory Key-Value Store with Transactions & TTL

"Build an in-memory key-value store with transactions" is one of the most common machine coding prompts for backend roles, and one of the most deceptive. A HashMap handles the first fifteen minutes. Then comes COUNT, which makes naive designs scan the whole map. Then nested transactions, which make naive designs copy the whole map. Then TTL, which interacts with rollback in ways most candidates never consider until the interviewer asks.

This guide treats it as the timed exercise it is: precise semantics first, a comparison of designs, then complete code with a scripted session you can diff against. For the general mechanics of these rounds, read the machine coding round guide; for how a distributed version of this store works, see key-value store system design.

The Machine Coding Prompt

Implement an in-memory key-value store that reads commands from standard input, one per line, and prints the result of each.

• SET key value [EX seconds]: set a value, optionally expiring after the given seconds. A plain SET removes any previous expiry.
• GET key: print the value, or NULL.
• DELETE key: delete; print 1 if it existed, else 0.
• COUNT value: print how many keys currently have exactly this value.
• EXPIRE key seconds, TTL key: set expiry on an existing key; print remaining seconds, -1 if no expiry, -2 if missing.
• BEGIN: open a transaction block. Blocks can be nested.
• ROLLBACK: undo all changes in the most recent open block. Print NO TRANSACTION if none.
• COMMIT: make the most recent block's changes part of its parent (or permanent if outermost). Print NO TRANSACTION if none.

Operations should be efficient: GET, SET, DELETE and COUNT should not depend on the total number of keys.

The last sentence is the real test. It rules out "copy the map on BEGIN" and "scan for COUNT".

Pin Down the Semantics Before Coding

Several behaviours are genuinely ambiguous. Ask, or state your assumption at the top of your file:

Three Ways to Implement Nested Transactions

1. Snapshot on BEGIN

Copy the whole map when a block opens; ROLLBACK restores the copy. Trivially correct, but BEGIN costs O(n) time and memory per nesting level. With a million keys and three nested blocks, that is three million entries copied to change one key. It fails the efficiency requirement.

2. Buffered writes (delta layers)

Each block has its own map of pending writes (with tombstones for deletes). Writes go to the top layer; reads check the top layer, then the next, down to the base. COMMIT merges the top layer into the one below; ROLLBACK discards it. This is how many explanations of the classic "simple database" problem work. The costs: GET becomes O(depth), and COUNT becomes hard, because the count of a value depends on every layer's overrides.

3. Undo log (write in place)

Writes go straight to the main map. The current block keeps an undo map: the first time a key is touched in this block, store its previous entry (or "absent"). ROLLBACK replays the undo map; COMMIT merges it into the parent's undo map, keeping the parent's entry when both have one (the parent's is older). This is the same idea database storage engines use for rollback segments.

The undo log wins on every operation the prompt cares about. The trade-off is that the main map always contains uncommitted data, which is fine for a single client but needs isolation for many clients (covered below). Choosing the design by walking through this table aloud is exactly the reasoning interviewers want to hear.

Making COUNT O(1) with a Secondary Index

Maintain valueCounts: value → number of keys holding it. Every mutation must update it: SET (decrement the old value, increment the new), DELETE, expiry, and every restore during ROLLBACK. The reliable way to guarantee that is a single internal method, rawPut(key, entryOrNull), through which every change flows. Public methods and rollback both call it; nothing else touches the map.

Decrementing to zero should remove the value from the index rather than leave a zero, or the index grows forever with dead values. In Java, computeIfPresent(v, (k, c) -> c == 1 ? null : c - 1) does exactly that.

TTL Expiry Design: Lazy Plus Active

Two classic strategies, and production stores use both. Redis documents this combination in its EXPIRE command reference:

• Lazy (passive) expiry: when a key is read, check its expiry and treat it as absent if passed. Cheap, always correct for reads, but expired keys that are never read stay in memory, and aggregates like COUNT would still see them.
• Active expiry: periodically find and delete expired keys. Redis samples keys with a TTL; a simpler and exact approach for an LLD is a min-heap ordered by expiry time.

This design purges from the heap at the start of every command: pop every heap entry whose time has passed and delete its key. Because keys can be overwritten or re-expired, the heap may contain stale entries. Each entry stores the expiry time it was created with; when popped, it deletes the key only if the key's current expiry equals that time. Stale entries are simply discarded. Each heap operation is O(log h), and every entry is pushed and popped once, so purging is amortised O(log h) per write.

Time comes from an injected Clock. In production, System::currentTimeMillis; in tests and the demo, a manual clock advanced by a SLEEP command. Never call Thread.sleep in a TTL test.

TTL Inside Transactions: The Cases Most Answers Miss

The principle that makes all four cases fall out naturally: expiry is a function of time, not an event. A key with an expiry in the past is absent, whether or not the purge has physically removed it yet. So the undo log never needs to record expirations; it only records writes.

Full Java Code

Java 17+, single file. java KvStoreDemo.java runs the scripted session below; pipe your own commands in with --stdin.

import java.util.*;

public class KvStoreDemo {

    static final long NO_EXPIRY = Long.MAX_VALUE;

    record Entry(String value, long expiresAt) { }

    private record Expiry(long at, String key) { }

    interface Clock { long millis(); }

    static final class ManualClock implements Clock {
        private long now;
        public long millis() { return now; }
        void advance(long ms) { now += ms; }
    }

    static final class KeyValueStore {
        private final Map<String, Entry> data = new HashMap<>();
        private final Map<String, Integer> valueCounts = new HashMap<>();
        private final PriorityQueue<Expiry> expiries = new PriorityQueue<>(Comparator.comparingLong(Expiry::at));
        /** Stack of undo logs. Each maps key -> entry before this block first touched it (null = absent). */
        private final Deque<Map<String, Entry>> transactions = new ArrayDeque<>();
        private final Clock clock;

        KeyValueStore(Clock clock) { this.clock = clock; }

        // ===================== public API =====================
        synchronized String get(String key) {
            purgeExpired();
            Entry e = data.get(key);
            return e == null ? null : e.value();
        }

        synchronized void set(String key, String value) {
            purgeExpired();
            write(key, new Entry(value, NO_EXPIRY));
        }

        synchronized void set(String key, String value, long ttlMillis) {
            if (ttlMillis <= 0) throw new IllegalArgumentException("TTL must be positive");
            purgeExpired();
            write(key, new Entry(value, clock.millis() + ttlMillis));
        }

        synchronized boolean delete(String key) {
            purgeExpired();
            if (!data.containsKey(key)) return false;
            write(key, null);
            return true;
        }

        synchronized boolean expire(String key, long ttlMillis) {
            purgeExpired();
            Entry e = data.get(key);
            if (e == null) return false;
            write(key, new Entry(e.value(), clock.millis() + ttlMillis));
            return true;
        }

        /** Remaining ms; -1 if the key has no expiry; -2 if it does not exist (Redis convention). */
        synchronized long ttl(String key) {
            purgeExpired();
            Entry e = data.get(key);
            if (e == null) return -2;
            if (e.expiresAt() == NO_EXPIRY) return -1;
            return e.expiresAt() - clock.millis();
        }

        synchronized int count(String value) {
            purgeExpired();
            return valueCounts.getOrDefault(value, 0);
        }

        synchronized void begin() { transactions.push(new HashMap<>()); }

        synchronized void rollback() {
            if (transactions.isEmpty()) throw new IllegalStateException("NO TRANSACTION");
            Map<String, Entry> undo = transactions.pop();
            undo.forEach(this::rawPut);                    // restore each key's pre-block entry
        }

        synchronized void commit() {
            if (transactions.isEmpty()) throw new IllegalStateException("NO TRANSACTION");
            Map<String, Entry> undo = transactions.pop();
            Map<String, Entry> parent = transactions.peek();
            if (parent == null) return;                    // outermost: changes are now permanent
            undo.forEach((k, old) -> {
                if (!parent.containsKey(k)) parent.put(k, old);   // parent's older entry wins
            });
        }

        // ===================== internals =====================
        private void write(String key, Entry entry) {
            Map<String, Entry> undo = transactions.peek();
            if (undo != null && !undo.containsKey(key)) undo.put(key, data.get(key));
            rawPut(key, entry);
        }

        /** The ONLY method that mutates data, so the count index and expiry heap never drift. */
        private void rawPut(String key, Entry entry) {
            Entry old = (entry == null) ? data.remove(key) : data.put(key, entry);
            if (old != null) valueCounts.computeIfPresent(old.value(), (v, c) -> c == 1 ? null : c - 1);
            if (entry != null) {
                valueCounts.merge(entry.value(), 1, Integer::sum);
                if (entry.expiresAt() != NO_EXPIRY) expiries.add(new Expiry(entry.expiresAt(), key));
            }
        }

        private void purgeExpired() {
            long now = clock.millis();
            while (!expiries.isEmpty() && expiries.peek().at() <= now) {
                Expiry x = expiries.poll();
                Entry e = data.get(x.key());
                if (e != null && e.expiresAt() == x.at()) rawPut(x.key(), null);   // skip stale heap entries
            }
        }
    }

    // ===================== command layer =====================
    static final class CommandProcessor {
        private final KeyValueStore store;
        private final ManualClock clock;

        CommandProcessor(KeyValueStore store, ManualClock clock) {
            this.store = store;
            this.clock = clock;
        }

        String execute(String line) {
            String[] t = line.trim().split("\\s+");
            String cmd = t[0].toUpperCase();
            try {
                return switch (cmd) {
                    case "SET" -> {
                        if (t.length == 5 && t[3].equalsIgnoreCase("EX")) store.set(t[1], t[2], Long.parseLong(t[4]) * 1000);
                        else if (t.length == 3) store.set(t[1], t[2]);
                        else yield "ERR usage: SET key value [EX seconds]";
                        yield "OK";
                    }
                    case "GET" -> {
                        String v = store.get(t[1]);
                        yield v == null ? "NULL" : v;
                    }
                    case "DELETE" -> store.delete(t[1]) ? "1" : "0";
                    case "COUNT" -> String.valueOf(store.count(t[1]));
                    case "EXPIRE" -> store.expire(t[1], Long.parseLong(t[2]) * 1000) ? "1" : "0";
                    case "TTL" -> {
                        long ms = store.ttl(t[1]);
                        yield ms < 0 ? String.valueOf(ms) : String.valueOf((ms + 999) / 1000);
                    }
                    case "BEGIN" -> { store.begin(); yield "OK"; }
                    case "ROLLBACK" -> { store.rollback(); yield "OK"; }
                    case "COMMIT" -> { store.commit(); yield "OK"; }
                    case "SLEEP" -> {                      // demo/test only: advances the manual clock
                        clock.advance(Long.parseLong(t[1]) * 1000);
                        yield "(clock +" + t[1] + "s)";
                    }
                    default -> "ERR unknown command " + cmd;
                };
            } catch (IllegalStateException e) {
                return e.getMessage();
            } catch (RuntimeException e) {
                return "ERR " + e.getMessage();
            }
        }
    }

    public static void main(String[] args) {
        ManualClock clock = new ManualClock();
        CommandProcessor cli = new CommandProcessor(new KeyValueStore(clock), clock);
        if (args.length > 0 && args[0].equals("--stdin")) {
            Scanner in = new Scanner(System.in);
            while (in.hasNextLine()) {
                String line = in.nextLine();
                if (!line.isBlank()) System.out.println(cli.execute(line));
            }
            return;
        }
        String script = """
            SET a 10
            SET b 10
            COUNT 10
            BEGIN
            SET a 20
            COUNT 10
            BEGIN
            DELETE b
            COUNT 10
            GET b
            ROLLBACK
            GET b
            COUNT 10
            COMMIT
            GET a
            ROLLBACK
            BEGIN
            SET c 1
            BEGIN
            SET c 2
            COMMIT
            GET c
            ROLLBACK
            GET c
            SET session abc EX 5
            TTL session
            SLEEP 3
            TTL session
            BEGIN
            EXPIRE session 10
            SLEEP 4
            GET session
            ROLLBACK
            GET session
            COUNT abc
            TTL a""";
        for (String line : script.split("\n")) System.out.println(line.trim() + "  ->  " + cli.execute(line));
    }
}

Sample Session and What It Proves

SET a 10  ->  OK
SET b 10  ->  OK
COUNT 10  ->  2
BEGIN  ->  OK
SET a 20  ->  OK
COUNT 10  ->  1
BEGIN  ->  OK
DELETE b  ->  1
COUNT 10  ->  0
GET b  ->  NULL
ROLLBACK  ->  OK
GET b  ->  10
COUNT 10  ->  1
COMMIT  ->  OK
GET a  ->  20
ROLLBACK  ->  NO TRANSACTION
BEGIN  ->  OK
SET c 1  ->  OK
BEGIN  ->  OK
SET c 2  ->  OK
COMMIT  ->  OK
GET c  ->  2
ROLLBACK  ->  OK
GET c  ->  NULL
SET session abc EX 5  ->  OK
TTL session  ->  5
SLEEP 3  ->  (clock +3s)
TTL session  ->  2
BEGIN  ->  OK
EXPIRE session 10  ->  1
SLEEP 4  ->  (clock +4s)
GET session  ->  abc
ROLLBACK  ->  OK
GET session  ->  NULL
COUNT abc  ->  0
TTL a  ->  -1

Reading the session

• Lines 1–16, nested rollback with COUNT. The inner DELETE drops COUNT 10 to 0; rolling back the inner block restores b and the count to 1 while keeping the outer block's a = 20. COMMIT at depth 1 makes it permanent, and a further ROLLBACK reports NO TRANSACTION.
• Lines 17–24, savepoint semantics. The inner COMMIT merges into the outer block, so the outer ROLLBACK still erases c entirely. The parent's undo entry ("c was absent") won over the child's ("c was 1"), which is the merge rule from the warning above.
• Lines 25–36, TTL meets rollback. At t=3s the key has 2s left. Inside a transaction its TTL is extended to t=13s, so at t=7s it is still readable. ROLLBACK restores the original entry with expiry t=5s, which is already in the past, so the next GET returns NULL and COUNT for its value is 0. a never had a TTL, so TTL a is -1.

The same session was run through the Python implementation below and produces identical output, which is a useful habit in a real round: when two independent implementations agree on a tricky script, you can trust both.

Python Version

If your machine coding round is in Python, the same design is shorter. heapq provides the expiry heap and a list serves as the transaction stack.

import heapq

NO_EXPIRY = float("inf")

class KeyValueStore:
    def __init__(self, clock):
        self.clock = clock          # callable returning current time in ms
        self.data = {}              # key -> (value, expires_at)
        self.counts = {}            # value -> number of live keys holding it
        self.expiries = []          # min-heap of (expires_at, key)
        self.tx = []                # stack of undo logs: key -> old entry or None

    def get(self, key):
        self._purge()
        e = self.data.get(key)
        return None if e is None else e[0]

    def set(self, key, value, ttl_ms=None):
        self._purge()
        exp = NO_EXPIRY if ttl_ms is None else self.clock() + ttl_ms
        self._write(key, (value, exp))

    def delete(self, key):
        self._purge()
        if key not in self.data:
            return False
        self._write(key, None)
        return True

    def expire(self, key, ttl_ms):
        self._purge()
        e = self.data.get(key)
        if e is None:
            return False
        self._write(key, (e[0], self.clock() + ttl_ms))
        return True

    def ttl(self, key):
        self._purge()
        e = self.data.get(key)
        if e is None:
            return -2
        return -1 if e[1] == NO_EXPIRY else e[1] - self.clock()

    def count(self, value):
        self._purge()
        return self.counts.get(value, 0)

    def begin(self):
        self.tx.append({})

    def rollback(self):
        if not self.tx:
            raise RuntimeError("NO TRANSACTION")
        for key, old in self.tx.pop().items():
            self._raw_put(key, old)

    def commit(self):
        if not self.tx:
            raise RuntimeError("NO TRANSACTION")
        undo = self.tx.pop()
        if self.tx:                                   # parent keeps its older entry
            parent = self.tx[-1]
            for key, old in undo.items():
                parent.setdefault(key, old)

    def _write(self, key, entry):
        if self.tx and key not in self.tx[-1]:
            self.tx[-1][key] = self.data.get(key)     # first touch in this block
        self._raw_put(key, entry)

    def _raw_put(self, key, entry):
        old = self.data.pop(key, None) if entry is None else self.data.get(key)
        if entry is not None:
            self.data[key] = entry
        if old is not None:
            self.counts[old[0]] -= 1
            if self.counts[old[0]] == 0:
                del self.counts[old[0]]
        if entry is not None:
            self.counts[entry[0]] = self.counts.get(entry[0], 0) + 1
            if entry[1] != NO_EXPIRY:
                heapq.heappush(self.expiries, (entry[1], key))

    def _purge(self):
        now = self.clock()
        while self.expiries and self.expiries[0][0] <= now:
            at, key = heapq.heappop(self.expiries)
            e = self.data.get(key)
            if e is not None and e[1] == at:          # ignore stale heap entries
                self._raw_put(key, None)

Python's dict.setdefault is safe for the commit merge because it checks key presence, not truthiness, unlike Java's putIfAbsent with null values.

Thread Safety and Multiple Clients

The Java store synchronises every public method, so concurrent calls are linearisable. That is the right answer for the prompt as stated (one client, one transaction stack). The interesting follow-up is "now several clients connect, each with its own transactions". The undo log design writes uncommitted data into the shared map, so client B would see client A's uncommitted writes, a dirty read. Options, in increasing sophistication:

1. Global transaction lock. BEGIN acquires a store-wide lock held until the outermost COMMIT or ROLLBACK. Correct (serialisable), simple, and kills concurrency. Acceptable as a first answer if you name the cost.
2. Per-session buffered writes. Switch to delta layers per session; commit applies the buffer under the store lock. Readers never see uncommitted data. Lost updates are possible unless you add conflict checks.
3. Optimistic concurrency. Record a version per key read or written during the transaction; at commit, verify versions are unchanged, else abort. This is what Redis's WATCH provides. Note that Redis MULTI/EXEC has no rollback at all: commands are queued and executed atomically.
4. MVCC. Keep multiple versions per key with commit timestamps; each transaction reads a snapshot. This is how PostgreSQL and many key-value engines provide snapshot isolation.

Being able to map these options to isolation levels (read committed, snapshot, serialisable) is the kind of depth that lifts an LLD answer; the ACID and isolation levels guide covers them, and Redis deep dive covers how Redis actually behaves.

Other extensions

• Memory limit with eviction: add an LRU list (doubly linked list + map), exactly LeetCode 146 LRU Cache; evictions must also go through rawPut and must not be recorded in undo logs.
• Persistence: append each committed command to a log file (like Redis AOF) and replay on startup; take periodic snapshots to bound replay time.
• Range queries: swap HashMap for TreeMap (O(log n) operations) or keep a sorted secondary index.
• Background expiry: a scheduled task that calls purgeExpired() every 100 ms reclaims memory even when no commands arrive.
• Heap compaction: if keys are re-expired very often, stale heap entries accumulate; rebuild the heap when stale entries exceed live ones.

How This Round Is Evaluated (and a Time Plan)

What evaluators look for

• COUNT that does not scan, and an index that stays correct through rollback and expiry.
• No full-map copies on BEGIN.
• A single mutation path (rawPut) rather than index updates sprinkled in five methods.
• An injected clock and deterministic TTL tests.
• Clear, stated semantics for COMMIT and TTL, and the ability to change them quickly when asked.

Common mistakes

• Copying the map on BEGIN.
• Recording expirations in the undo log, which resurrects expired keys on rollback.
• Restoring only the value on rollback and losing the original TTL.
• Using putIfAbsent with null-valued undo entries.
• Counting expired keys in COUNT because only lazy expiry was implemented.
• Thread.sleep in tests.

For more practice in the same format, try the snake and ladder machine coding problem for game-loop structure, and see the LLD interview guide for which problems are most commonly asked.

Frequently Asked Questions

Pranjal Jain

Ex-Microsoft SDE · IIT Kanpur · Founder of Prepflix. Helps engineers crack startup and product company interviews across India.