commonMain.dev.kord.cache.api.meta.MetricsCache.kt Maven / Gradle / Ivy

Go to download
package dev.kord.cache.api.meta

import dev.kord.cache.api.DataCache
import dev.kord.cache.api.DataEntryCache
import dev.kord.cache.api.QueryBuilder
import dev.kord.cache.api.Query
import kotlinx.atomicfu.atomic
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.onEach
import kotlin.reflect.KType

/**
 * Wraps this cache in a [MetricsCache], enabling logging of metrics.
 */
fun DataCache.withMetrics(): MetricsCache = MetricsCache(this)

/**
 * A [DataCache] that logs usage [statistics] via its [logger].
 */
class MetricsCache(
        private val delegate: DataCache,
        private val logger: StatisticsLogger = StatisticsLogger()
) : DataCache by delegate {

    val statistics: CacheStatistics get() = logger.metaData

    override fun  getEntry(type: KType): DataEntryCache? {
        val cache = getEntry(type) ?: return null
        return MetricsEntryCache(cache, logger.getForType(type))
    }

}

private class MetricsEntryCache(
        private val delegate: DataEntryCache,
        private val logger: TypeStatisticsLogger
) : DataEntryCache by delegate {

    override fun query(): QueryBuilder {
        val builder = delegate.query()
        return MetricsQueryBuilder(builder, logger)
    }

}

private class MetricsQueryBuilder(
        private val delegate: QueryBuilder,
        private val logger: TypeStatisticsLogger
) : QueryBuilder by delegate {

    override fun build(): Query {
        val query = delegate.build()
        return MetricsQuery(query, logger)
    }

}

private class MetricsQuery(
        private val delegate: Query,
        private val logger: TypeStatisticsLogger
) : Query by delegate {

    inner class FlowConverter {
        val first = atomic(false)

        fun convert() = delegate.asFlow().onEach {
            if (first.compareAndSet(expect = false, update = true)) logger.logHit()
        }
    }

    override fun asFlow(): Flow {
        logger.logQuery()
        return FlowConverter().convert()
    }

    override suspend fun single(): V {
        logger.logQuery()
        val result = runCatching { delegate.single() }

        result.onSuccess { logger.logQuery() }
        return result.getOrElse { throw it }
    }

    override suspend fun singleOrNull(): V? {
        logger.logQuery()
        return delegate.singleOrNull()?.also {
            logger.logHit()
        }
    }

    override suspend fun toCollection(): Collection {
        logger.logQuery()
        val collection = super.toCollection()

        if (collection.isNotEmpty()) logger.logHit()
        return collection
    }

}