io.github.freya022.botcommands.internal.commands.ratelimit.RateLimitAnnotations.kt Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of BotCommands Show documentation

A Kotlin-first (and Java) framework that makes creating Discord bots a piece of cake, using the JDA library.

There is a newer version: 3.0.0-alpha.18

package io.github.freya022.botcommands.internal.commands.ratelimit

import io.github.freya022.botcommands.api.commands.annotations.*
import io.github.freya022.botcommands.api.commands.builder.CommandBuilder
import io.github.freya022.botcommands.api.commands.ratelimit.AnnotatedRateLimiterFactory
import io.github.freya022.botcommands.api.commands.ratelimit.RateLimiter
import io.github.freya022.botcommands.api.commands.ratelimit.bucket.Buckets
import io.github.freya022.botcommands.api.commands.ratelimit.bucket.toSupplier
import io.github.freya022.botcommands.api.core.service.getService
import io.github.freya022.botcommands.api.core.utils.findAnnotationRecursive
import io.github.freya022.botcommands.internal.utils.ReflectionUtils.declaringClass
import io.github.freya022.botcommands.internal.utils.annotationRef
import io.github.freya022.botcommands.internal.utils.requireAt
import java.time.Duration
import kotlin.reflect.KFunction
import io.github.bucket4j.Bandwidth as BucketBandwidth

private fun Bandwidth.toRealBandwidth(): BucketBandwidth =
    BucketBandwidth.builder()
        .capacity(capacity)
        .let {
            val duration = Duration.of(refill.period, refill.periodUnit)
            when (refill.type) {
                RefillType.GREEDY -> it.refillGreedy(refill.tokens, duration)
                RefillType.INTERVAL -> it.refillIntervally(refill.tokens, duration)
            }
        }
        .build()

internal fun CommandBuilder.readRateLimit(func: KFunction<*>): RateLimiter? {
    val rateLimitAnnotation = func.findAnnotationRecursive() ?: func.declaringClass.findAnnotationRecursive()
    val cooldownAnnotation = func.findAnnotationRecursive() ?: func.declaringClass.findAnnotationRecursive()
    requireAt(cooldownAnnotation == null || rateLimitAnnotation == null, func) {
        "Cannot use both ${annotationRef()} and ${annotationRef()}"
    }

    return if (rateLimitAnnotation != null) {
        val bucketConfigurationSupplier = Buckets.custom(rateLimitAnnotation.bandwidths.map { it.toRealBandwidth() }).toSupplier()
        val annotatedRateLimiterFactory = context.getService()
        annotatedRateLimiterFactory.create(rateLimitAnnotation.scope, bucketConfigurationSupplier, rateLimitAnnotation.deleteOnRefill)
    } else if (cooldownAnnotation != null) {
        val bucketConfigurationSupplier = Buckets.ofCooldown(Duration.of(cooldownAnnotation.cooldown, cooldownAnnotation.unit)).toSupplier()
        val annotatedRateLimiterFactory = context.getService()
        annotatedRateLimiterFactory.create(cooldownAnnotation.scope, bucketConfigurationSupplier, cooldownAnnotation.deleteOnRefill)
    } else {
        null
    }
}