BoundedLocalCache.java

package io.github.czhaodev.tinylfu.internal;

import io.github.czhaodev.tinylfu.Cache;
import io.github.czhaodev.tinylfu.Ticker;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.Function;

/**
 * The concurrent engine: keeps {@link #getIfPresent} effectively lock-free
 * by never touching the eviction policy or timer wheel inline. Reads are
 * recorded into a {@link StripedReadBuffer} (droppable, advisory-only);
 * writes update the backing map immediately (visible right away) and enqueue
 * a bookkeeping event onto a {@link WriteBuffer} (never dropped). A single
 * {@link DrainStatus} state machine ensures at most one thread drains at a
 * time. <strong>The map is always right, everything else is eventually
 * right</strong> -- {@code get}/{@code put}/{@code invalidate} are visible
 * the instant they return, regardless of whether any bookkeeping has run.
 *
 * <p>Unlike the Rust sibling's {@code DashMap}-backed design, this class
 * never needs to release a map guard before triggering a drain:
 * {@link ConcurrentHashMap#compute} holds its per-bin lock as an ordinary
 * (reentrant) Java monitor for the callback's entire duration, so a nested,
 * same-thread call back into the map -- even one that happens to hash to
 * the same bin -- cannot deadlock the way a non-reentrant lock would. This
 * also means an {@code UPDATE} event can never be drained before its own
 * node's {@code ADD} event: a second writer's {@code compute()} call for the
 * same key cannot even begin until the first writer's callback (which
 * already enqueued the {@code ADD}) has returned. The staleness checks
 * mirroring the Rust source's races are kept anyway, as defense in depth.
 */
public final class BoundedLocalCache<K, V> implements Cache<K, V> {

    private final Ticker ticker;
    private final long expireAfterWriteNanos; // -1 if not configured
    private final long expireAfterAccessNanos; // -1 if not configured

    private final ConcurrentHashMap<K, Node<K, V>> data = new ConcurrentHashMap<>();
    private final StripedReadBuffer<K, V> readBuffer = new StripedReadBuffer<>();
    private final WriteBuffer<K, V> writeBuffer = new WriteBuffer<>();
    private final DrainStatus drainStatus = new DrainStatus();

    /** Guards {@link #policy} and {@link #timerWheel}; only ever acquired via tryLock, except in {@link #cleanUp}. */
    private final ReentrantLock evictionLock = new ReentrantLock();
    private final WindowTinyLfuPolicy<K, V> policy;
    private final TimerWheel<K, V> timerWheel;

    public BoundedLocalCache(long maximumSize, Ticker ticker, long expireAfterWriteNanos, long expireAfterAccessNanos) {
        this.ticker = ticker;
        this.expireAfterWriteNanos = expireAfterWriteNanos;
        this.expireAfterAccessNanos = expireAfterAccessNanos;
        this.policy = new WindowTinyLfuPolicy<>(maximumSize);
        this.timerWheel = new TimerWheel<>(ticker.read());
    }

    @Override
    public Optional<V> getIfPresent(K key) {
        Node<K, V> node = data.get(key);
        if (node == null || isExpired(node)) {
            return Optional.empty();
        }
        afterRead(node);
        return Optional.of(node.value);
    }

