From 07720e23cb35fe7906dc26ea8de91c385321d0c9 Mon Sep 17 00:00:00 2001 From: Lanius Trolling Date: Tue, 23 Apr 2024 11:33:49 -0400 Subject: [PATCH] Improve rate-limit handling --- .../info/mechyrdia/robot/RobotRateLimiter.kt | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/src/jvmMain/kotlin/info/mechyrdia/robot/RobotRateLimiter.kt b/src/jvmMain/kotlin/info/mechyrdia/robot/RobotRateLimiter.kt index c4f2f4d..e5f5812 100644 --- a/src/jvmMain/kotlin/info/mechyrdia/robot/RobotRateLimiter.kt +++ b/src/jvmMain/kotlin/info/mechyrdia/robot/RobotRateLimiter.kt @@ -48,9 +48,9 @@ val RobotRateLimiter = createClientPlugin("RobotRateLimiter") { val tokensResetAt = AtomicLong(0) onRequest { request, _ -> - val requestDelay = calculateRateLimitDelayDouble(requestsRemaining.getAcquire(), requestsResetAt.getAcquire()) - val tokenDelay = request.attributes.getTokens()?.let { _ -> - calculateRateLimitDelayDouble(tokensRemaining.getAcquire(), tokensResetAt.getAcquire()) + val requestDelay = calculateRateLimitDelayDouble(requestsRemaining.getAndAdd(-1), requestsResetAt.get()) + val tokenDelay = request.attributes.getTokens()?.let { tokens -> + calculateRateLimitDelayDouble(tokensRemaining.getAndAdd(-tokens), tokensResetAt.get()) } combineDelays(requestDelay, tokenDelay)?.seconds?.let { delay(it) } @@ -58,15 +58,10 @@ val RobotRateLimiter = createClientPlugin("RobotRateLimiter") { @Suppress("UastIncorrectHttpHeaderInspection") onResponse { response -> - val newRequestsRemaining = response.headers["X-Ratelimit-Remaining-Requests"]?.toIntOrNull() ?: -1 - val newRequestsResetAt = response.headers["X-Ratelimit-Reset-Requests"]?.parseDurationToSeconds()?.secondFromNow() ?: 0 - val newTokensRemaining = response.headers["X-Ratelimit-Remaining-Tokens"]?.toIntOrNull() ?: -1 - val newTokensResetAt = response.headers["X-Ratelimit-Reset-Tokens"]?.parseDurationToSeconds()?.secondFromNow() ?: 0 - - requestsRemaining.setRelease(newRequestsRemaining) - requestsResetAt.setRelease(newRequestsResetAt) - tokensRemaining.setRelease(newTokensRemaining) - tokensResetAt.setRelease(newTokensResetAt) + response.headers["X-Ratelimit-Remaining-Requests"]?.toIntOrNull()?.let(requestsRemaining::set) + response.headers["X-Ratelimit-Reset-Requests"]?.parseDurationToSeconds()?.secondFromNow()?.let(requestsResetAt::set) + response.headers["X-Ratelimit-Remaining-Tokens"]?.toIntOrNull()?.let(tokensRemaining::set) + response.headers["X-Ratelimit-Reset-Tokens"]?.parseDurationToSeconds()?.secondFromNow()?.let(tokensResetAt::set) } } -- 2.25.1