LoadingCache.kt

import java.time.Instant
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicLong

class LoadingCache<K : Any, V : Any>(
    private val maxReuses: Int, // how many times it's possible to SEQUENTIALLY reuse previously cached value in case of loading failure
    private val loader: (key: K, oldValue: V?) -> V,
    private val ttlResolver: (state: State<V>, timestamp: Instant) -> Instant
) {
    private val cache = ConcurrentHashMap<K, Slot<V>>()
    private val lastStatsLoggedAtMillis: AtomicLong = AtomicLong()

    operator fun get(key: K): V? {
        val slot: Slot<V> = cache.compute(key, ::computeSlot)!!
        slot.hitsCount.incrementAndGet()
        logStats()

        return slot.state.data?.value
    }

    private fun logStats() {
        val timestampMs = System.currentTimeMillis()
        val lastTimestampMs = lastStatsLoggedAtMillis.get()

        if (timestampMs - lastTimestampMs >= STATS_LOG_INTERVAL_MILLIS) {
            if (lastStatsLoggedAtMillis.compareAndSet(lastTimestampMs, timestampMs)) {
                log.info { "Total keys count: ${cache.size}" }
            }
        }
    }

    private fun computeSlot(key: K, oldSlot: Slot<V>?): Slot<V> {
        val timestamp = Instant.now()
        return if (oldSlot != null && timestamp <= oldSlot.actualTill) {
            oldSlot
        } else {
            val state = try {
                State.Loaded(loader(key, oldSlot?.state?.data?.value), timestamp)
            } catch (e: Exception) {
                if (oldSlot == null || oldSlot.state is State.Failed) {
                    log.error(e) { "Failed to load value by key [$key]; value is undefined" }
                    STATE_FAILED.cast<V>()
                } else {
                    val reusesCount: Int = when (oldSlot.state) {
                        is State.Loaded -> 0
                        is State.Reused -> oldSlot.state.reusesCount
                        is State.Failed -> throw RuntimeException("Unreachable condition, Kotlin")
                    }
                    if (reusesCount >= maxReuses) {
                        log.error(e) {
                            "Failed to load value by key [$key]; previously cached value can't be reused due" +
                                " to the max reuses limit ($maxReuses); now value will become undefined"
                        }
                        STATE_FAILED.cast<V>()
                    } else {
                        log.error(e) { "Failed to load value by key [$key]; previously cached value will be reused" }
                        with(oldSlot.state.data!!) {
                            State.Reused(value, loadedAt, reusesCount + 1)
                        }
                    }
                }
            }

            val actualTill = ttlResolver(state, timestamp)
            if (actualTill < timestamp) {
                throw RuntimeException("Resolved expiration timestamp [$actualTill]" +
                    " can't be less than current timestamp [$timestamp]")
            }
            Slot(state, actualTill).applyStats(oldSlot)
        }
    }

    private fun Slot<V>.applyStats(source: Slot<V>?): Slot<V> {
        if (source != null) {
            hitsCount.addAndGet(source.hitsCount.get())
        }
        return this
    }

    private fun <V : Any> State.Failed<*>.cast() = this as State.Failed<V>

    interface Data<V : Any> {
        val value: V
        val loadedAt: Instant
    }

    sealed class State<V : Any> {
        val data: Data<V>?
            get() = if (this is Data<*>) this as Data<V> else null

        data class Loaded<V : Any>(
            override val value: V,
            override val loadedAt: Instant
        ) : State<V>(), Data<V>

        data class Reused<V : Any>(
            override val value: V,
            override val loadedAt: Instant,
            val reusesCount: Int
        ) : State<V>(), Data<V>

        class Failed<V : Any> : State<V>()
    }

    private data class Slot<V : Any>(
        val state: State<V>,
        val actualTill: Instant,
        val hitsCount: AtomicLong = AtomicLong()
    )

    private companion object : CompanionDevKit() {
        private val STATS_LOG_INTERVAL_MILLIS = TimeUnit.MINUTES.toMillis(1)
        private val STATE_FAILED = State.Failed<Any>()
    }
}