Command timeouts and context coroutines

Author: mrkirby153

src/main/kotlin/com/mrkirby153/botcore/command/slashcommand/dsl/Context.kt

@@ -97,7 +97,8 @@ class SlashContext(
  *
  * @param event The [UserContextInteractionEvent] that this context wraps
  */
-class UserContext(event: UserContextInteractionEvent) : UserContextInteraction by event {
+class UserContext(event: UserContextInteractionEvent, scope: CoroutineScope) :
+    UserContextInteraction by event, CoroutineScope by scope {
 
 }
 
@@ -106,10 +107,10 @@ class UserContext(event: UserContextInteractionEvent) : UserContextInteraction b
  *
  * @param event The [MessageContextInteractionEvent] that this context wraps
  */
-class MessageContext(private val event: MessageContextInteractionEvent) :
-    MessageContextInteraction by event {
+class MessageContext(private val event: MessageContextInteractionEvent, scope: CoroutineScope) :
+    MessageContextInteraction by event, CoroutineScope by scope {
     /**
-     * The message that the interaction is being ran from
+     * The message that the interaction is being run from
      */
     val message
         get() = event.target

src/main/kotlin/com/mrkirby153/botcore/command/slashcommand/dsl/ContextCommands.kt

@@ -2,6 +2,7 @@ package com.mrkirby153.botcore.command.slashcommand.dsl
 
 import com.mrkirby153.botcore.utils.PrerequisiteCheck
 import net.dv8tion.jda.api.interactions.commands.context.ContextInteraction
+import java.util.concurrent.TimeUnit
 
 /**
  * Top level class for context menu interactions
@@ -16,10 +17,12 @@ open class ContextCommand<Event : ContextInteraction<*>>(
      * The name of the command
      */
 
-    private lateinit var commandAction: (Event) -> Unit
+    private lateinit var commandAction: suspend (Event) -> Unit
 
     private val checks = mutableListOf<PrerequisiteCheck<Event>.() -> Unit>()
 
+    internal var timeout: Long = 30_000
+
     /**
      * If this command should be enabled by default
      */
@@ -28,7 +31,7 @@ open class ContextCommand<Event : ContextInteraction<*>>(
     /**
      * The function ran when this interaction is invoked
      */
-    fun action(action: (Event) -> Unit) {
+    fun action(action: suspend (Event) -> Unit) {
         this.commandAction = action
     }
 
@@ -42,7 +45,7 @@ open class ContextCommand<Event : ContextInteraction<*>>(
     /**
      * Executes this command with the provided [event]
      */
-    fun execute(event: Event) {
+    suspend fun execute(event: Event) {
         val checkCtx = PrerequisiteCheck(event)
         checks.forEach {
             it(checkCtx)
@@ -59,6 +62,13 @@ open class ContextCommand<Event : ContextInteraction<*>>(
                 .setEphemeral(true).queue()
         }
     }
+
+    /**
+     * Determines how long an invocation can run for
+     */
+    fun timeout(timeout: Long, unit: TimeUnit) {
+        this.timeout = unit.toMillis(timeout)
+    }
 }
 
 /**

src/main/kotlin/com/mrkirby153/botcore/command/slashcommand/dsl/DslCommandExecutor.kt

@@ -10,6 +10,7 @@ import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.withTimeout
 import net.dv8tion.jda.api.JDA
 import net.dv8tion.jda.api.entities.Guild
 import net.dv8tion.jda.api.events.GenericEvent
@@ -291,32 +292,49 @@ class DslCommandExecutor private constructor(
         override suspend fun onEvent(event: GenericEvent) {
             when (event) {
                 is SlashCommandInteractionEvent -> {
-                    if (getSlashCommand(event) != null) {
+                    val slashCommand = getSlashCommand(event)
+                    if (slashCommand != null) {
                         scope.launch {
-                            execute(event, this)
+                            log.trace("Running command ${slashCommand.name} with a timeout of ${slashCommand.timeout}")
+                            withTimeout(slashCommand.timeout) {
+                                execute(event, this)
+                            }
+
                         }
                     }
                 }
 
                 is CommandAutoCompleteInteractionEvent -> {
-                    if (getSlashCommand(event) != null) {
+                    val slashCommand = getSlashCommand(event)
+                    if (slashCommand != null) {
                         scope.launch {
-                            handleAutocomplete(event)
+                            log.trace("Running autocomplete handler ${slashCommand.name} with a timeout of ${slashCommand.timeout}")
+                            withTimeout(slashCommand.autocompleteTimeout) {
+                                handleAutocomplete(event)
+                            }
                         }
                     }
                 }
 
                 is UserContextInteractionEvent -> {
                     scope.launch {
                         userContextCommands.firstOrNull { it.name == event.name }
-                            ?.execute(UserContext(event))
+                            ?.apply {
+                                withTimeout(timeout) {
+                                    execute(UserContext(event, scope))
+                                }
+                            }
                     }
                 }
 
                 is MessageContextInteractionEvent -> {
                     scope.launch {
                         messageContextCommands.firstOrNull { it.name == event.name }
-                            ?.execute(MessageContext(event))
+                            ?.apply {
+                                withTimeout(timeout) {
+                                    execute(MessageContext(event, scope))
+                                }
+                            }
                     }
                 }
             }

src/main/kotlin/com/mrkirby153/botcore/command/slashcommand/dsl/SlashCommand.kt

@@ -6,6 +6,7 @@ import net.dv8tion.jda.api.events.interaction.command.CommandAutoCompleteInterac
 import net.dv8tion.jda.api.events.interaction.command.SlashCommandInteractionEvent
 import net.dv8tion.jda.api.interactions.commands.Command
 import net.dv8tion.jda.api.interactions.commands.DefaultMemberPermissions
+import java.util.concurrent.TimeUnit
 
 open class AbstractSlashCommand(
     val name: String
@@ -15,6 +16,9 @@ open class AbstractSlashCommand(
 
     internal var action: (suspend SlashContext.() -> Unit)? = null
 
+    internal var timeout: Long = 30_000 // 30 seconds
+    internal var autocompleteTimeout: Long = 5_000 // 5 seconds
+
     /**
      * Executes the slash command. If [body] is null, this is a no-op
      */
@@ -40,6 +44,20 @@ open class AbstractSlashCommand(
     internal open fun addArgument(name: String, argument: ArgumentContainer<*, *>) {
         arguments[name] = argument
     }
+
+    /**
+     * Sets how long this command is allowed to execute before it times out
+     */
+    fun commandTimeout(time: Long, timeUnit: TimeUnit) {
+        timeout = timeUnit.toMillis(time)
+    }
+
+    /**
+     * Sets how long any autocomplete query can take
+     */
+    fun autocompleteTimeout(time: Long, timeUnit: TimeUnit) {
+        autocompleteTimeout = timeUnit.toMillis(time)
+    }
 }
 
 @SlashDsl