    @Override
    public Optional<V> getWith(K key, Function<? super K, ? extends V> mappingFunction) {
        Node<K, V> existing = data.get(key);
        if (existing != null && !isExpired(existing)) {
            afterRead(existing);
            return Optional.of(existing.value);
        }

        long now = ticker.read();
        AtomicReference<V> result = new AtomicReference<>();
        AtomicReference<Node<K, V>> removedNode = new AtomicReference<>();
        AtomicReference<Node<K, V>> addedNode = new AtomicReference<>();
        data.compute(key, (k, current) -> {
            if (current != null && !isExpired(current)) {
                result.set(current.value);
                return current;
            }
            if (current != null) {
                // An expired entry occupied this key; it's discarded either way. Never
                // reused, since it may still be referenced by an in-flight buffered
                // event or scheduled in the timer wheel with a stale deadline.
                removedNode.set(current);
            }
            V computed = mappingFunction.apply(key);
            if (computed == null) {
                return null;
            }
            Node<K, V> fresh = new Node<>(key, computed);
            applyWriteExpiration(fresh, now);
            addedNode.set(fresh);
            result.set(computed);
            return fresh;
        });

        // afterRemove/afterAdd (and the synchronous drain they may trigger) must run
        // only once compute() has returned and released this key's bin lock: compute()
        // reinstalls its return value unconditionally when the callback exits, so a
        // nested drain that touches this same key while still inside the callback
        // would race against that reinstallation.
        if (removedNode.get() != null) {
            afterRemove(removedNode.get());
        }
        if (addedNode.get() != null) {
            afterAdd(addedNode.get());
        }

        return Optional.ofNullable(result.get());
    }

    @Override
    public void put(K key, V value) {
        rawPut(key, value);
    }

    @Override
    public void invalidate(K key) {
        rawInvalidate(key);
    }

    @Override
    public void invalidateAll() {
        data.clear();
        writeBuffer.add(WriteEvent.clear());
        scheduleDrain();
    }

    @Override
    public void invalidateEntries(Iterable<? extends K> keys) {
        for (K key : keys) {
            rawInvalidate(key);
        }
    }

    @Override
    public long estimatedSize() {
        return data.mappingCount();
    }

    @Override
    public Map<K, V> asMapSnapshot() {
        Map<K, V> snapshot = new HashMap<>();
        for (Map.Entry<K, Node<K, V>> entry : data.entrySet()) {
            Node<K, V> node = entry.getValue();
            if (!isExpired(node)) {
                snapshot.put(entry.getKey(), node.value);
            }
        }
        return snapshot;
    }

    @Override
    public void cleanUp() {
        evictionLock.lock();
        try {
            drainStatus.set(DrainStatus.State.PROCESSING);
            performDrain();
        } finally {
            drainStatus.set(DrainStatus.State.IDLE);
            evictionLock.unlock();
        }
    }

    private V rawPut(K key, V value) {
        long now = ticker.read();
        AtomicReference<V> previous = new AtomicReference<>();
        AtomicReference<Node<K, V>> addedNode = new AtomicReference<>();
        AtomicReference<Node<K, V>> updatedNode = new AtomicReference<>();
        data.compute(key, (k, existing) -> {
            if (existing != null) {
                previous.set(existing.value);
                existing.value = value;
                applyWriteExpiration(existing, now);
                updatedNode.set(existing);
                return existing;
            }
            Node<K, V> node = new Node<>(key, value);
            applyWriteExpiration(node, now);
            addedNode.set(node);
            return node;
        });
        // See getWith's comment: afterAdd/afterUpdate must run only after compute()
        // has returned, so a nested drain never races against compute()'s own
        // (unconditional, on-return) installation of this key's mapping.
        if (addedNode.get() != null) {
            afterAdd(addedNode.get());
        } else if (updatedNode.get() != null) {
            afterUpdate(updatedNode.get());
        }
        return previous.get();
    }

    private V rawInvalidate(K key) {
        Node<K, V> removed = data.remove(key);
        if (removed == null) {
            return null;
        }
        afterRemove(removed);
        return removed.value;
    }

    private boolean hasExpiration() {
        return expireAfterWriteNanos >= 0 || expireAfterAccessNanos >= 0;
    }

    private boolean isExpired(Node<K, V> node) {
        return hasExpiration() && node.expirationTimeNanos <= ticker.read();
    }

    private void applyWriteExpiration(Node<K, V> node, long now) {
        if (!hasExpiration()) {
            return;
        }
        if (expireAfterWriteNanos >= 0) {
            node.writeExpirationTimeNanos = now + expireAfterWriteNanos;
        }
        node.expirationTimeNanos = effectiveDeadline(node, now);
    }

