In-memory key-value store

import java.util.*;
import java.util.concurrent.atomic.LongAdder;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/** Time behind an interface — the one seam that makes TTL testable without sleeping. */
interface Clock { long nowMillis(); }

final class SystemClock implements Clock {
    public long nowMillis() { return System.currentTimeMillis(); }
}

/** A clock you move by hand. Every TTL test in this file runs in microseconds because of it. */
final class FakeClock implements Clock {
    private long now;
    FakeClock(long startMillis) { this.now = startMillis; }
    public long nowMillis() { return now; }
    void advance(long millis) { now += millis; }
}

/**
 * Immutable value object. Holds an ABSOLUTE instant, never a countdown:
 * nothing has to tick it, and a million keys still means zero running timers.
 */
final class Entry {
    static final long NEVER = Long.MAX_VALUE;

    private final String value;
    private final long expiresAtMillis;

    Entry(String value, long expiresAtMillis) {
        this.value = Objects.requireNonNull(value);
        this.expiresAtMillis = expiresAtMillis;
    }

    String value()            { return value; }
    long expiresAtMillis()    { return expiresAtMillis; }
    boolean neverExpires()    { return expiresAtMillis == NEVER; }
    boolean isExpired(long now) { return now >= expiresAtMillis; }
}

/** Three genuinely different answers. A single int that means all three is a bad API. */
sealed interface TtlResult {
    record Remaining(long millis) implements TtlResult {
        public String toString() { return millis + "ms left"; }
    }
    record NoExpiry() implements TtlResult {
        public String toString() { return "no expiry"; }
    }
    record NoSuchKey() implements TtlResult {
        public String toString() { return "no such key"; }
    }
}

final class KeyValueStore {
    private static final int SAMPLE_SIZE = 20;
    private static final int MAX_SWEEP_ROUNDS = 16;   // so one pass ALWAYS terminates

    private final Map<String, Entry> map = new HashMap<>();

    // Only the keys that HAVE a deadline. The sweeper samples from here, never from
    // the whole keyspace — otherwise almost every draw is wasted on permanent keys.
    private final List<String> volatileKeys = new ArrayList<>();
    private final Map<String, Integer> volatileIndex = new HashMap<>();

    private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
    private final Clock clock;
    private final int maxEntries;
    private final Random random = new Random(7);

    private final LongAdder hits = new LongAdder(), misses = new LongAdder();
    private final LongAdder lazyExpired = new LongAdder(), activeExpired = new LongAdder();
    private final LongAdder evicted = new LongAdder();

    KeyValueStore(Clock clock, int maxEntries) {
        this.clock = clock;
        this.maxEntries = maxEntries;
    }

    /* ------------------------------------------------------------------ reads */

    /**
     * A get really IS a read here — until it finds an expired key, and then it must WRITE.
     * Java's ReadWriteLock cannot upgrade in place, so we drop the read lock, take the
     * write lock, and RE-CHECK: another thread may have written a fresh value in the gap.
     */
    Optional<String> get(String key) {
        long now = clock.nowMillis();
        lock.readLock().lock();
        try {
            Entry e = map.get(key);
            if (e == null) { misses.increment(); return Optional.empty(); }
            if (!e.isExpired(now)) { hits.increment(); return Optional.of(e.value()); }
        } finally {
            lock.readLock().unlock();
        }

        lock.writeLock().lock();
        try {
            Entry again = map.get(key);
            if (again != null && again.isExpired(clock.nowMillis())) {
                removeLocked(key);
                lazyExpired.increment();
            }
            misses.increment();          // this call misses either way — it saw a dead key
            return Optional.empty();
        } finally {
            lock.writeLock().unlock();
        }
    }

    boolean exists(String key) { return get(key).isPresent(); }   // SAME expiry rule as get

