System Design: Building an In-Memory Key-Value Store (JS)

Design a simple in memory key-value store which provides, basic get, set and delete operations. Once done, add transactional support to it, where all changes for a transaction will be applied either all as a whole or all changes can be rolled back for a transaction. Also share some thoughts on if multiple transactions are happening in parallel then how your system would behave.

Introduction

Key-value stores are simple yet powerful constructs, often used as in-memory databases or caches due to their efficient data retrieval capabilities. They associate each key with one, and only one, value, hence the term “key-value store”.

In this blog, we’ll demonstrate how to create a basic in-memory key-value store in JavaScript. Our store will offer three primary functions: get(key), set(key, value), and delete(key).

Basic Key-Value Data Structure

Let’s begin by defining our basic KeyValueStore class.

class KeyValueStore {
    constructor() {
        this.store = {};
    }

    get(key) {
        // To be implemented.
    }

    set(key, value) {
        // To be implemented.
    }

    delete(key) {
        // To be implemented.
    }
}

In the constructor, we initialize an empty JavaScript object ({}), which will serve as our data store. The get(), set(), and delete() methods are currently placeholders that we will soon flesh out.

The get() Function

The get() function retrieves the value associated with a given key. If the key doesn't exist in the store, it should return null.

class KeyValueStore {
    // ...
    get(key) {
        return this.store[key] || null;
    }
    // ...
}

The get() method checks if the key exists in our store and returns its value. If it doesn't exist, it returns null.

The set() Function

The set() function is responsible for adding new key-value pairs to our store or updating the value of an existing key.

class KeyValueStore {
    // ...
    set(key, value) {
        this.store[key] = value;
    }
    // ...
}

The set() method updates the value of a key if it already exists, or creates a new key with the provided value if it doesn't.

The delete() Function

Finally, the delete() function removes a key-value pair from our store. If the key doesn't exist, it should do nothing.

class KeyValueStore {
    // ...
    delete(key) {
        delete this.store[key];
    }
}

The delete() method removes a key from our store if it exists. If the key doesn't exist, the delete operation is a no-op.

With this, we now have a simple yet functional in-memory key-value store in JavaScript. This can serve as a starting point for more complex data structures, caches, or in-memory databases in JavaScript.

Do note that the above key-value store does not persist data, meaning if your program terminates or crashes, all data will be lost. If you need persistent data storage, consider using a database or a file system-based approach.

Adding Transactional Support to the Store

To implement transactions, we’ll use a stack for each active transaction. Each element in the stack will represent a transaction with a unique transaction ID. This stack will maintain all changes made during a particular transaction.

Each transaction will store a map of keys to their original values at the time the transaction began. This allows us to “undo” changes if a transaction is rolled back.

Let’s break down the transactional methods:

• beginTransaction(): This method creates a new transaction and adds it to the stack. A transaction consists of an id and a changes object, which will hold the original state of any keys modified during the transaction. The transaction id is simply the current timestamp, but in a real system you might want a more robust way of generating unique ids.
• commitTransaction(transactionId): This method applies the changes from the specified transaction. It first finds the transaction in the stack and then applies any changes to the main store. It then removes the committed transaction and all transactions that were started after it from the transaction stack.
• rollbackTransaction(transactionId): This method undoes the changes from the specified transaction. It first finds the transaction in the stack, then restores the original values of any keys that were modified. It then removes the rolled back transaction and all transactions that were started after it from the transaction stack.

Here’s the updated KeyValueStore class with transaction support:

class KeyValueStore {
    constructor() {
        this.store = {};
        this.transactionStack = [];
    }

    beginTransaction() {
        const transaction = { id: Date.now(), changes: {} };
        this.transactionStack.push(transaction);
        return transaction.id;
    }

    commitTransaction(transactionId) {
        // To be implemented.
    }

    rollbackTransaction(transactionId) {
        // To be implemented.
    }

    get(key) {
        // To be implemented with transactional support.
    }

    set(key, value) {
        // To be implemented with transactional support.
    }

    delete(key) {
        // To be implemented with transactional support.
    }
}

beginTransaction()

The beginTransaction method is responsible for starting a new transaction. When called, it creates a new transaction object that consists of two properties:

• id: This is a unique identifier for the transaction. We're using the current timestamp (Date.now()) as the id, but in a real system, you might want a more sophisticated way to generate unique ids.
• changes: This is an object that will hold the original state of any keys modified during the transaction. It starts out empty because no changes have been made when the transaction is first started.

After creating the transaction object, it’s added to the transaction stack and its id is returned. The transaction stack allows us to keep track of multiple ongoing transactions.

avascribeginTransaction() {
    const transaction = { id: Date.now(), changes: {} };
    this.transactionStack.push(transaction);
    return transaction.id;
}

commitTransaction(transactionId)

The commitTransaction method is responsible for making all the changes made during a transaction permanent. It takes one parameter, transactionId, which is used to find the transaction in the transaction stack.

When this method is called, it finds the specified transaction in the transaction stack. It then applies any changes from that transaction to the main store. Because of how we’ve implemented the set and delete methods, the changes are actually applied to the main store as soon as they're made, so there's nothing to do here.

