calamity 0.1.28.0 → 0.1.28.1

raw patch · 4 files changed

+270/−87 lines, 4 filesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

API changes (from Hackage documentation)

Files

• Calamity/HTTP/Internal/Ratelimit.hs +221/−78
• Calamity/HTTP/Internal/Types.hs +42/−7
• ChangeLog.md +5/−0
• calamity.cabal +2/−2

@@ -9,20 +9,19 @@ import Calamity.HTTP.Internal.Types import Calamity.Internal.Utils +import Control.Applicative import Control.Concurrent import Control.Concurrent.Event (Event) import qualified Control.Concurrent.Event as E import Control.Concurrent.STM-import Control.Concurrent.STM.Lock (Lock) import qualified Control.Concurrent.STM.Lock as L import Control.Lens import Control.Monad  import Data.Aeson import Data.Aeson.Lens-import Data.ByteString (ByteString) import qualified Data.ByteString.Lazy as LB-import Data.Functor+import qualified Data.ByteString          as B import Data.Maybe import qualified Data.Text.Lazy as LT import Data.Time@@ -30,10 +29,7 @@  import Fmt -import Focus- import Network.HTTP.Client (responseStatus)-import Network.HTTP.Date import Network.HTTP.Req import Network.HTTP.Types @@ -42,25 +38,108 @@ import qualified Polysemy.Async as P  import Prelude hiding (error)+import qualified Prelude  import qualified Control.Exception.Safe as Ex import qualified StmContainers.Map as SC  newRateLimitState :: IO RateLimitState-newRateLimitState = RateLimitState <$> SC.newIO <*> E.newSet+newRateLimitState = RateLimitState <$> SC.newIO <*> SC.newIO <*> E.newSet -lookupOrInsertDefaultM :: Monad m => m a -> Focus a m a-lookupOrInsertDefaultM aM =-  casesM-    ( do-        a <- aM-        pure (a, Set a)-    )-    (\a -> pure (a, Leave))+data Ratelimit+  = KnownRatelimit Bucket+  | UnknownRatelimit Route -getRateLimit :: RateLimitState -> Route -> STM Lock-getRateLimit s h = SC.focus (lookupOrInsertDefaultM L.new) h (rateLimits s)+{- performing rate limits+  1. try to fetch the lock from the ratelimiter state+  1.1. if we know the bucket then we retrieve the lock (KnownRatelimit)+  1.2. if we don't know the bucket, we just store the route (UnknownRatelimit)+  2.1. wait fot the global event (the global event is set unless we're globally ratelimited)+  2.2. wait for the local lock (locks are unlocked unless the bucket is exhausted)+  2.3. try to wait for the ratelimit to expire if it's exhausted+  3. we then perform the request, after this we'll know the bucket,+     so we update the ratelimit state with the bucket. if while performing+     the current request another request completed that made the bucket known,+     we don't try to overwrite it+  4.1. if we exceeded the local ratelimit we take the lock and release it after a delay+  4.2. if we exceeded the global ratelimit, we unset the global rl event,+       wait for it to expire, and then try again+-} +getRateLimit :: RateLimitState -> Route -> STM Ratelimit+getRateLimit s h = do+  bucketKey <- SC.lookup h $ bucketKeys s+  bucket <- join <$> sequenceA (flip SC.lookup (buckets s) <$> bucketKey)+  case bucket of+    Just bucket' ->+      pure $ KnownRatelimit bucket'+    Nothing ->+      pure $ UnknownRatelimit h++-- | Knowing the bucket for a route, and the ratelimit info, map the route to+-- the bucket key and retrieve the bucket+updateBucket :: RateLimitState -> Route -> B.ByteString -> BucketState -> STM Bucket+updateBucket s h b bucketState = do+  bucketKey <- SC.lookup h $ bucketKeys s+  case bucketKey of+    Just bucketKey' -> do+      -- if we know the bucket key here, then the bucket has already been made+      bucket <- SC.lookup bucketKey' $ buckets s+      case bucket of+        Just bucket' -> do+          modifyTVar' (bucket' ^. #state) (`mergeBucket` bucketState)+          pure bucket'+        Nothing -> Prelude.error "Not possible"+    Nothing -> do+      -- the bucket key wasn't known, make a new bucket and insert it+      lock <- L.new+      bs <- newTVar bucketState+      let bucket = Bucket lock bs+      SC.insert bucket b $ buckets s+      SC.insert b h $ bucketKeys s+      pure bucket+  where+    mergeBucket :: BucketState -> BucketState -> BucketState+    mergeBucket old new = new { remaining = new ^. #remaining <|> old ^. #remaining }++resetBucket :: Bucket -> STM ()+resetBucket bucket = modifyTVar' (bucket ^. #state) (#remaining .~ Nothing)++-- | Maybe wait for a bucket, updating its state to say we used it+useBucketOnce :: Bucket -> IO ()+useBucketOnce bucket = do+  now <- getCurrentTime+  mWaitUntil <- atomically $ do+    -- first wait on the lock+    L.wait $ bucket ^. #lock++    -- now try to estimate an expiry+    s <- readTVar $ bucket ^. #state++    let remaining = if now > s ^. #resetTime+          then Nothing+          else s ^. #remaining++    case remaining of+      Just n | n > 0 -> do+        modifyTVar (bucket ^. #state) (#remaining ?~ n - 1)+        pure Nothing+      Just _ -> do+        -- expired, lock and then reset the remaining+        L.acquire $ bucket ^. #lock+        modifyTVar (bucket ^. #state) (#remaining .~ Nothing)+        pure . Just $ s ^. #resetTime+      Nothing ->+        -- unknown, just assume we're good+        pure Nothing++  case mWaitUntil of+    Just when -> do+      threadDelayUntil when+      atomically . L.release $ bucket ^. #lock+    Nothing ->+      pure ()+ doDiscordRequest :: BotC r => IO LbsResponse -> Sem r DiscordResponseType doDiscordRequest r = do   r'' <- P.embed $ Ex.catchAny (Right <$> r) (pure . Left . Ex.displayException)@@ -71,18 +150,27 @@           | statusIsSuccessful status -> do             let resp = responseBody r'             debug $ "Got good response from discord: " +|| status ||+ ""-            pure $-              if isExhausted r'-                then case parseRateLimitHeader r' of-                  Just !sleepTime -> ExhaustedBucket resp sleepTime-                  Nothing -> ServerError (statusCode status)-                else Good resp+            now <- P.embed getCurrentTime+            if isExhausted r'+              then case (parseRateLimitHeader now r', buildBucketState now r') of+                (Just !when, Just (!bs, !key)) ->+                  pure $ ExhaustedBucket resp when bs key+                _ ->+                  pure $ ServerError (statusCode status)+              else case buildBucketState now r' of+                Just (!bs, !key) ->+                  pure $ Good resp bs key+                Nothing ->+                  pure $ ServerError (statusCode status)           | status == status429 -> do             debug "Got 429 from discord, retrying."+            now <- P.embed getCurrentTime             let resp = responseBody r'-            case resp ^? _Value of-              Just rv -> pure $ Ratelimited (parseRetryAfter rv) (isGlobal rv)-              Nothing -> pure $ ClientError (statusCode status) "429 with invalid json???"+            case (resp ^? _Value, buildBucketState now r') of+              (Just !rv, bs) ->+                pure $ Ratelimited (parseRetryAfter now rv) (isGlobal rv) bs+              _ ->+                pure $ ServerError (statusCode status)           | statusIsClientError status -> do             let err = responseBody r'             error $ "Something went wrong: " +|| err ||+ " response: " +|| r' ||+ ""@@ -94,26 +182,35 @@       error $ "Something went wrong with the http client: " +| LT.pack e |+ ""       pure . InternalResponseError $ LT.pack e -parseDiscordTime :: ByteString -> Maybe UTCTime-parseDiscordTime s = httpDateToUTC <$> parseHTTPDate s+-- | Parse a ratelimit header returning when it unlocks+parseRateLimitHeader :: HttpResponse r => UTCTime -> r -> Maybe UTCTime+parseRateLimitHeader now r = computedEnd <|> end+  where+    computedEnd :: Maybe UTCTime+    computedEnd = flip addUTCTime now <$> resetAfter -computeDiscordTimeDiff :: Double -> UTCTime -> Int-computeDiscordTimeDiff end !now = round . (* 1000.0) $! diffUTCTime end' now- where-  !end' = end & toRational & fromRational & posixSecondsToUTCTime+    resetAfter :: Maybe NominalDiffTime+    resetAfter = realToFrac <$> responseHeader r "X-Ratelimit-Reset-After" ^? _Just . _Double --- | Parse a ratelimit header returning the number of milliseconds until it resets-parseRateLimitHeader :: HttpResponse r => r -> Maybe Int-parseRateLimitHeader r = computeDiscordTimeDiff <$> end <*> now- where-  end = responseHeader r "X-Ratelimit-Reset" ^? _Just . _Double-  now = parseDiscordTime =<< responseHeader r "Date"+    end :: Maybe UTCTime+    end = posixSecondsToUTCTime . realToFrac <$>+      responseHeader r "X-Ratelimit-Reset" ^? _Just . _Double +buildBucketState :: HttpResponse r => UTCTime -> r -> Maybe (BucketState, B.ByteString)+buildBucketState now r = (bs,) <$> bucketKey+  where+    remaining = responseHeader r "X-RateLimit-Remaining" ^? _Just . _Integral+    bs = BucketState now remaining+    bucketKey = responseHeader r "X-RateLimit-Bucket"+ isExhausted :: HttpResponse r => r -> Bool isExhausted r = responseHeader r "X-RateLimit-Remaining" == Just "0" -parseRetryAfter :: Value -> Int-parseRetryAfter r = r ^?! key "retry_after" . _Integral+-- | Parse the retry after field, returning when to retry+parseRetryAfter :: UTCTime -> Value -> UTCTime+parseRetryAfter now r = addUTCTime retryAfter now+  where+    retryAfter = realToFrac $ r ^?! key "retry_after" . _Double  isGlobal :: Value -> Bool isGlobal r = r ^? key "global" . _Bool == Just True@@ -130,80 +227,126 @@   Int ->   -- | action to perform   Sem r (ShouldRetry a b) ->-  -- | action to run if max number of retries was reached-  Sem r () ->   Sem r (Either a b)-retryRequest max_retries action failAction = retryInner 0+retryRequest maxRetries action = retryInner 0  where-  retryInner num_retries = do+  retryInner numRetries = do     res <- action     case res of-      Retry r | num_retries > max_retries -> do-        debug $ "Request failed after " +| max_retries |+ " retries."-        doFail $ Left r-      Retry _ -> retryInner (num_retries + 1)+      Retry r | numRetries > maxRetries -> do+        debug $ "Request failed after " +| maxRetries |+ " retries."+        pure $ Left r+      Retry _ ->+        retryInner (numRetries + 1)       RFail r -> do         debug "Request failed due to error response."-        doFail $ Left r-      RGood r -> pure $ Right r-   where-    doFail v = failAction $> v+        pure $ Left r+      RGood r ->+        pure $ Right r +threadDelayMS :: Int -> IO ()+threadDelayMS ms = threadDelay (1000 * ms)++threadDelayUntil :: UTCTime -> IO ()+threadDelayUntil when = do+  now <- getCurrentTime+  let msUntil = round . (* 1000) . realToFrac @_ @Double $ diffUTCTime when now+  threadDelayMS msUntil+ -- Run a single request--- NOTE: this function will only unlock the ratelimit lock if the request--- gave a response, otherwise it will stay locked so that it can be retried again doSingleRequest ::   BotC r =>+  RateLimitState ->+  Route ->   -- | Global lock   Event ->-  -- | Local lock-  Lock ->   -- | Request action   IO LbsResponse ->   Sem r (ShouldRetry RestError LB.ByteString)-doSingleRequest gl l r = do+doSingleRequest rlstate route gl r = do+  P.embed $ E.wait (globalLock rlstate)++  rl <- P.embed . atomically $ getRateLimit rlstate route++  case rl of+    KnownRatelimit bucket ->+      P.embed $ useBucketOnce bucket++    _ -> pure ()+   r' <- doDiscordRequest r+   case r' of-    Good v -> do-      P.embed . atomically $ L.release l+    Good v bs bk -> do+      void . P.embed . atomically $ updateBucket rlstate route bk bs       pure $ RGood v-    ExhaustedBucket v d -> do-      debug $ "Exhausted bucket, unlocking after " +| d |+ "ms"++    ExhaustedBucket v unlockWhen bs bk -> do+      debug $ "Exhausted bucket, unlocking at" +| unlockWhen |+ ""++      bucket <- P.embed . atomically $ do+        bucket <- updateBucket rlstate route bk bs+        L.acquire $ bucket ^. #lock+        pure bucket+       void . P.async $ do         P.embed $ do-          threadDelay $ 1000 * d-          atomically $ L.release l+          threadDelayUntil unlockWhen+          atomically $ do+            L.release $ bucket ^. #lock+            resetBucket bucket         debug "unlocking bucket"+       pure $ RGood v-    Ratelimited d False -> do-      debug $ "429 ratelimited on route, sleeping for " +| d |+ " ms"-      P.embed . threadDelay $ 1000 * d++    Ratelimited unlockWhen False (Just (bs, bk)) -> do+      debug $ "429 ratelimited on route, retrying at " +| unlockWhen |+ ""++      bucket <- P.embed . atomically $ do+        bucket <- updateBucket rlstate route bk bs+        L.acquire $ bucket ^. #lock+        pure bucket++      P.embed $ do+        threadDelayUntil unlockWhen+        atomically $ do+          L.release $ bucket ^. #lock+          resetBucket bucket+       pure $ Retry (HTTPError 429 Nothing)-    Ratelimited d True -> do++    Ratelimited unlockWhen False _ -> do+      debug "Internal error (ratelimited but no headers), retrying"+      P.embed $ threadDelayUntil unlockWhen+      pure $ Retry (HTTPError 429 Nothing)++    Ratelimited unlockWhen True bs -> do       debug "429 ratelimited globally"+       P.embed $ do+        case bs of+          Just (bs', bk) ->+            void . atomically $ updateBucket rlstate route bk bs'+          Nothing ->+            pure ()+         E.clear gl-        threadDelay $ 1000 * d+        threadDelayUntil unlockWhen         E.set gl       pure $ Retry (HTTPError 429 Nothing)+     ServerError c -> do       debug "Server failed, retrying"       pure $ Retry (HTTPError c Nothing)+     InternalResponseError c -> do       debug "Internal error, retrying"       pure $ Retry (InternalClientError c)+     ClientError c v -> pure $ RFail (HTTPError c $ decode v)  doRequest :: BotC r => RateLimitState -> Route -> IO LbsResponse -> Sem r (Either RestError LB.ByteString)-doRequest rlState route action = do-  P.embed $ E.wait (globalLock rlState)--  ratelimit <- P.embed . atomically $ do-    lock <- getRateLimit rlState route-    L.acquire lock-    pure lock-+doRequest rlstate route action =   retryRequest     5-    (doSingleRequest (globalLock rlState) ratelimit action)-    (P.embed . atomically $ L.release ratelimit)+    (doSingleRequest rlstate route (globalLock rlstate) action)