    /** Three states, never overloaded onto one number. */
    TtlResult ttl(String key) {
        long now = clock.nowMillis();
        lock.readLock().lock();
        try {
            Entry e = map.get(key);
            if (e == null || e.isExpired(now)) return new TtlResult.NoSuchKey();
            if (e.neverExpires())               return new TtlResult.NoExpiry();
            return new TtlResult.Remaining(e.expiresAtMillis() - now);
        } finally {
            lock.readLock().unlock();
        }
    }

    /** The PHYSICAL entry count — expired-but-not-yet-reclaimed keys are included, on purpose. */
    int size() {
        lock.readLock().lock();
        try { return map.size(); } finally { lock.readLock().unlock(); }
    }

    /** What a caller could actually read right now. Usually smaller than size(). */
    int liveSize() {
        long now = clock.nowMillis();
        lock.readLock().lock();
        try {
            int n = 0;
            for (Entry e : map.values()) if (!e.isExpired(now)) n++;
            return n;
        } finally { lock.readLock().unlock(); }
    }

    /** O(n), and it holds the lock the whole way. This is exactly why SCAN with a cursor exists. */
    List<String> keys(String pattern) {
        long now = clock.nowMillis();
        lock.readLock().lock();
        try {
            List<String> out = new ArrayList<>();
            for (Map.Entry<String, Entry> e : map.entrySet())
                if (!e.getValue().isExpired(now) && glob(pattern, e.getKey())) out.add(e.getKey());
            Collections.sort(out);
            return out;
        } finally { lock.readLock().unlock(); }
    }

    /* ----------------------------------------------------------------- writes */

    /** Plain set CLEARS any existing deadline — the key becomes permanent. Real stores do this too. */
    void set(String key, String value) { set(key, value, false); }

    void set(String key, String value, boolean keepTtl) {
        lock.writeLock().lock();
        try {
            Entry old = map.get(key);
            long expiresAt = Entry.NEVER;
            if (keepTtl && old != null && !old.isExpired(clock.nowMillis()))
                expiresAt = old.expiresAtMillis();
            putLocked(key, new Entry(value, expiresAt));
            evictIfNeededLocked();
        } finally { lock.writeLock().unlock(); }
    }

    void setWithTtl(String key, String value, long ttlMillis) {
        lock.writeLock().lock();
        try {
            putLocked(key, new Entry(value, clock.nowMillis() + ttlMillis));   // ONE addition, once
            evictIfNeededLocked();
        } finally { lock.writeLock().unlock(); }
    }

    boolean delete(String key) {
        long now = clock.nowMillis();
        lock.writeLock().lock();
        try {
            Entry e = map.get(key);
            removeLocked(key);
            return e != null && !e.isExpired(now);       // do not claim to have deleted a ghost
        } finally { lock.writeLock().unlock(); }
    }

    boolean expire(String key, long ttlMillis) {
        lock.writeLock().lock();
        try {
            long now = clock.nowMillis();
            Entry e = map.get(key);
            if (e == null || e.isExpired(now)) return false;
            putLocked(key, new Entry(e.value(), now + ttlMillis));
            return true;
        } finally { lock.writeLock().unlock(); }
    }

    boolean persist(String key) {
        lock.writeLock().lock();
        try {
            long now = clock.nowMillis();
            Entry e = map.get(key);
            if (e == null || e.isExpired(now) || e.neverExpires()) return false;
            putLocked(key, new Entry(e.value(), Entry.NEVER));
            return true;
        } finally { lock.writeLock().unlock(); }
    }

    /* -------------------------------------------------- atomic read-modify-write */

    /** Two clients doing get-then-set would lose an increment. Doing it inside the store cannot. */
    long incr(String key, long delta) {
        lock.writeLock().lock();
        try {
            long now = clock.nowMillis();
            Entry e = map.get(key);
            boolean live = e != null && !e.isExpired(now);
            long current = live ? Long.parseLong(e.value()) : 0;      // missing key counts as 0
            long next = current + delta;
            putLocked(key, new Entry(Long.toString(next), live ? e.expiresAtMillis() : Entry.NEVER));
            return next;
        } finally { lock.writeLock().unlock(); }
    }

