Track token rate limits per endpoint

src/api.nim

@@ -6,57 +6,57 @@ import types, query, formatters, consts, apiutils, parser
 proc getGraphListBySlug*(name, list: string): Future[List] {.async.} =
   let
     variables = %*{"screenName": name, "listSlug": list, "withHighlightedLabel": false}
-    js = await fetch(graphList ? {"variables": $variables})
-  result = parseGraphList(js)
+    url = graphListBySlug ? {"variables": $variables}
+  result = parseGraphList(await fetch(url, Api.listBySlug))
 
 proc getGraphList*(id: string): Future[List] {.async.} =
   let
     variables = %*{"listId": id, "withHighlightedLabel": false}
-    js = await fetch(graphListId ? {"variables": $variables})
-  result = parseGraphList(js)
+    url = graphList ? {"variables": $variables}
+  result = parseGraphList(await fetch(url, Api.list))
 
 proc getListTimeline*(id: string; after=""): Future[Timeline] {.async.} =
   if id.len == 0: return
   let
     ps = genParams({"list_id": id, "ranking_mode": "reverse_chronological"}, after)
     url = listTimeline ? ps
-  result = parseTimeline(await fetch(url), after)
+  result = parseTimeline(await fetch(url, Api.timeline), after)
 
 proc getListMembers*(list: List; after=""): Future[Result[Profile]] {.async.} =
   if list.id.len == 0: return
   let
     ps = genParams({"list_id": list.id}, after)
     url = listMembers ? ps
-  result = parseListMembers(await fetch(url, oldApi=true), after)
+  result = parseListMembers(await fetch(url, Api.listMembers), after)
 
 proc getProfile*(username: string): Future[Profile] {.async.} =
   let
     ps = genParams({"screen_name": username})
-    js = await fetch(userShow ? ps, oldApi=true)
+    js = await fetch(userShow ? ps, Api.userShow)
   result = parseUserShow(js, username=username)
 
 proc getProfileById*(userId: string): Future[Profile] {.async.} =
   let
     ps = genParams({"user_id": userId})
-    js = await fetch(userShow ? ps, oldApi=true)
+    js = await fetch(userShow ? ps, Api.userShow)
   result = parseUserShow(js, id=userId)
 
 proc getTimeline*(id: string; after=""; replies=false): Future[Timeline] {.async.} =
   let
     ps = genParams({"userId": id, "include_tweet_replies": $replies}, after)
     url = timeline / (id & ".json") ? ps
-  result = parseTimeline(await fetch(url), after)
+  result = parseTimeline(await fetch(url, Api.timeline), after)
 
 proc getMediaTimeline*(id: string; after=""): Future[Timeline] {.async.} =
   let url = mediaTimeline / (id & ".json") ? genParams(cursor=after)
-  result = parseTimeline(await fetch(url), after)
+  result = parseTimeline(await fetch(url, Api.timeline), after)
 
 proc getPhotoRail*(name: string): Future[PhotoRail] {.async.} =
   let
     ps = genParams({"screen_name": name, "trim_user": "true"},
                    count="18", ext=false)
     url = photoRail ? ps
-  result = parsePhotoRail(await fetch(url, oldApi=true))
+  result = parsePhotoRail(await fetch(url, Api.photoRail))
 
 proc getSearch*[T](query: Query; after=""): Future[Result[T]] {.async.} =
   when T is Profile:
@@ -74,14 +74,14 @@ proc getSearch*[T](query: Query; after=""): Future[Result[T]] {.async.} =
 
   let url = search ? genParams(searchParams & @[("q", q), searchMode], after)
   try:
-    result = parse(await fetch(url), after)
+    result = parse(await fetch(url, Api.search), after)
     result.query = query
   except InternalError:
     return Result[T](beginning: true, query: query)
 
 proc getTweetImpl(id: string; after=""): Future[Conversation] {.async.} =
   let url = tweet / (id & ".json") ? genParams(cursor=after)
-  result = parseConversation(await fetch(url), id)
+  result = parseConversation(await fetch(url, Api.tweet), id)
 
 proc getReplies*(id, after: string): Future[Result[Chain]] {.async.} =
   result = (await getTweetImpl(id, after)).replies

src/apiutils.nim

@@ -3,7 +3,9 @@ import httpclient, asyncdispatch, options, times, strutils, uri
 import packedjson, zippy
 import types, tokens, consts, parserutils, http_pool
 
-const rl = "x-rate-limit-"
+const
+  rlRemaining = "x-rate-limit-remaining"
+  rlReset = "x-rate-limit-reset"
 
 var pool: HttpPool
 
@@ -38,11 +40,11 @@ proc genHeaders*(token: Token = nil): HttpHeaders =
     "DNT": "1"
   })
 
-proc fetch*(url: Uri; oldApi=false): Future[JsonNode] {.async.} =
+proc fetch*(url: Uri; api: Api): Future[JsonNode] {.async.} =
   once:
     pool = HttpPool()
 
-  var token = await getToken()
+  var token = await getToken(api)
   if token.tok.len == 0:
     raise rateLimitError()
 
@@ -65,9 +67,14 @@ proc fetch*(url: Uri; oldApi=false): Future[JsonNode] {.async.} =
       echo resp.status, ": ", body
       result = newJNull()
 
-    if not oldApi and resp.headers.hasKey(rl & "reset"):
-      token.remaining = parseInt(resp.headers[rl & "remaining"])
-      token.reset = fromUnix(parseInt(resp.headers[rl & "reset"]))
+    if api != Api.search and resp.headers.hasKey(rlRemaining):
+      let
+        remaining = parseInt(resp.headers[rlRemaining])
+        reset = parseInt(resp.headers[rlReset])
+      token.setRateLimit(api, remaining, reset)
+      echo api, " ", remaining, " ", url.path
+    else:
+      echo api, " ", url.path
 
     if result.getError notin {invalidToken, forbidden, badToken}:
       token.lastUse = getTime()

src/consts.nim

@@ -19,8 +19,8 @@ const
   tweet* = timelineApi / "conversation"
 
   graphql = api / "graphql"
-  graphList* = graphql / "ErWsz9cObLel1BF-HjuBlA/ListBySlug"
-  graphListId* = graphql / "JADTh6cjebfgetzvF3tQvQ/List"
+  graphListBySlug* = graphql / "ErWsz9cObLel1BF-HjuBlA/ListBySlug"
+  graphList* = graphql / "JADTh6cjebfgetzvF3tQvQ/List"
 
   timelineParams* = {
     "include_profile_interstitial_type": "0",

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)

src/types.nim

@@ -8,12 +8,25 @@ type
   RateLimitError* = object of CatchableError
   InternalError* = object of CatchableError
 
-  Token* = ref object
-    tok*: string
+  Api* {.pure.} = enum
+    userShow
+    photoRail
+    timeline
+    search
+    tweet
+    list
+    listBySlug
+    listMembers
+
+  RateLimit* = object
     remaining*: int
     reset*: Time
+
+  Token* = ref object
+    tok*: string
     init*: Time
     lastUse*: Time
+    apis*: Table[Api, RateLimit]
 
   Error* = enum
     null = 0