KVStore.java

package com.sayandip.kvstore;

import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.ReentrantLock;

/**
 * Small in-memory key-value store with two eviction paths:
 *
 *   - LRU: once the store is at capacity, the least-recently-used entry
 *     is dropped to make room for a new one.
 *   - TTL: entries can carry an expiry. A read of an expired key is a
 *     miss (lazy expiry), and a background sweep also purges expired
 *     entries proactively (active expiry) so dead keys nobody ever reads
 *     again don't sit in memory forever.
 *
 * Lookups go through CustomHashMap (key -> Node), and each Node also
 * sits in a doubly linked list that tracks recency order. That
 * combination gives O(1) get/put/delete while still supporting O(1)
 * "move to front" on every access -- the same trick as the classic LRU
 * Cache problem, just with TTL and a real backing map layered on top.
 */
public class KVStore<K, V> {

    private final int capacity;
    private final CustomHashMap<K, Node<K, V>> map;
    private final ReentrantLock lock = new ReentrantLock();

    private Node<K, V> head; // most recently used
    private Node<K, V> tail; // least recently used

    private final ScheduledExecutorService sweeper;

    public KVStore(int capacity) {
        this(capacity, true);
    }

    public KVStore(int capacity, boolean enableActiveExpirySweep) {
        if (capacity <= 0) throw new IllegalArgumentException("capacity must be positive");
        this.capacity = capacity;
        this.map = new CustomHashMap<>();

        if (enableActiveExpirySweep) {
            sweeper = Executors.newSingleThreadScheduledExecutor(r -> {
                Thread t = new Thread(r, "kvstore-ttl-sweeper");
                t.setDaemon(true);
                return t;
            });
            sweeper.scheduleAtFixedRate(this::sweepExpired, 1, 1, TimeUnit.SECONDS);
        } else {
            sweeper = null;
        }
    }

    public void put(K key, V value) {
        put(key, value, -1);
    }

    /** ttlSeconds < 0 means "no expiry". */
    public void put(K key, V value, long ttlSeconds) {
        long expiresAt = ttlSeconds < 0 ? -1 : System.currentTimeMillis() + ttlSeconds * 1000;
        lock.lock();
        try {
            Node<K, V> existing = map.get(key);
            if (existing != null) {
                existing.value = value;
                existing.expiresAtMillis = expiresAt;
                moveToFront(existing);
                return;
            }
            Node<K, V> node = new Node<>(key, value, expiresAt);
            map.put(key, node);
            addToFront(node);

            if (map.size() > capacity) {
                evictLRU();
            }
        } finally {
            lock.unlock();
        }
    }

    public V get(K key) {
        lock.lock();
        try {
            Node<K, V> node = map.get(key);
            if (node == null) return null;
            if (node.isExpired()) {
                removeNode(node);
                map.remove(key);
                return null;
            }
            moveToFront(node);
            return node.value;
        } finally {
            lock.unlock();
        }
    }

    public void delete(K key) {
        lock.lock();
        try {
            Node<K, V> node = map.remove(key);
            if (node != null) removeNode(node);
        } finally {
            lock.unlock();
        }
    }

    public int size() {
        lock.lock();
        try {
            return map.size();
        } finally {
            lock.unlock();
        }
    }

    public void shutdown() {
        if (sweeper != null) sweeper.shutdownNow();
    }

    // ---- internal doubly linked list bookkeeping (recency order) ----

    private void addToFront(Node<K, V> node) {
        node.prev = null;
        node.next = head;
        if (head != null) head.prev = node;
        head = node;
        if (tail == null) tail = node;
    }

    private void moveToFront(Node<K, V> node) {
        if (head == node) return;
        removeNode(node);
        addToFront(node);
    }

    private void removeNode(Node<K, V> node) {
        if (node.prev != null) node.prev.next = node.next;
        else head = node.next;
        if (node.next != null) node.next.prev = node.prev;
        else tail = node.prev;
        node.prev = null;
        node.next = null;
    }

    private void evictLRU() {
        if (tail == null) return;
        K evictedKey = tail.key;
        removeNode(tail);
        map.remove(evictedKey);
    }

    // O(number of expired entries near the tail) per sweep -- fine at
    // demo scale; a production version would use a min-heap keyed by
    // expiry instead of scanning, worth mentioning as the next step.
    private void sweepExpired() {
        lock.lock();
        try {
            Node<K, V> curr = tail;
            while (curr != null) {
                Node<K, V> prevNode = curr.prev;
                if (curr.isExpired()) {
                    removeNode(curr);
                    map.remove(curr.key);
                }
                curr = prevNode;
            }
        } finally {
            lock.unlock();
        }
    }
}