    /** Optimistic concurrency: write only if nobody changed it since you read it. */
    boolean compareAndSet(String key, String expected, String next) {
        lock.writeLock().lock();
        try {
            long now = clock.nowMillis();
            Entry e = map.get(key);
            if (e == null || e.isExpired(now) || !e.value().equals(expected)) return false;
            putLocked(key, new Entry(next, e.expiresAtMillis()));
            return true;
        } finally { lock.writeLock().unlock(); }
    }

    /* ---------------------------------------------------------- active expiry */

    /**
     * One pass of the sweeper. Draw 20 random keys that have a TTL, delete the dead ones,
     * and go again only while the sample stays dense. Bounded work, then it returns.
     */
    int sweepOnce() {
        int totalRemoved = 0;
        for (int round = 0; round < MAX_SWEEP_ROUNDS; round++) {
            List<String> sample;
            lock.readLock().lock();
            try { sample = sampleVolatileLocked(); } finally { lock.readLock().unlock(); }
            if (sample.isEmpty()) break;

            int expiredInSample = 0;
            lock.writeLock().lock();
            try {
                long now = clock.nowMillis();
                for (String key : sample) {
                    Entry e = map.get(key);
                    if (e != null && e.isExpired(now)) { removeLocked(key); expiredInSample++; }
                }
            } finally { lock.writeLock().unlock(); }

            totalRemoved += expiredInSample;
            activeExpired.add(expiredInSample);

            // more than 25% dead? go again immediately. Integer form of expired/n > 0.25
            if (expiredInSample * 4 <= sample.size()) break;
        }
        return totalRemoved;
    }

    private List<String> sampleVolatileLocked() {
        int n = Math.min(SAMPLE_SIZE, volatileKeys.size());
        List<String> out = new ArrayList<>(n);
        for (int i = 0; i < n; i++) out.add(volatileKeys.get(random.nextInt(volatileKeys.size())));
        return out;
    }

    /* ------------------------------------------------------------- eviction */

    /**
     * Dying of SPACE, not of time. Policy here is volatile-ttl: among keys that already
     * have a deadline, drop the one expiring soonest. The full machinery of allkeys-lru
     * is a problem of its own — this store only needs a pickVictim() seam.
     */
    private void evictIfNeededLocked() {
        while (map.size() > maxEntries) {
            String victim = null;
            long soonest = Long.MAX_VALUE;
            for (String key : volatileKeys) {
                Entry e = map.get(key);
                if (e != null && e.expiresAtMillis() < soonest) { soonest = e.expiresAtMillis(); victim = key; }
            }
            if (victim == null) victim = map.keySet().iterator().next();   // nothing volatile: random
            removeLocked(victim);
            evicted.increment();
        }
    }

    /* -------------------------------------------------------------- internals */

    private void putLocked(String key, Entry e) {
        map.put(key, e);
        if (e.neverExpires()) unmarkVolatile(key); else markVolatile(key);
    }

    private void removeLocked(String key) {
        map.remove(key);
        unmarkVolatile(key);
    }

    /** List + index, so a random pick and a removal are both O(1). */
    private void markVolatile(String key) {
        if (volatileIndex.containsKey(key)) return;
        volatileIndex.put(key, volatileKeys.size());
        volatileKeys.add(key);
    }

    private void unmarkVolatile(String key) {
        Integer i = volatileIndex.remove(key);
        if (i == null) return;
        int last = volatileKeys.size() - 1;
        if (i != last) {
            String moved = volatileKeys.get(last);
            volatileKeys.set(i, moved);
            volatileIndex.put(moved, i);
        }
        volatileKeys.remove(last);
    }

    /** One-star glob, which is all keys(pattern) ever really needs in an interview. */
    private static boolean glob(String pattern, String key) {
        int star = pattern.indexOf('*');
        if (star < 0) return pattern.equals(key);
        String head = pattern.substring(0, star), tail = pattern.substring(star + 1);
        return key.length() >= head.length() + tail.length()
            && key.startsWith(head) && key.endsWith(tail);
    }

