Improve rate limit handling, minor refactor

This commit is contained in:
Zed 2023-08-30 03:04:22 +02:00
parent 986b91ac73
commit 898b19b92f
3 changed files with 47 additions and 50 deletions

View File

@ -61,13 +61,6 @@ proc genHeaders*(url, oauthToken, oauthTokenSecret: string): HttpHeaders =
"DNT": "1" "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.} = template fetchImpl(result, fetchBody) {.dirty.} =
once: once:
pool = HttpPool() pool = HttpPool()
@ -89,28 +82,46 @@ template fetchImpl(result, fetchBody) {.dirty.} =
badClient = true badClient = true
raise newException(BadClientError, "Bad client") 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 result.len > 0:
if resp.headers.getOrDefault("content-encoding") == "gzip": if resp.headers.getOrDefault("content-encoding") == "gzip":
result = uncompress(result, dfGzip) 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 fetchBody
release(account, used=true)
if resp.status == $Http400: if resp.status == $Http400:
raise newException(InternalError, $url) raise newException(InternalError, $url)
except InternalError as e: except InternalError as e:
raise e raise e
except BadClientError as e: except BadClientError as e:
release(account, used=true) raise e
except OSError as e:
raise e raise e
except Exception as e: except Exception as e:
echo "error: ", e.name, ", msg: ", e.msg, ", accountId: ", account.id, ", url: ", url 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() raise rateLimitError()
finally:
release(account)
proc fetch*(url: Uri; api: Api): Future[JsonNode] {.async.} = proc fetch*(url: Uri; api: Api): Future[JsonNode] {.async.} =
var body: string var body: string
@ -121,36 +132,14 @@ proc fetch*(url: Uri; api: Api): Future[JsonNode] {.async.} =
echo resp.status, ": ", body, " --- url: ", url echo resp.status, ": ", body, " --- url: ", url
result = newJNull() result = newJNull()
updateAccount()
let error = result.getError let error = result.getError
if error in {invalidToken, badToken}: if error in {expiredToken, badToken}:
echo "fetch error: ", result.getError echo "fetchBody error: ", error
release(account, invalid=true) invalidate(account)
raise rateLimitError() 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.} = proc fetchRaw*(url: Uri; api: Api): Future[string] {.async.} =
fetchImpl result: fetchImpl result:
if not (result.startsWith('{') or result.startsWith('[')): if not (result.startsWith('{') or result.startsWith('[')):
echo resp.status, ": ", result, " --- url: ", url echo resp.status, ": ", result, " --- url: ", url
result.setLen(0) 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()

View File

@ -11,7 +11,7 @@ var
accountPool: seq[GuestAccount] accountPool: seq[GuestAccount]
enableLogging = false enableLogging = false
template log(str) = template log(str: varargs[string, `$`]) =
if enableLogging: echo "[accounts] ", str if enableLogging: echo "[accounts] ", str
proc getPoolJson*(): JsonNode = proc getPoolJson*(): JsonNode =
@ -91,7 +91,7 @@ proc isLimited(account: GuestAccount; api: Api): bool =
if limit.limited and (epochTime().int - limit.limitedAt) > dayInSeconds: if limit.limited and (epochTime().int - limit.limitedAt) > dayInSeconds:
account.apis[api].limited = false 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) return limit.limited or (limit.remaining <= 10 and limit.reset > epochTime().int)
else: else:
@ -100,14 +100,17 @@ proc isLimited(account: GuestAccount; api: Api): bool =
proc isReady(account: GuestAccount; api: Api): bool = proc isReady(account: GuestAccount; api: Api): bool =
not (account.isNil or account.pending > maxConcurrentReqs or account.isLimited(api)) 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 account.isNil: return
if invalid: log "invalidating expired account: ", account.id
log "discarding invalid account: " & account.id
# TODO: This isn't sufficient, but it works for now
let idx = accountPool.find(account) let idx = accountPool.find(account)
if idx > -1: accountPool.delete(idx) if idx > -1: accountPool.delete(idx)
elif used: account = nil
proc release*(account: GuestAccount; invalid=false) =
if account.isNil: return
dec account.pending dec account.pending
proc getGuestAccount*(api: Api): Future[GuestAccount] {.async.} = proc getGuestAccount*(api: Api): Future[GuestAccount] {.async.} =
@ -119,9 +122,14 @@ proc getGuestAccount*(api: Api): Future[GuestAccount] {.async.} =
if not result.isNil and result.isReady(api): if not result.isNil and result.isReady(api):
inc result.pending inc result.pending
else: else:
log "no accounts available for API: " & $api log "no accounts available for API: ", api
raise rateLimitError() 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) = proc setRateLimit*(account: GuestAccount; api: Api; remaining, reset: int) =
# avoid undefined behavior in race conditions # avoid undefined behavior in race conditions
if api in account.apis: if api in account.apis:

View File

@ -56,7 +56,7 @@ type
userNotFound = 50 userNotFound = 50
suspended = 63 suspended = 63
rateLimited = 88 rateLimited = 88
invalidToken = 89 expiredToken = 89
listIdOrSlug = 112 listIdOrSlug = 112
tweetNotFound = 144 tweetNotFound = 144
tweetNotAuthorized = 179 tweetNotAuthorized = 179