After applying the changes, it removes the committed transaction and all transactions that were started after it from the transaction stack.

commitTransaction(transactionId) {
    const transactionIndex = this.transactionStack.findIndex(t => t.id === transactionId);
    if (transactionIndex !== -1) {
        // All changes in transaction are already applied to the store
        this.transactionStack = this.transactionStack.slice(0, transactionIndex);
    } else {
        throw new Error(`Transaction with ID ${transactionId} not found`);
    }
}

rollbackTransaction(transactionId)

The rollbackTransaction method is responsible for discarding all the changes made during a transaction. Like commitTransaction, it takes the transactionId as a parameter and uses it to find the transaction in the transaction stack.

When this method is called, it first finds the specified transaction in the transaction stack. It then goes through each key in the transaction’s changes object and restores the original value in the main store. If the original value was null, it means the key was newly added during the transaction, so it's deleted from the store.

After rolling back the changes, it removes the rolled back transaction and all transactions that were started after it from the transaction stack.javascriptCopy code

rollbackTransaction(transactionId) {
    const transactionIndex = this.transactionStack.findIndex(t => t.id === transactionId);
    if (transactionIndex !== -1) {
        // Rollback changes
        const transaction = this.transactionStack[transactionIndex];
        for (let key in transaction.changes) {
            if (transaction.changes[key] === null) {
                delete this.store[key];
            } else {
                this.store[key] = transaction.changes[key];
            }
        }
        this.transactionStack = this.transactionStack.slice(0, transactionIndex);
    } else {
        throw new Error(`Transaction with ID ${transactionId} not found`);
    }
}

Full Implementation

Here is the full implementation of all the pieces together for a key-value store with transactional support.

class KeyValueStore {
    constructor() {
        this.store = {};
        this.transactionStack = [];
    }

    get(key) {
        return this.store[key];
    }

    set(key, value) {
        this.trackChange(key);
        this.store[key] = value;
    }

    delete(key) {
        this.trackChange(key);
        delete this.store[key];
    }

    trackChange(key) {
        if (this.transactionStack.length > 0) {
            const transaction = this.transactionStack[this.transactionStack.length - 1];
            if (!(key in transaction.changes)) {
                // Keep track of original value for rollback purposes
                transaction.changes[key] = this.store[key] || null;
            }
        }
    }

    beginTransaction() {
        const transaction = { id: Date.now(), changes: {} };
        this.transactionStack.push(transaction);
        return transaction.id;
    }

    commitTransaction(transactionId) {
        const transactionIndex = this.transactionStack.findIndex(t => t.id === transactionId);
        if (transactionIndex !== -1) {
            // All changes in transaction are already applied to the store
            this.transactionStack = this.transactionStack.slice(0, transactionIndex);
        } else {
            throw new Error(`Transaction with ID ${transactionId} not found`);
        }
    }

    rollbackTransaction(transactionId) {
        const transactionIndex = this.transactionStack.findIndex(t => t.id === transactionId);
        if (transactionIndex !== -1) {
            // Rollback changes
            const transaction = this.transactionStack[transactionIndex];
            for (let key in transaction.changes) {
                if (transaction.changes[key] === null) {
                    delete this.store[key];
                } else {
                    this.store[key] = transaction.changes[key];
                }
            }
            this.transactionStack = this.transactionStack.slice(0, transactionIndex);
        } else {
            throw new Error(`Transaction with ID ${transactionId} not found`);
        }
    }
}

Thoughts on Multiple Transactions in Parallel

In our current implementation, transactions follow a stack structure, or last-in-first-out (LIFO) order, meaning the last transaction that started needs to be the first one to either commit or rollback. If a transaction in the middle of the stack tries to commit or rollback, it will throw an error because the transaction isn’t at the top of the stack.

This LIFO order is because of the way we use Array.prototype.findIndex to locate the transaction within the transaction stack and the way we slice the transaction stack when committing or rolling back. If a transaction isn’t at the top of the stack (i.e., it’s not the most recent transaction), it will have transactions above it in the stack that would also get removed when slicing the stack, which could cause unexpected behavior.

The current design assumes that the user will manage the order of transactions appropriately, making sure to close transactions in the reverse order they were started. It’s a simple design suitable for scenarios where you can easily control the transaction order, but it might not be appropriate for more complex situations where transactions need to be able to commit or rollback independently.

In a more advanced system, you would likely want to support nested transactions or implement concurrency control mechanisms such as locks or multiversion concurrency control (MVCC) to handle more complex transaction scenarios. These concepts are a bit beyond the scope of our simple key-value store but are crucial for understanding real-world database transaction management.

Conclusion

With this, we now have a JavaScript key-value store with basic transactional support. As you can see, transactions can add quite a bit of complexity, but they also provide powerful capabilities for ensuring data integrity and consistency.

This is a basic implementation, primarily designed for learning purposes. In a production environment, more advanced data structures and concurrency control mechanisms would likely be needed for optimal performance and reliability. Nevertheless, this simple key-value store serves as a great introduction to how transactions can be implemented in an in-memory database.