Improve rate-limit handling

diff --git a/src/jvmMain/kotlin/info/mechyrdia/robot/RobotRateLimiter.kt b/src/jvmMain/kotlin/info/mechyrdia/robot/RobotRateLimiter.kt
index c4f2f4d1a7fd0e1a9f3ebd3f775a1900ee810f95..e5f581229ead225b19f607e9ae3aa5398997b434 100644 (file)
--- 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)
        }
 }