@@ -3,6 +3,8 @@     ( RestError(..)     , RateLimitState(..)     , DiscordResponseType(..)+    , Bucket(..)+    , BucketState(..)     , GatewayResponse     , BotGatewayResponse ) where @@ -11,9 +13,12 @@  import           Control.Concurrent.Event      ( Event ) import           Control.Concurrent.STM.Lock   ( Lock )+import           Control.Concurrent.STM.TVar   ( TVar ) +import Data.Time import           Data.Aeson import qualified Data.ByteString.Lazy          as LB+import qualified Data.ByteString          as B import           Data.Text.Lazy  import           GHC.Generics@@ -30,19 +35,49 @@   | InternalClientError Text   deriving ( Show, Generic ) +data BucketState = BucketState+  { resetTime :: UTCTime+    -- ^ The time when the bucket resets, used to heuristically wait out ratelimits+  , remaining  :: Maybe Int+    -- ^ The number of uses left in the bucket, used to heuristically wait out ratelimits+  }+  deriving ( Generic )++data Bucket = Bucket+  { lock :: Lock+  , state :: TVar BucketState+  }+  deriving ( Generic )+ data RateLimitState = RateLimitState-  { rateLimits :: SC.Map Route Lock+  { bucketKeys :: SC.Map Route B.ByteString+  , buckets    :: SC.Map B.ByteString Bucket   , globalLock :: Event   }   deriving ( Generic )  data DiscordResponseType-  = Good LB.ByteString -- ^ A good response-  | ExhaustedBucket -- ^ We got a response but also exhausted the bucket-      LB.ByteString Int -- ^ Retry after (milliseconds)-  | Ratelimited -- ^ We hit a 429, no response and ratelimited-      Int -- ^ Retry after (milliseconds)-      Bool -- ^ Global ratelimit+  = Good+    -- ^ A good response+    LB.ByteString+    BucketState+    B.ByteString+    -- ^ The bucket the route is in+  | ExhaustedBucket+    -- ^ We got a response but also exhausted the bucket+    LB.ByteString+    UTCTime+    -- ^ Retry after+    BucketState+    B.ByteString+    -- ^ The bucket the route is in+  | Ratelimited+    -- ^ We hit a 429, no response and ratelimited+    UTCTime+    -- ^ Retry after+    Bool+    -- ^ Global ratelimit+    (Maybe (BucketState, B.ByteString))   | ServerError Int -- ^ Discord's error, we should retry (HTTP 5XX)   | ClientError Int LB.ByteString -- ^ Our error, we should fail   | InternalResponseError Text -- ^ Something went wrong with the http client

@@ -1,5 +1,10 @@ # Changelog for Calamity +## 0.1.28.1+++ Rework the ratelimit implementation to use X-Ratelimit-Bucket++ Fix incorrect interpretation of the retry-after for 429s+ ## 0.1.28.0  + Added support for message types 19 (reply) and 20 (application command)

@@ -4,10 +4,10 @@ -- -- see: https://github.com/sol/hpack ----- hash: 3c5b52b22636e00ca4b9cc43f1349fce19deee44cd599e47312642322256f9ab+-- hash: 2482ee40637f3fe615af0564b33ebb2dcbb4e67d1ec98d152f59206b977c3686  name:           calamity-version:        0.1.28.0+version:        0.1.28.1 synopsis:       A library for writing discord bots in haskell description:    Please see the README on GitHub at <https://github.com/simmsb/calamity#readme> category:       Network, Web