63 lines
2.3 KiB
Kotlin
63 lines
2.3 KiB
Kotlin
package dev.inmo.tgbotapi.bot.settings.limiters
|
|
|
|
import dev.inmo.micro_utils.coroutines.safely
|
|
import dev.inmo.tgbotapi.bot.exceptions.TooMuchRequestsException
|
|
import dev.inmo.tgbotapi.types.*
|
|
import io.ktor.client.features.ClientRequestException
|
|
import io.ktor.http.HttpStatusCode
|
|
import kotlinx.coroutines.delay
|
|
import kotlinx.coroutines.flow.*
|
|
|
|
/**
|
|
* This limiter will limit requests only after getting a [RetryAfterError] or [ClientRequestException] with
|
|
* [HttpStatusCode.TooManyRequests] status code. Important thing is that in case if some of block has been blocked, all
|
|
* the others will wait until it will be possible to be called
|
|
*
|
|
* @param defaultTooManyRequestsDelay This parameter will be used in case of getting [ClientRequestException] with
|
|
* [HttpStatusCode.TooManyRequests] as a parameter for delay like it would be [TooMuchRequestsException]. The reason of
|
|
* it is that in [ClientRequestException] there is no information about required delay between requests
|
|
*/
|
|
class ExceptionsOnlyLimiter(
|
|
private val defaultTooManyRequestsDelay: MilliSeconds = 1000L
|
|
) : RequestLimiter {
|
|
private val lockState = MutableStateFlow(false)
|
|
private suspend fun lock(timeMillis: MilliSeconds) {
|
|
try {
|
|
safely {
|
|
lockState.emit(true)
|
|
delay(timeMillis)
|
|
}
|
|
} finally {
|
|
lockState.emit(false)
|
|
}
|
|
}
|
|
|
|
override suspend fun <T> limit(block: suspend () -> T): T {
|
|
while (true) {
|
|
lockState.first { !it }
|
|
val result = safely({
|
|
when (it) {
|
|
is TooMuchRequestsException -> {
|
|
lock(it.retryAfter.leftToRetry)
|
|
Result.failure(it)
|
|
}
|
|
is ClientRequestException -> {
|
|
if (it.response.status == HttpStatusCode.TooManyRequests) {
|
|
lock(defaultTooManyRequestsDelay)
|
|
} else {
|
|
throw it
|
|
}
|
|
Result.failure(it)
|
|
}
|
|
else -> throw it
|
|
}
|
|
}) {
|
|
Result.success(block())
|
|
}
|
|
if (result.isSuccess) {
|
|
return result.getOrNull()!!
|
|
}
|
|
}
|
|
}
|
|
}
|