Design a In-Memory Key-Value DB with Transactions (Redis-lite with Rollback)

Problem Statement

Design a thread-safe in-memory key-value database supporting nested transactions (similar to Redis or database savepoint stacks). The database must support standard operations: get(key), put(key, value), and delete(key). Additionally, it must support transaction commands: begin() (start a transaction context), commit() (apply all changes globally), and rollback() (discard changes and restore the previous state).

Asked In Companies

Uber Grab Directi

---

Design Decisions & Patterns Used

Managing state during nested transactions requires maintaining temporary write buffers. If a user writes to key “A” inside a transaction, subsequent reads for “A” must return the new temporary value, while other threads outside the transaction must still see the old global value. If the transaction is rolled back, the temporary write buffer is discarded. If committed, the updates are merged globally. We implement this cleanly using a Stack of maps to isolate transaction scopes.

We will utilize the following Design Patterns:

• Memento Pattern: Pushing snapshots (temporary change maps) onto a transaction stack and popping them off to rollback states.
• Command Pattern: Wrapping operations like PUT and DELETE as transactional updates that execute within isolated scopes.

---

Functional Requirements

• Perform basic key-value operations: put(key, value), get(key), and delete(key).
• Start new transaction blocks with begin(). Support nesting multiple transactions.
• Rollback changes using rollback(), restoring key values to their pre-transaction state.
• Commit changes using commit(), applying the updates from the active transaction block globally (or to the parent transaction if nested).
• Ensure read isolation: transactions must not overwrite global variables until committed.

---

Objects Required

• TransactionKeyValueDb (Core engine managing the global store, transaction stacks, and coordination logic)

---

TransactionKeyValueDb Class

The TransactionKeyValueDb class implements the database logic. It uses a stack of maps to track temporary changes, and a special marker string to represent deleted keys inside a transaction.

import java.util.HashMap;
import java.util.Map;
import java.util.Stack;
import java.util.concurrent.ConcurrentHashMap;

public class TransactionKeyValueDb {
    private final Map<String, String> globalStore;
    private final Stack<Map<String, String>> transactionStack;
    
    // Sentinel marker to indicate key deletions within a transaction scope
    private static final String DELETE_MARKER = "__DELETED_KEY_SENTINEL__";

    public TransactionKeyValueDb() {
        this.globalStore = new ConcurrentHashMap<>();
        this.transactionStack = new Stack<>();
    }

    public synchronized void begin() {
        // Push a new transaction map onto the stack
        transactionStack.push(new HashMap<>());
        System.out.println("Transaction started. Stack level: " + transactionStack.size());
    }

    public synchronized void commit() {
        if (transactionStack.isEmpty()) {
            throw new IllegalStateException("No active transaction found to commit.");
        }

        Map<String, String> activeTx = transactionStack.pop();
        
        if (transactionStack.isEmpty()) {
            // Apply changes directly to the global database
            for (Map.Entry<String, String> entry : activeTx.entrySet()) {
                if (entry.getValue().equals(DELETE_MARKER)) {
                    globalStore.remove(entry.getKey());
                } else {
                    globalStore.put(entry.getKey(), entry.getValue());
                }
            }
        } else {
            // Nested Transaction: Merge updates into the parent transaction map
            Map<String, String> parentTx = transactionStack.peek();
            parentTx.putAll(activeTx);
        }
        System.out.println("Transaction committed. Stack level: " + transactionStack.size());
    }

    public synchronized void rollback() {
        if (transactionStack.isEmpty()) {
            throw new IllegalStateException("No active transaction found to rollback.");
        }
        // Discard the active transaction map
        transactionStack.pop();
        System.out.println("Transaction rolled back. Stack level: " + transactionStack.size());
    }

    public synchronized void put(String key, String value) {
        if (transactionStack.isEmpty()) {
            globalStore.put(key, value);
        } else {
            // Write to the active transaction map on top of the stack
            Map<String, String> activeTx = transactionStack.peek();
            activeTx.put(key, value);
        }
    }

    public synchronized String get(String key) {
        // Walk down the transaction stack from top to bottom
        for (int i = transactionStack.size() - 1; i >= 0; i--) {
            Map<String, String> txMap = transactionStack.get(i);
            if (txMap.containsKey(key)) {
                String value = txMap.get(key);
                if (value.equals(DELETE_MARKER)) {
                    return null; // Key was deleted in this transaction
                }
                return value;
            }
        }
        return globalStore.get(key); // Fallback to the global database
    }

    public synchronized void delete(String key) {
        if (transactionStack.isEmpty()) {
            globalStore.remove(key);
        } else {
            // Mark the key as deleted in the active transaction map
            Map<String, String> activeTx = transactionStack.peek();
            activeTx.put(key, DELETE_MARKER);
        }
    }
}

Here is an explanation of the core operations in the TransactionKeyValueDb class:

• The constructor configures the global data map and initializes the transaction stack.
• begin() starts a new transaction by pushing an empty HashMap onto the stack.
• commit() pops the top transaction map. If it was the last transaction, it merges the updates into the global map. If it was a nested transaction, it merges the updates into the parent transaction map below it on the stack.
• rollback() rolls back changes by popping the top map off the stack, discarding all edits in that transaction block.
• put() updates the top map on the stack if a transaction is active; otherwise, it writes directly to the global database.
• get() resolves reads by walking down the transaction stack from top to bottom. If the key is not found in the stack, it falls back to the global database.
• delete() removes the key from the global database if no transaction is active; otherwise, it sets a delete sentinel on the active transaction map to mask the key.

---

Main Driver Class

This class tests our in-memory transactional key-value database. It performs writes, tests transaction rollbacks, commits updates, and verifies nested transaction behaviors.

public class Main {
    public static void main(String[] args) {
        TransactionKeyValueDb db = new TransactionKeyValueDb();

        System.out.println("--- Scenario 1: Basic Operations ---");
        db.put("key1", "value1");
        System.out.println("Global key1 balance: " + db.get("key1"));

        System.out.println("\n--- Scenario 2: Rollback Transaction ---");
        db.begin();
        db.put("key1", "transactional_value");
        System.out.println("Key1 inside transaction: " + db.get("key1"));
        db.rollback();
        System.out.println("Key1 post rollback (Expected: value1): " + db.get("key1"));

        System.out.println("\n--- Scenario 3: Commit Transaction ---");
        db.begin();
        db.put("key2", "value2");
        db.delete("key1");
        db.commit();
        System.out.println("Key2 post commit: " + db.get("key2"));
        System.out.println("Key1 post commit (Expected: null): " + db.get("key1"));

        System.out.println("\n--- Scenario 4: Nested Transactions ---");
        db.put("num", "10");
        db.begin(); // Transaction Level 1
        db.put("num", "20");
        System.out.println("Num at Level 1: " + db.get("num"));

        db.begin(); // Transaction Level 2 (Nested)
        db.put("num", "30");
        System.out.println("Num at Level 2: " + db.get("num"));
        db.rollback(); // Discard Level 2 changes

        System.out.println("Num at Level 1 post Level 2 rollback (Expected: 20): " + db.get("num"));
        db.commit(); // Commit Level 1 changes

        System.out.println("Global Num post Level 1 commit (Expected: 20): " + db.get("num"));
    }
}

The main() driver configures the database, initiates nested transactions, performs updates, tests rollback restorations, and commits changes to confirm correct behavior.

---

Also See

• Most Commonly Asked Low Level Design Interview Questions
• Top 10 Low Level Design Interview Questions