Track pending token requests to limit concurrency

This commit is contained in:
Zed 2022-01-05 23:38:46 +01:00
parent f9c9b0d3a4
commit 34964f9e56
3 changed files with 32 additions and 19 deletions

View file

@ -72,15 +72,12 @@ proc fetch*(url: Uri; api: Api): Future[JsonNode] {.async.} =
remaining = parseInt(resp.headers[rlRemaining]) remaining = parseInt(resp.headers[rlRemaining])
reset = parseInt(resp.headers[rlReset]) reset = parseInt(resp.headers[rlReset])
token.setRateLimit(api, remaining, reset) token.setRateLimit(api, remaining, reset)
echo api, " ", remaining, " ", url.path
else:
echo api, " ", url.path
if result.getError notin {invalidToken, forbidden, badToken}: if result.getError notin {invalidToken, forbidden, badToken}:
token.lastUse = getTime() release(token, used=true)
else: else:
echo "fetch error: ", result.getError echo "fetch error: ", result.getError
release(token, true) release(token, invalid=true)
raise rateLimitError() raise rateLimitError()
if resp.status == $Http400: if resp.status == $Http400:
@ -90,5 +87,5 @@ proc fetch*(url: Uri; api: Api): Future[JsonNode] {.async.} =
except Exception as e: except Exception as e:
echo "error: ", e.name, ", msg: ", e.msg, ", token: ", token[], ", url: ", url echo "error: ", e.name, ", msg: ", e.msg, ", token: ", token[], ", url: ", url
if "length" notin e.msg and "descriptor" notin e.msg: if "length" notin e.msg and "descriptor" notin e.msg:
release(token, true) release(token, invalid=true)
raise rateLimitError() raise rateLimitError()

View file

@ -5,8 +5,9 @@ import zippy
import types, agents, consts, http_pool import types, agents, consts, http_pool
const const
maxAge = 3.hours # tokens expire after 3 hours maxConcurrentReqs = 5 # max requests at a time per token, to avoid race conditions
maxLastUse = 1.hours # if a token is unused for 60 minutes, it expires 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) failDelay = initDuration(minutes=30)
var var
@ -19,6 +20,7 @@ proc getPoolJson*: string =
for token in tokenPool: for token in tokenPool:
list[token.tok] = %*{ list[token.tok] = %*{
"apis": newJObject(), "apis": newJObject(),
"pending": token.pending,
"init": $token.init, "init": $token.init,
"lastUse": $token.lastUse "lastUse": $token.lastUse
} }
@ -69,35 +71,48 @@ proc isLimited(token: Token; api: Api): bool =
if api in token.apis: if api in token.apis:
let limit = token.apis[api] let limit = token.apis[api]
return (limit.remaining <= 5 and limit.reset > getTime()) return (limit.remaining <= 10 and limit.reset > getTime())
else: else:
return false return false
proc release*(token: Token; invalid=false) = proc isReady(token: Token; api: Api): bool =
if not token.isNil and (invalid or token.expired): not (token.isNil or token.pending > maxConcurrentReqs or token.isLimited(api))
proc release*(token: Token; used=false; invalid=false) =
if token.isNil: return
if invalid or token.expired:
let idx = tokenPool.find(token) let idx = tokenPool.find(token)
if idx > -1: tokenPool.delete(idx) if idx > -1: tokenPool.delete(idx)
elif used:
dec token.pending
token.lastUse = getTime()
proc getToken*(api: Api): Future[Token] {.async.} = proc getToken*(api: Api): Future[Token] {.async.} =
for i in 0 ..< tokenPool.len: for i in 0 ..< tokenPool.len:
if not (result.isNil or result.isLimited(api)): if result.isReady(api): break
break
release(result) release(result)
result = tokenPool.sample() result = tokenPool.sample()
if result.isNil or result.isLimited(api): if not result.isReady(api):
release(result) release(result)
result = await fetchToken() result = await fetchToken()
tokenPool.add result tokenPool.add result
if result.isNil: if not result.isNil:
inc result.pending
else:
raise rateLimitError() raise rateLimitError()
proc setRateLimit*(token: Token; api: Api; remaining, reset: int) = proc setRateLimit*(token: Token; api: Api; remaining, reset: int) =
token.apis[api] = RateLimit( let reset = fromUnix(reset)
remaining: remaining,
reset: fromUnix(reset) # avoid undefined behavior in race conditions
) if api in token.apis:
let limit = token.apis[api]
if limit.reset >= reset and limit.remaining < remaining:
return
token.apis[api] = RateLimit(remaining: remaining, reset: reset)
proc poolTokens*(amount: int) {.async.} = proc poolTokens*(amount: int) {.async.} =
var futs: seq[Future[Token]] var futs: seq[Future[Token]]

View file

@ -26,6 +26,7 @@ type
tok*: string tok*: string
init*: Time init*: Time
lastUse*: Time lastUse*: Time
pending*: int
apis*: Table[Api, RateLimit] apis*: Table[Api, RateLimit]
Error* = enum Error* = enum