From 898b19b92f3121b2185ba3a59381be2ea8299f44 Mon Sep 17 00:00:00 2001
From: Zed <zedeus@pm.me>
Date: Wed, 30 Aug 2023 03:04:22 +0200
Subject: [PATCH] Improve rate limit handling, minor refactor

---
 src/apiutils.nim | 67 ++++++++++++++++++++----------------------------
 src/tokens.nim   | 28 ++++++++++++--------
 src/types.nim    |  2 +-
 3 files changed, 47 insertions(+), 50 deletions(-)

diff --git a/src/apiutils.nim b/src/apiutils.nim
index 453b36a..37afded 100644
--- a/src/apiutils.nim
+++ b/src/apiutils.nim
@@ -61,13 +61,6 @@ proc genHeaders*(url, oauthToken, oauthTokenSecret: string): HttpHeaders =
     "DNT": "1"
   })
 
-template updateAccount() =
-  if resp.headers.hasKey(rlRemaining):
-    let
-      remaining = parseInt(resp.headers[rlRemaining])
-      reset = parseInt(resp.headers[rlReset])
-    account.setRateLimit(api, remaining, reset)
-
 template fetchImpl(result, fetchBody) {.dirty.} =
   once:
     pool = HttpPool()
@@ -89,28 +82,46 @@ template fetchImpl(result, fetchBody) {.dirty.} =
         badClient = true
         raise newException(BadClientError, "Bad client")
 
+    if resp.headers.hasKey(rlRemaining):
+      let
+        remaining = parseInt(resp.headers[rlRemaining])
+        reset = parseInt(resp.headers[rlReset])
+      account.setRateLimit(api, remaining, reset)
+
     if result.len > 0:
       if resp.headers.getOrDefault("content-encoding") == "gzip":
         result = uncompress(result, dfGzip)
-      else:
-        echo "non-gzip body, url: ", url, ", body: ", result
+
+      if result.startsWith("{\"errors"):
+        let errors = result.fromJson(Errors)
+        if errors in {expiredToken, badToken}:
+          echo "fetch error: ", errors
+          invalidate(account)
+          raise rateLimitError()
+        elif errors in {rateLimited}:
+          # rate limit hit, resets after 24 hours
+          setLimited(account, api)
+          raise rateLimitError()
+      elif result.startsWith("429 Too Many Requests"):
+        account.apis[api].remaining = 0
+        # rate limit hit, resets after the 15 minute window
+        raise rateLimitError()
 
     fetchBody
 
-    release(account, used=true)
-
     if resp.status == $Http400:
       raise newException(InternalError, $url)
   except InternalError as e:
     raise e
   except BadClientError as e:
-    release(account, used=true)
+    raise e
+  except OSError as e:
     raise e
   except Exception as e:
     echo "error: ", e.name, ", msg: ", e.msg, ", accountId: ", account.id, ", url: ", url
-    if "length" notin e.msg and "descriptor" notin e.msg:
-      release(account, invalid=true)
     raise rateLimitError()
+  finally:
+    release(account)
 
 proc fetch*(url: Uri; api: Api): Future[JsonNode] {.async.} =
   var body: string
@@ -121,36 +132,14 @@ proc fetch*(url: Uri; api: Api): Future[JsonNode] {.async.} =
       echo resp.status, ": ", body, " --- url: ", url
       result = newJNull()
 
-    updateAccount()
-
     let error = result.getError
-    if error in {invalidToken, badToken}:
-      echo "fetch error: ", result.getError
-      release(account, invalid=true)
+    if error in {expiredToken, badToken}:
+      echo "fetchBody error: ", error
+      invalidate(account)
       raise rateLimitError()
 
