Skip to main content

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

Comments

Post a Comment

Popular posts from this blog

Designing a Parking Lot - Low Level Design

Problem Statement Design a parking lot that can handle vehicles entering and leaving while managing parking across multiple floors. Each vehicle should be assigned a suitable parking spot based on its type, and the spot should be freed once the vehicle exits. The design should also support generating a ticket at entry and optionally calculating the parking fee based on the duration of stay. Asked In Companies Amazon Google Microsoft Uber Walmart Flipkart Meta PayPal Oracle Salesforce Adobe Apple Intuit LinkedIn Atlassian Functional Requirements The design should support multiple vehicle types such as bikes, cars, and trucks A vehicle must be assigned a parking spot compatible with its type A parking spot cannot be assigned to more than one vehicle at a time The parking lot should support multiple levels (floors) The design should search and allocate an availa...
Post a Comment
Read more

Most Frequently Asked Low Level Design(LLD) Interview Questions

Below are the curated list of most commonly asked Low Level Design (LLD) interview problems. Each problem includes a short description and a link to the complete solution with code and class diagrams. Design Parking Lot System The system should handle parking for different vehicle types such as bikes, cars, and trucks. It should manage slot allocation, availability tracking, and entry/exit flow. The design also ensures efficient usage of parking space under varying load conditions. View Solution Design Elevator / Lift System The system should support multiple elevators operating across floors with request handling logic. It focuses on scheduling algorithms to minimize wait time and optimize movement. It also manages direction control and concurrent floor requests. View Solution Design Movie Ticket Booking System The system should allow users to browse movies, select shows, and book seats. It handles seat ...
Post a Comment
Read more

Software Design Patterns for LLD Interviews: A Complete Guide

Software Design Patterns for LLD Interviews: A Complete Guide In Software Development Engineer (SDE) interviews—especially for mid-level and senior roles—low-level design (LLD) rounds assess your ability to write clean, reusable, maintainable, and extensible code. The foundation of resolving these architectural challenges lies in the standard Gang of Four (GoF) Design Patterns. Rather than memorizing theoretical definitions, interviewers expect you to apply these patterns to real-world scenarios, identifying the trade-offs of each. Below is a comprehensive guide to the 12 most frequently asked design patterns in LLD interviews, categorized by their classification (Creational, Structural, and Behavioral). Each pattern contains a concrete, real-world Java implementation and a detailed breakdown of design decisions. Creational Design Patterns Creational design patterns deal with object creation mechanisms. They abstract the instantiation process, making a system independent of how...
Post a Comment
Read more