Design an InMemory Database library

Overview

Databases are used widely across the software industry to store the application data. You must be familiar with some common databases like MySQL, MongoDB, DynamoDb, Redis, etc. In this article we will try to implement Low level design of a library for an InMemory database.

What is InMemory database?

The databases which stores data in the internal memory like Random Access Memory(RAM) rather than an external drive like Solid State Drive(SSD) attached to them are known as InMemory database. These databases offer’s sub-millisecond latency to fetch the data as it eliminates the need to access the external drive’s. These types of databases are highly used in creating cache’s because of their sub-millisecond latency and high throughput.

Requirements

1. Library should store the data in <Key, Value> pair.
2. Library should provide a function to add/update a new key-value pair.
3. Library should provide a function to delete a key-value pair.
4. Library should provide a function to fetch the value corresponding to a given key.
5. Library should be able to handle concurrent access to database. In other words, it should support multi-threaded environment.
6. Library should support different eviction policies.

Design

Class Diagram

1. To give flexibility in choosing the type of storage and eviction policy, we have used strategy pattern. When a user will initialize the client for our library, it can specify what type of storage and eviction policy they want to use.
2. Providing an interface to Storage and eviction policy gives extensibility to our design. If the user want to create a custom type of storage, eviction policy it can create by simply implementing the interface and use the library.
3. To support multi-threaded environment, we have used ReadWriteLock provided java. This provides us with different types of lock for read and write operations. Using write lock it will allow only 1 thread at a time to execute a particular code block.

Code

Interfaces

InMemoryDb

import java.util.List;

public interface InMemoryDb<K, V> {

    public void put(K key, V value);

    public V get(K key);

    public void remove(K key);

}

Storage

public interface Storage<K, V> {

    public void add(K key, V value);

    public V get(K key);

    public V remove(K key);
}

EvictionPolicy

public interface EvictionPolicy<K> {

    public void keyAccessed(K key);

    public void evict();
}

Implementation

InMemoryDbImpl

import inMemoryLibrary.evictionPolicies.EvictionPolicy;
import inMemoryLibrary.exception.StorageFullException;
import inMemoryLibrary.storage.Storage;

import java.util.List;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

public class InMemoryDbImpl<K, V> implements InMemoryDb<K, V> {

    private final Storage<K, V> storage;
    private final EvictionPolicy<K> evictionPolicy;
    private final ReadWriteLock readWriteLock;

    public InMemoryDbImpl(final Storage<K, V> storage, final EvictionPolicy<K> evictionPolicy) {
        this.storage = storage;
        this.evictionPolicy = evictionPolicy;
        this.readWriteLock = new ReentrantReadWriteLock(true);
    }

    @Override
    public void put(K key, V value) {
        readWriteLock.writeLock().lock();
        try {
            System.out.println("Adding key " + key);
            storage.add(key, value);
            evictionPolicy.keyAccessed(key);
            System.out.println("Added key " + key);
        } catch (final StorageFullException e) {
            final K evictedKey = evictionPolicy.evict();
            System.out.println("evicted key " + evictedKey);
            storage.remove(evictedKey);
            put(key, value);
        } finally {
            readWriteLock.writeLock().unlock();
        }
    }

    @Override
    public V get(K key) {
        final V value;
        try {
            readWriteLock.readLock().lock();
            value = storage.get(key);
        } finally {
            readWriteLock.readLock().unlock();
        }
        try {
            readWriteLock.writeLock().lock();
            evictionPolicy.keyAccessed(key);
        } finally {
            readWriteLock.writeLock().unlock();
        }
        return value;
    }

    @Override
    public V remove(K key) {
        readWriteLock.writeLock().lock();
        try {
            System.out.println("Removing key " + key);
            final V value = storage.remove(key);
            evictionPolicy.keyRemoved(key);
            System.out.println("Removed key " + key);
            return value;
        } finally {
            readWriteLock.writeLock().unlock();
        }
    }
}

Multi-threading handling — To handle concurrent access to the library, I used ReentrantReadWriteLock provided by Java. In the given implementation writeLock is used which will give access to the shared code exclusively to only 1 thread at a time. While using the locks remember to unlock the lock acquired so that other threads can access the resource.

HashMapBasedStorage

import java.util.HashMap;
import java.util.Map;

public class HashMapBasedStorage<K, V> implements Storage<K, V> {

    private final int capacity;
    private final Map<K, V> storageMap;

    public HashMapBasedStorage(final int capacity) {
        this.capacity = capacity;
        this.storageMap = new HashMap<>(capacity);
    }

    @Override
    public void add(K key, V value) {
        if (storageMap.size() >= capacity) {
            throw new StorageFullException();
        }
        storageMap.put(key, value);
    }

    @Override
    public V get(K key) {
        if (storageMap.containsKey(key)) {
            return storageMap.get(key);
        }
        return null;
    }

    @Override
    public V remove(K key) {
        if (storageMap.containsKey(key)) {
            return storageMap.remove(key);
        }
        throw new KeyNotFoundException();
    }
}

LeastRecentlyUsedEvictionPolicy (LRUEvictionPolicy)

import java.util.LinkedList;

public class LRUEvictionPolicy<K> implements EvictionPolicy<K> {
    private final LinkedList<K> keyList;

    public LRUEvictionPolicy() {
        keyList = new LinkedList<>();
    }

    @Override
    public void keyAccessed(final K key) {
        if (keyList.contains(key)) {
            keyList.remove(key);
        }
        keyList.add(key);
    }

    @Override
    public void evict() {
        keyList.poll();
    }
}

I tried to use LinkedList library provided by Java. This is a double ended linked list provided java.

In this article we tried to explain the design using 1 eviction policy, but there can be more types of eviction policies like LFU, FIFO and etc.

Future extensions

1. Library should be able to give the top N keys based on most updated keys.
2. Expand the design to handle distributed InMemoryDatabase.

This is one of the approaches that I documented in this article, there is always room for improvements, feel free to drop your suggestions in the comment box. Hope you find this article helpful, keep watching the space and follow for more such future articles.