From 51b1567af67d53591a9afeaeb4aaf8c04ab824ae Mon Sep 17 00:00:00 2001
From: Zed <zedeus@pm.me>
Date: Wed, 13 Jan 2021 14:32:26 +0100
Subject: [PATCH] Improve token pool to prevent rate limits

---
 src/apiutils.nim | 15 ++++++++-------
 src/tokens.nim   | 43 ++++++++++++++++++++++++++++++-------------
 2 files changed, 38 insertions(+), 20 deletions(-)

diff --git a/src/apiutils.nim b/src/apiutils.nim
index 89c7687..3d05074 100644
--- a/src/apiutils.nim
+++ b/src/apiutils.nim
@@ -31,9 +31,6 @@ proc genHeaders*(token: Token = nil): HttpHeaders =
     "DNT": "1"
   })
 
-proc rateLimitError(): ref RateLimitError =
-  newException(RateLimitError, "rate limited with " & getPoolInfo())
-
 proc fetch*(url: Uri; oldApi=false): Future[JsonNode] {.async.} =
   once:
     pool = HttpPool()
@@ -54,12 +51,16 @@ proc fetch*(url: Uri; oldApi=false): Future[JsonNode] {.async.} =
       echo resp.status, ": ", body
       result = newJNull()
 
-    if not oldApi and resp.headers.hasKey(rl & "limit"):
-      token.remaining = parseInt(resp.headers[rl & "remaining"])
-      token.reset = fromUnix(parseInt(resp.headers[rl & "reset"]))
+    if not oldApi and resp.headers.hasKey(rl & "reset"):
+      let time = fromUnix(parseInt(resp.headers[rl & "reset"]))
+      if token.reset != time:
+        token.remaining = parseInt(resp.headers[rl & "limit"])
+      token.reset = time
 
     if result.getError notin {invalidToken, forbidden, badToken}:
-      token.release()
+      token.lastUse = getTime()
+    else:
+      echo "fetch error: ", result.getError
   except Exception:
     echo "error: ", url
     raise rateLimitError()
diff --git a/src/tokens.nim b/src/tokens.nim
index 5741b9d..799f849 100644
--- a/src/tokens.nim
+++ b/src/tokens.nim
@@ -1,4 +1,4 @@
-import asyncdispatch, httpclient, times, sequtils, json, math
+import asyncdispatch, httpclient, times, sequtils, json, math, random
 import strutils, strformat
 import types, agents, consts, http_pool
 
@@ -6,11 +6,20 @@ var
   clientPool {.threadvar.}: HttpPool
   tokenPool {.threadvar.}: seq[Token]
   lastFailed: Time
-  minFail = initDuration(seconds=10)
+  minFail = initDuration(minutes=30)
+
+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())
 
 proc fetchToken(): Future[Token] {.async.} =
   if getTime() - lastFailed < minFail:
-    return Token()
+    raise rateLimitError()
 
   let headers = newHttpHeaders({
     "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
@@ -33,7 +42,6 @@ proc fetchToken(): Future[Token] {.async.} =
                    init: time, lastUse: time)
   except Exception as e:
     lastFailed = getTime()
-    result = Token()
     echo "fetching token failed: ", e.msg
 
 proc expired(token: Token): bool {.inline.} =
@@ -49,19 +57,25 @@ proc isLimited(token: Token): bool {.inline.} =
     token.expired
 
 proc release*(token: Token) =
-  if token != nil and not token.expired:
-    token.lastUse = getTime()
-    tokenPool.insert(token)
+  if token != nil and token.expired:
+    tokenPool.delete(tokenPool.find(token))
 
 proc getToken*(): Future[Token] {.async.} =
   for i in 0 ..< tokenPool.len:
     if not result.isLimited: break
     result.release()
-    result = tokenPool.pop()
+    result = tokenPool.sample()
 
   if result.isLimited:
     result.release()
     result = await fetchToken()
+    tokenPool.add result
+    echo getPoolInfo()
+
+  if result == nil:
+    raise rateLimitError()
+
+  dec result.remaining
 
 proc poolTokens*(amount: int) {.async.} =
   var futs: seq[Future[Token]]
@@ -69,7 +83,14 @@ proc poolTokens*(amount: int) {.async.} =
     futs.add fetchToken()
 
   for token in futs:
-    release(await token)
+    var newToken: Token
+
+    try: newToken = await token
+    except: discard
+
+    if newToken != nil:
+      tokenPool.add newToken
+      echo getPoolInfo()
 
 proc initTokenPool*(cfg: Config) {.async.} =
   clientPool = HttpPool()
@@ -78,7 +99,3 @@ proc initTokenPool*(cfg: Config) {.async.} =
     if tokenPool.countIt(not it.isLimited) < cfg.minTokens:
       await poolTokens(min(4, cfg.minTokens - tokenPool.len))
     await sleepAsync(2000)
-
-proc getPoolInfo*: string =
-  let avg = tokenPool.mapIt(it.remaining).sum()
-  return &"{tokenPool.len} tokens, average remaining: {avg}"