From dd71e60f358952b3442b47c73078d9f8c9d15878 Mon Sep 17 00:00:00 2001 From: Zed Date: Wed, 5 Jan 2022 22:48:45 +0100 Subject: [PATCH] Track token rate limits per endpoint --- src/api.nim | 26 +++++++++--------- src/apiutils.nim | 19 +++++++++---- src/consts.nim | 4 +-- src/tokens.nim | 71 ++++++++++++++++++++++++------------------------ src/types.nim | 17 ++++++++++-- 5 files changed, 79 insertions(+), 58 deletions(-) diff --git a/src/api.nim b/src/api.nim index 260435e..9a6e70a 100644 --- a/src/api.nim +++ b/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 diff --git a/src/apiutils.nim b/src/apiutils.nim index 4f9af56..ada68a9 100644 --- a/src/apiutils.nim +++ b/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() diff --git a/src/consts.nim b/src/consts.nim index c7651b4..efa0420 100644 --- a/src/consts.nim +++ b/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", diff --git a/src/tokens.nim b/src/tokens.nim index ba90ad2..da9f36a 100644 --- a/src/tokens.nim +++ b/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) diff --git a/src/types.nim b/src/types.nim index 56f54dc..53b171c 100644 --- a/src/types.nim +++ b/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