command cooldown

src/main/kotlin/me/devoxin/flight/api/CommandClient.kt

@@ -2,10 +2,13 @@ package me.devoxin.flight.api
 
 import kotlinx.coroutines.GlobalScope
 import kotlinx.coroutines.launch
+import me.devoxin.flight.api.annotations.Cooldown
+import me.devoxin.flight.api.entities.BucketType
 import me.devoxin.flight.internal.arguments.ArgParser
 import me.devoxin.flight.api.exceptions.BadArgument
 import me.devoxin.flight.internal.entities.WaitingEvent
 import me.devoxin.flight.api.entities.Cog
+import me.devoxin.flight.api.entities.CooldownProvider
 import me.devoxin.flight.api.hooks.CommandEventAdapter
 import me.devoxin.flight.api.entities.PrefixProvider
 import me.devoxin.flight.internal.entities.CommandRegistry
@@ -26,6 +29,7 @@ import kotlin.reflect.KParameter
 class CommandClient(
     parsers: HashMap<Class<*>, Parser<*>>,
     private val prefixProvider: PrefixProvider,
+    private val cooldownProvider: CooldownProvider,
     private val ignoreBots: Boolean,
     private val eventListeners: List<CommandEventAdapter>,
     customOwnerIds: MutableSet<Long>?
@@ -98,6 +102,22 @@ class CommandClient(
                 ?: return
 
         val ctx = Context(this, event, trigger)
+
+        if (cmd.cooldown != null) {
+            val entityId = when (cmd.cooldown.bucket) {
+                BucketType.USER -> ctx.author.idLong
+                BucketType.GUILD -> ctx.guild?.idLong
+                BucketType.GLOBAL -> -1
+            }
+
+            if (entityId != null) {
+                if (cooldownProvider.isOnCooldown(entityId, cmd.cooldown.bucket, cmd)) {
+                    val time = cooldownProvider.getCooldownTime(entityId, cmd.cooldown.bucket, cmd)
+                    return eventListeners.forEach { it.onCommandCooldown(ctx, cmd, time) }
+                }
+            }
+        }
+
         val props = cmd.properties
 
         if (props.developerOnly && !ownerIds.contains(event.author.idLong)) {
@@ -159,6 +179,19 @@ class CommandClient(
             eventListeners.forEach { it.onCommandPostInvoke(ctx, cmd, !success) }
         }
 
+        if (cmd.cooldown != null && cmd.cooldown.duration > 0) {
+            val entityId = when (cmd.cooldown.bucket) {
+                BucketType.USER -> ctx.author.idLong
+                BucketType.GUILD -> ctx.guild?.idLong
+                BucketType.GLOBAL -> -1
+            }
+
+            if (entityId != null) {
+                val time = cmd.cooldown.timeUnit.toMillis(cmd.cooldown.duration)
+                cooldownProvider.setCooldown(entityId, cmd.cooldown.bucket, time, cmd)
+            }
+        }
+
         if (cmd.async) {
             GlobalScope.launch {
                 cmd.executeAsync(ctx, arguments, complete = cb)

src/main/kotlin/me/devoxin/flight/api/CommandClientBuilder.kt

@@ -1,12 +1,9 @@
 package me.devoxin.flight.api
 
+import me.devoxin.flight.api.entities.*
+import me.devoxin.flight.api.entities.Invite
 import me.devoxin.flight.internal.arguments.types.Snowflake
-import me.devoxin.flight.api.entities.DefaultHelpCommand
-import me.devoxin.flight.api.entities.DefaultPrefixProvider
 import me.devoxin.flight.api.hooks.CommandEventAdapter
-import me.devoxin.flight.api.entities.Emoji
-import me.devoxin.flight.api.entities.Invite
-import me.devoxin.flight.api.entities.PrefixProvider
 import me.devoxin.flight.internal.parsers.*
 import net.dv8tion.jda.api.entities.*
 import java.net.URL
@@ -19,6 +16,7 @@ class CommandClientBuilder {
     private var helpCommandConfig: DefaultHelpCommandConfig = DefaultHelpCommandConfig()
     private var ignoreBots: Boolean = true
     private var prefixProvider: PrefixProvider? = null
+    private var cooldownProvider: CooldownProvider? = null
     private var eventListeners: MutableList<CommandEventAdapter> = mutableListOf()
     private var ownerIds: MutableSet<Long>? = null
 
@@ -51,6 +49,14 @@ class CommandClientBuilder {
         return this
     }
 
+    /**
+     * Sets the provider used for cool-downs.
+     */
+    fun setCooldownProvider(provider: CooldownProvider): CommandClientBuilder {
+        this.cooldownProvider = provider
+        return this
+    }
+
     /**
      * Whether the bot will allow mentions to be used as a prefix.
      *
@@ -188,7 +194,8 @@ class CommandClientBuilder {
     @ExperimentalStdlibApi
     fun build(): CommandClient {
         val prefixProvider = this.prefixProvider ?: DefaultPrefixProvider(prefixes, allowMentionPrefix)
-        val commandClient = CommandClient(parsers, prefixProvider, ignoreBots, eventListeners.toList(), ownerIds)
+        val cooldownProvider = this.cooldownProvider ?: DefaultCooldownProvider()
+        val commandClient = CommandClient(parsers, prefixProvider, cooldownProvider, ignoreBots, eventListeners.toList(), ownerIds)
 
         if (helpCommandConfig.enabled) {
             commandClient.registerCommands(DefaultHelpCommand(helpCommandConfig.showParameterTypes))

src/main/kotlin/me/devoxin/flight/api/CommandFunction.kt

@@ -1,6 +1,7 @@
 package me.devoxin.flight.api
 
 import me.devoxin.flight.api.annotations.Command
+import me.devoxin.flight.api.annotations.Cooldown
 import me.devoxin.flight.internal.arguments.Argument
 import me.devoxin.flight.internal.entities.Jar
 import me.devoxin.flight.api.entities.Cog
@@ -14,6 +15,7 @@ class CommandFunction(
     val arguments: List<Argument>,
     val category: String,
     val properties: Command,
+    val cooldown: Cooldown?,
     val async: Boolean,
     val method: KFunction<*>,
     val cog: Cog,

src/main/kotlin/me/devoxin/flight/api/annotations/Cooldown.kt

@@ -0,0 +1,12 @@
+package me.devoxin.flight.api.annotations
+
+import me.devoxin.flight.api.entities.BucketType
+import java.util.concurrent.TimeUnit
+
+@Retention(AnnotationRetention.RUNTIME)
+@Target(AnnotationTarget.FUNCTION)
+annotation class Cooldown(
+    val duration: Long,
+    val timeUnit: TimeUnit = TimeUnit.MILLISECONDS,
+    val bucket: BucketType
+)

src/main/kotlin/me/devoxin/flight/api/entities/BucketType.kt

@@ -0,0 +1,7 @@
+package me.devoxin.flight.api.entities
+
+enum class BucketType {
+    USER,
+    GUILD,
+    GLOBAL
+}

src/main/kotlin/me/devoxin/flight/api/entities/CooldownProvider.kt

@@ -0,0 +1,66 @@
+package me.devoxin.flight.api.entities
+
+import me.devoxin.flight.api.CommandFunction
+
+interface CooldownProvider {
+
+    /**
+     * Checks whether the entity associated with the provided ID is on cool-down.
+     * When BucketType is `GUILD` and the command was invoked in a private context, this
+     * method won't be called.
+     *
+     * @param id
+     *        The ID of the entity. If the bucket type is USER, this will be a user ID.
+     *        If the bucket type is GUILD, this will be the guild id.
+     *        If the bucket type is GLOBAL, this will be -1.
+     *
+     * @param bucket
+     *        The type of bucket the cool-down belongs to.
+     *        For example, one bucket for each entity type; USER, GUILD, GLOBAL.
+     *        If this parameter is GUILD, theoretically you would do `bucket[type].get(id) != null`
+     *
+     * @param command
+     *        The command that was invoked.
+     *
+     * @returns True, if the entity associated with the ID is on cool-down and the command should
+     *          not be executed.
+     */
+    fun isOnCooldown(id: Long, bucket: BucketType, command: CommandFunction): Boolean
+
+    /**
+     * Gets the remaining time of the cooldown in milliseconds.
+     * This may either return 0L, or throw an exception if an entry isn't present, however
+     * this should not happen as `isOnCooldown` should be called prior to this.
+     *
+     * @param id
+     *        The ID of the entity. The ID could belong to a user or guild, or be -1 if the bucket is GLOBAL.
+     *
+     * @param bucket
+     *        The type of bucket to check the cool-down of.
+     *
+     * @param command
+     *        The command to get the cool-down time of.
+     */
+    fun getCooldownTime(id: Long, bucket: BucketType, command: CommandFunction): Long
+
+    /**
+     * Adds a cooldown for the given entity ID.
+     * It is up to you whether this passively, or actively removes expired cool-downs.
+     * When BucketType is `GUILD` and the command was invoked in a private context, this
+     * method won't be called.
+     *
+     * @param id
+     *        The ID of the entity, that the cool-down should be associated with.
+     *        This ID could belong to a user or guild. If bucket is GLOBAL, this will be -1.
+     *
+     * @param bucket
+     *        The type of bucket the cool-down belongs to.
+     *
+     * @param time
+     *        How long the cool-down should last for, in milliseconds.
+     *
+     * @param command
+     *        The command to set cool-down for.
+     */
+    fun setCooldown(id: Long, bucket: BucketType, time: Long, command: CommandFunction)
+}

src/main/kotlin/me/devoxin/flight/api/entities/DefaultCooldownProvider.kt

@@ -0,0 +1,62 @@
+package me.devoxin.flight.api.entities
+
+import me.devoxin.flight.api.CommandFunction
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.Executors
+import java.util.concurrent.TimeUnit
+import kotlin.math.abs
+
+class DefaultCooldownProvider : CooldownProvider {
+    private val buckets = ConcurrentHashMap<BucketType, Bucket>()
+
+    override fun isOnCooldown(id: Long, bucket: BucketType, command: CommandFunction): Boolean {
+        return buckets[bucket]?.isOnCooldown(id, command.name) ?: false
+    }
+
+    override fun getCooldownTime(id: Long, bucket: BucketType, command: CommandFunction): Long {
+        return buckets[bucket]?.getCooldownRemainingTime(id, command.name) ?: 0
+    }
+
+    override fun setCooldown(id: Long, bucket: BucketType, time: Long, command: CommandFunction) {
+        buckets.computeIfAbsent(bucket) { Bucket() }.setCooldown(id, time, command.name)
+    }
+
+    inner class Bucket {
+        private val sweeperThread = Executors.newSingleThreadScheduledExecutor()
+        private val cooldowns = ConcurrentHashMap<Long, MutableSet<Cooldown>>() // EntityID => [Commands...]
+
+        fun isOnCooldown(id: Long, commandName: String): Boolean {
+            return getCooldownRemainingTime(id, commandName) > 0
+        }
+
+        fun getCooldownRemainingTime(id: Long, commandName: String): Long {
+            val cd = cooldowns[id]?.firstOrNull { it.name == commandName }
+                ?: return 0
+
+            return abs(cd.expires - System.currentTimeMillis())
+        }
+
+        fun setCooldown(id: Long, time: Long, commandName: String) {
+            val cds = cooldowns.computeIfAbsent(id) { mutableSetOf() }
+            val cooldown = Cooldown(commandName, System.currentTimeMillis() + time)
+            cds.add(cooldown)
+
+            sweeperThread.schedule({ cds.remove(cooldown) }, time, TimeUnit.MILLISECONDS)
+        }
+    }
+
+    inner class Cooldown(val name: String, val expires: Long) {
+        override fun equals(other: Any?): Boolean {
+            if (this === other) return true
+            if (javaClass != other?.javaClass) return false
+
+            other as Cooldown
+
+            return name == other.name
+        }
+
+        override fun hashCode(): Int {
+            return 31 * name.hashCode()
+        }
+    }
+}

src/main/kotlin/me/devoxin/flight/api/hooks/CommandEventAdapter.kt

@@ -41,6 +41,18 @@ interface CommandEventAdapter {
      */
     fun onCommandError(ctx: Context, command: CommandFunction, error: Throwable)
 
+    /**
+     * Invoked when a command is executed while on cool-down.
+     *
+     * @param ctx
+     *        The command context.
+     * @param command
+     *        The command that encountered the cool-down.
+     * @param cooldown
+     *        The remaining time of the cool-down, in milliseconds.
+     */
+    fun onCommandCooldown(ctx: Context, command: CommandFunction, cooldown: Long)
+
     /**
      * Invoked when a user lacks permissions to execute a command
      */

src/main/kotlin/me/devoxin/flight/api/hooks/DefaultCommandEventAdapter.kt

@@ -7,15 +7,23 @@ import net.dv8tion.jda.api.Permission
 
 abstract class DefaultCommandEventAdapter : CommandEventAdapter {
 
-    override fun onBadArgument(ctx: Context, command: CommandFunction, error: BadArgument) {}
+    override fun onBadArgument(ctx: Context, command: CommandFunction, error: BadArgument) {
+        error.printStackTrace()
+    }
 
-    override fun onCommandError(ctx: Context, command: CommandFunction, error: Throwable) {}
+    override fun onCommandError(ctx: Context, command: CommandFunction, error: Throwable) {
+        error.printStackTrace()
+    }
 
     override fun onCommandPostInvoke(ctx: Context, command: CommandFunction, failed: Boolean) {}
 
     override fun onCommandPreInvoke(ctx: Context, command: CommandFunction) = true
 
-    override fun onParseError(ctx: Context, command: CommandFunction, error: Throwable) {}
+    override fun onParseError(ctx: Context, command: CommandFunction, error: Throwable) {
+        error.printStackTrace()
+    }
+
+    override fun onCommandCooldown(ctx: Context, command: CommandFunction, cooldown: Long) {}
 
     override fun onBotMissingPermissions(ctx: Context, command: CommandFunction, permissions: List<Permission>) {}
 

src/main/kotlin/me/devoxin/flight/internal/utils/Indexer.kt

@@ -3,6 +3,7 @@ package me.devoxin.flight.internal.utils
 import me.devoxin.flight.api.annotations.Command
 import me.devoxin.flight.api.CommandFunction
 import me.devoxin.flight.api.Context
+import me.devoxin.flight.api.annotations.Cooldown
 import me.devoxin.flight.internal.arguments.Argument
 import me.devoxin.flight.api.annotations.Greedy
 import me.devoxin.flight.api.annotations.Name
@@ -79,6 +80,7 @@ class Indexer {
         val category = cog.name()
         val name = meth.name.toLowerCase()
         val properties = meth.findAnnotation<Command>()!!
+        val cooldown = meth.findAnnotation<Cooldown>()
         val async = meth.isSuspend
         val ctxParam = meth.valueParameters.firstOrNull { it.type.classifier?.equals(Context::class) == true }
 
@@ -99,7 +101,7 @@ class Indexer {
             arguments.add(Argument(pName, type, greedy, optional, isNullable, p))
         }
 
-        return CommandFunction(name, arguments, category, properties, async, meth, cog, jar, ctxParam)
+        return CommandFunction(name, arguments, category, properties, cooldown, async, meth, cog, jar, ctxParam)
     }
 
     companion object {