-    if body.startsWith("{\"errors"):
-      let errors = body.fromJson(Errors)
-      if errors in {invalidToken, badToken}:
-        echo "fetch error: ", errors
-        release(account, invalid=true)
-        raise rateLimitError()
-      elif errors in {rateLimited}:
-        account.apis[api].limited = true
-        account.apis[api].limitedAt = epochTime().int
-        echo "[accounts] rate limited, api: ", api, ", reqs left: ", account.apis[api].remaining, ", id: ", account.id
-
 proc fetchRaw*(url: Uri; api: Api): Future[string] {.async.} =
   fetchImpl result:
     if not (result.startsWith('{') or result.startsWith('[')):
       echo resp.status, ": ", result, " --- url: ", url
       result.setLen(0)
-
-    updateAccount()
-
-    if result.startsWith("{\"errors"):
-      let errors = result.fromJson(Errors)
-      if errors in {invalidToken, badToken}:
-        echo "fetch error: ", errors
-        release(account, invalid=true)
-        raise rateLimitError()
diff --git a/src/tokens.nim b/src/tokens.nim
index a3af9bf..c620bc7 100644
--- a/src/tokens.nim
+++ b/src/tokens.nim
@@ -11,7 +11,7 @@ var
   accountPool: seq[GuestAccount]
   enableLogging = false
 
-template log(str) =
+template log(str: varargs[string, `$`]) =
   if enableLogging: echo "[accounts] ", str
 
 proc getPoolJson*(): JsonNode =
@@ -91,7 +91,7 @@ proc isLimited(account: GuestAccount; api: Api): bool =
 
     if limit.limited and (epochTime().int - limit.limitedAt) > dayInSeconds:
       account.apis[api].limited = false
-      log "resetting limit, api: " & $api & ", id: " & $account.id
+      log "resetting limit, api: ", api, ", id: ", account.id
 
     return limit.limited or (limit.remaining <= 10 and limit.reset > epochTime().int)
   else:
@@ -100,15 +100,18 @@ proc isLimited(account: GuestAccount; api: Api): bool =
 proc isReady(account: GuestAccount; api: Api): bool =
   not (account.isNil or account.pending > maxConcurrentReqs or account.isLimited(api))
 
-proc release*(account: GuestAccount; used=false; invalid=false) =
+proc invalidate*(account: var GuestAccount) =
   if account.isNil: return
-  if invalid:
-    log "discarding invalid account: " & account.id
+  log "invalidating expired account: ", account.id
 
-    let idx = accountPool.find(account)
-    if idx > -1: accountPool.delete(idx)
-  elif used:
-    dec account.pending
+  # TODO: This isn't sufficient, but it works for now
+  let idx = accountPool.find(account)
+  if idx > -1: accountPool.delete(idx)
+  account = nil
+
+proc release*(account: GuestAccount; invalid=false) =
+  if account.isNil: return
+  dec account.pending
 
 proc getGuestAccount*(api: Api): Future[GuestAccount] {.async.} =
   for i in 0 ..< accountPool.len:
@@ -119,9 +122,14 @@ proc getGuestAccount*(api: Api): Future[GuestAccount] {.async.} =
   if not result.isNil and result.isReady(api):
     inc result.pending
   else:
-    log "no accounts available for API: " & $api
+    log "no accounts available for API: ", api
     raise rateLimitError()
 
+proc setLimited*(account: GuestAccount; api: Api) =
+  account.apis[api].limited = true
+  account.apis[api].limitedAt = epochTime().int
+  log "rate limited, api: ", api, ", reqs left: ", account.apis[api].remaining, ", id: ", account.id
+
 proc setRateLimit*(account: GuestAccount; api: Api; remaining, reset: int) =
   # avoid undefined behavior in race conditions
   if api in account.apis:
diff --git a/src/types.nim b/src/types.nim
index fcb24c0..8a7a66e 100644
--- a/src/types.nim
+++ b/src/types.nim
@@ -56,7 +56,7 @@ type
     userNotFound = 50
     suspended = 63
     rateLimited = 88
-    invalidToken = 89
+    expiredToken = 89
     listIdOrSlug = 112
     tweetNotFound = 144
     tweetNotAuthorized = 179