From: Lanius Trolling <lanius@laniustrolling.dev>
Date: Tue, 23 Apr 2024 15:33:49 +0000 (-0400)
Subject: Improve rate-limit handling
X-Git-Url: https://gitweb.starshipfights.net/?a=commitdiff_plain;h=07720e23cb35fe7906dc26ea8de91c385321d0c9;p=factbooks

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 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)
 	}
 }