Track token rate limits per endpoint

src/tokens.nim

@@ -1,13 +1,12 @@
 # SPDX-License-Identifier: AGPL-3.0-only
-import asyncdispatch, httpclient, times, sequtils, json, math, random
-import strutils, strformat
+import asyncdispatch, httpclient, times, sequtils, json, random
+import strutils, tables
 import zippy
 import types, agents, consts, http_pool
 
 const
-  expirationTime = 3.hours
-  maxLastUse = 1.hours
-  resetPeriod = 15.minutes
+  maxAge = 3.hours      # tokens expire after 3 hours
+  maxLastUse = 1.hours  # if a token is unused for 60 minutes, it expires
   failDelay = initDuration(minutes=30)
 
 var
@@ -15,14 +14,9 @@ var
   tokenPool: seq[Token]
   lastFailed: Time
 
-proc getPoolInfo*: string =
-  if tokenPool.len == 0: return "token pool empty"
-
-  let avg = tokenPool.mapIt(it.remaining).sum() div tokenPool.len
-  return &"{tokenPool.len} tokens, average remaining: {avg}"
 
 proc rateLimitError*(): ref RateLimitError =
-  newException(RateLimitError, "rate limited with " & getPoolInfo())
+  newException(RateLimitError, "rate limited")
 
 proc fetchToken(): Future[Token] {.async.} =
   if getTime() - lastFailed < failDelay:
@@ -37,51 +31,58 @@ proc fetchToken(): Future[Token] {.async.} =
     "authorization": auth
   })
 
-  var
-    resp: string
-    tokNode: JsonNode
-    tok: string
-
   try:
-    resp = clientPool.use(headers): await c.postContent(activate)
-    tokNode = parseJson(uncompress(resp))["guest_token"]
-    tok = tokNode.getStr($(tokNode.getInt))
+    let
+      resp = clientPool.use(headers): await c.postContent(activate)
+      tokNode = parseJson(uncompress(resp))["guest_token"]
+      tok = tokNode.getStr($(tokNode.getInt))
+      time = getTime()
 
-    let time = getTime()
-    result = Token(tok: tok, remaining: 187, reset: time + resetPeriod,
-                   init: time, lastUse: time)
+    return Token(tok: tok, init: time, lastUse: time)
   except Exception as e:
     lastFailed = getTime()
     echo "fetching token failed: ", e.msg
 
-template expired(token: Token): untyped =
+proc expired(token: Token): bool =
   let time = getTime()
-  token.init < time - expirationTime or
-    token.lastUse < time - maxLastUse
+  token.init < time - maxAge or token.lastUse < time - maxLastUse
 
-template isLimited(token: Token): untyped =
-  token == nil or (token.remaining <= 5 and token.reset > getTime()) or
-    token.expired
+proc isLimited(token: Token; api: Api): bool =
+  if token.isNil or token.expired:
+    return true
+
+  if api in token.apis:
+    let limit = token.apis[api]
+    return (limit.remaining <= 5 and limit.reset > getTime())
+  else:
+    return false
 
 proc release*(token: Token; invalid=false) =
-  if token != nil and (invalid or token.expired):
+  if not token.isNil and (invalid or token.expired):
     let idx = tokenPool.find(token)
     if idx > -1: tokenPool.delete(idx)
 
-proc getToken*(): Future[Token] {.async.} =
+proc getToken*(api: Api): Future[Token] {.async.} =
   for i in 0 ..< tokenPool.len:
-    if not result.isLimited: break
+    if not (result.isNil or result.isLimited(api)):
+      break
     release(result)
     result = tokenPool.sample()
 
-  if result.isLimited:
+  if result.isNil or result.isLimited(api):
     release(result)
     result = await fetchToken()
     tokenPool.add result
 
-  if result == nil:
+  if result.isNil:
     raise rateLimitError()
 
+proc setRateLimit*(token: Token; api: Api; remaining, reset: int) =
+  token.apis[api] = RateLimit(
+    remaining: remaining,
+    reset: fromUnix(reset)
+  )
+
 proc poolTokens*(amount: int) {.async.} =
   var futs: seq[Future[Token]]
   for i in 0 ..< amount:
@@ -93,13 +94,13 @@ proc poolTokens*(amount: int) {.async.} =
     try: newToken = await token
     except: discard
 
-    if newToken != nil:
+    if not newToken.isNil:
       tokenPool.add newToken
 
 proc initTokenPool*(cfg: Config) {.async.} =
   clientPool = HttpPool()
 
   while true:
-    if tokenPool.countIt(not it.isLimited) < cfg.minTokens:
+    if tokenPool.countIt(not it.isLimited(Api.timeline)) < cfg.minTokens:
       await poolTokens(min(4, cfg.minTokens - tokenPool.len))
     await sleepAsync(2000)