    private long effectiveDeadline(Node<K, V> node, long now) {
        long deadline = Long.MAX_VALUE;
        if (expireAfterWriteNanos >= 0) {
            deadline = Math.min(deadline, node.writeExpirationTimeNanos);
        }
        if (expireAfterAccessNanos >= 0) {
            deadline = Math.min(deadline, now + expireAfterAccessNanos);
        }
        return deadline;
    }

    private void afterRead(Node<K, V> node) {
        if (readBuffer.recordRead(node)) {
            scheduleDrain();
        }
    }

    private void afterAdd(Node<K, V> node) {
        writeBuffer.add(WriteEvent.add(node));
        scheduleDrain();
    }

    private void afterUpdate(Node<K, V> node) {
        writeBuffer.add(WriteEvent.update(node));
        scheduleDrain();
    }

    private void afterRemove(Node<K, V> node) {
        writeBuffer.add(WriteEvent.remove(node));
        scheduleDrain();
    }

    private void scheduleDrain() {
        if (drainStatus.compareAndSet(DrainStatus.State.IDLE, DrainStatus.State.REQUIRED)) {
            tryDrain();
        }
    }

    private void tryDrain() {
        if (!drainStatus.compareAndSet(DrainStatus.State.REQUIRED, DrainStatus.State.PROCESSING)) {
            return;
        }
        if (!evictionLock.tryLock()) {
            drainStatus.compareAndSet(DrainStatus.State.PROCESSING, DrainStatus.State.REQUIRED);
            return;
        }
        try {
            performDrain();
        } finally {
            drainStatus.set(DrainStatus.State.IDLE);
            evictionLock.unlock();
        }
    }

    /** Must run only while holding {@link #evictionLock}. */
    private void performDrain() {
        // Writes drain before reads: a buffered read can only ever reference a node
        // whose ADD event is already enqueued, so this ordering guarantees every
        // node a buffered read names has already been linked by the time reads
        // replay.
        WriteEvent<K, V> event;
        while ((event = writeBuffer.poll()) != null) {
            applyWriteEvent(event);
        }

        long now = ticker.read();

        readBuffer.drainTo(node -> {
            policy.recordAccess(node);
            if (node.queueType != null && expireAfterAccessNanos >= 0) {
                long newDeadline = effectiveDeadline(node, now);
                if (newDeadline != node.expirationTimeNanos) {
                    timerWheel.deschedule(node);
                    node.expirationTimeNanos = newDeadline;
                    timerWheel.schedule(node, newDeadline);
                }
            }
        });

        if (hasExpiration()) {
            List<Node<K, V>> expired = new ArrayList<>();
            timerWheel.advance(now, expired);
            for (Node<K, V> node : expired) {
                if (data.remove(node.key, node)) {
                    policy.remove(node);
                }
            }
        }

        for (Node<K, V> evicted : policy.evict()) {
            if (hasExpiration()) {
                timerWheel.deschedule(evicted);
            }
            data.remove(evicted.key, evicted);
        }
    }

    private void applyWriteEvent(WriteEvent<K, V> event) {
        Node<K, V> node = event.node;
        switch (event.type) {
            case ADD -> {
                policy.recordAdd(node);
                if (hasExpiration()) {
                    timerWheel.schedule(node, node.expirationTimeNanos);
                }
            }
            case UPDATE -> {
                if (node.queueType == null) {
                    return; // this node's own ADD hasn't drained yet -- safe no-op.
                }
                policy.recordAccess(node);
                if (hasExpiration()) {
                    timerWheel.deschedule(node);
                    timerWheel.schedule(node, node.expirationTimeNanos);
                }
            }
            case REMOVE -> {
                if (node.queueType != null) {
                    policy.remove(node);
                }
                if (hasExpiration()) {
                    timerWheel.deschedule(node);
                }
            }
            case CLEAR -> {
                policy.clear();
                if (hasExpiration()) {
                    timerWheel.clear();
                }
            }
        }
    }
}