    String stats() {
        return "keys=" + size() + " live=" + liveSize()
             + " hits=" + hits.sum() + " misses=" + misses.sum()
             + " expired-lazily=" + lazyExpired.sum()
             + " expired-actively=" + activeExpired.sum()
             + " evicted=" + evicted.sum();
    }
}

public class Main {
    public static void main(String[] args) {
        FakeClock clock = new FakeClock(0);
        KeyValueStore store = new KeyValueStore(clock, 100);

        store.setWithTtl("session:1", "abc", 5_000);
        store.set("user:7", "nina");                       // no deadline — permanent
        System.out.println("t=0     get(session:1) = " + store.get("session:1").orElse("(nil)"));
        System.out.println("t=0     size()=" + store.size() + " live=" + store.liveSize());

        clock.advance(6_000);                              // one second past the deadline
        System.out.println();
        System.out.println("t=6000  the key expired 1s ago. NOTHING RAN.");
        System.out.println("t=6000  size()=" + store.size() + " live=" + store.liveSize() + "   <- still counted");
        System.out.println("t=6000  get(session:1) = " + store.get("session:1").orElse("(nil)"));
        System.out.println("t=6000  size()=" + store.size() + "   <- the GET removed it (lazy expiry)");

        store.setWithTtl("report:2019", "...", 1_000);     // a key nobody will ever read
        clock.advance(60_000);
        System.out.println();
        System.out.println("t=66000 cold key expired 59s ago, never read: size()=" + store.size());
        System.out.println("t=66000 sweeper reclaimed " + store.sweepOnce() + " key(s), size()=" + store.size());

        store.setWithTtl("cart:42", "[a]", 30_000);
        System.out.println();
        System.out.println("ttl(cart:42) = " + store.ttl("cart:42"));
        System.out.println("ttl(user:7)  = " + store.ttl("user:7"));
        System.out.println("ttl(nope)    = " + store.ttl("nope"));

        store.set("cart:42", "[a,b]");                     // plain set — the TTL is gone
        System.out.println("after a plain set,   ttl(cart:42) = " + store.ttl("cart:42"));
        store.expire("cart:42", 30_000);
        store.set("cart:42", "[a,b,c]", true);             // keepTtl
        System.out.println("with keepTtl=true,   ttl(cart:42) = " + store.ttl("cart:42"));

        System.out.println();
        System.out.println("incr(visits) x3      = " + store.incr("visits", 1) + ", "
                                                     + store.incr("visits", 1) + ", "
                                                     + store.incr("visits", 1));
        System.out.println("cas(visits, 3 -> 10) = " + store.compareAndSet("visits", "3", "10"));
        System.out.println("cas(visits, 3 -> 99) = " + store.compareAndSet("visits", "3", "99") + "  (stale)");

        System.out.println();
        System.out.println("keys(\"c*\")          = " + store.keys("c*"));
        System.out.println(store.stats());
    }
}

/* -------------------------- expected output --------------------------
t=0     get(session:1) = abc
t=0     size()=2 live=2

t=6000  the key expired 1s ago. NOTHING RAN.
t=6000  size()=2 live=1   <- still counted
t=6000  get(session:1) = (nil)
t=6000  size()=1   <- the GET removed it (lazy expiry)

t=66000 cold key expired 59s ago, never read: size()=2
t=66000 sweeper reclaimed 1 key(s), size()=1

ttl(cart:42) = 30000ms left
ttl(user:7)  = no expiry
ttl(nope)    = no such key
after a plain set,   ttl(cart:42) = no expiry
with keepTtl=true,   ttl(cart:42) = 30000ms left

incr(visits) x3      = 1, 2, 3
cas(visits, 3 -> 10) = true
cas(visits, 3 -> 99) = false  (stale)

keys("c*")          = [cart:42]
keys=3 live=3 hits=1 misses=1 expired-lazily=1 expired-actively=1 evicted=0
--------------------------------------------------------------------- */