75 lines
2.7 KiB
Haskell
75 lines
2.7 KiB
Haskell
{-# Language DerivingVia #-}
|
|
{-# Language OverloadedStrings #-}
|
|
module HomeAssistant.Runtime.RateLimit where
|
|
import Control.Monad.IO.Class (MonadIO, liftIO)
|
|
import Data.Time (NominalDiffTime)
|
|
import Control.Concurrent.STM (newTVarIO, atomically, readTVar, writeTVar)
|
|
import Data.Time.Clock.POSIX (getPOSIXTime)
|
|
import Katip.Monadic (KatipContext(..))
|
|
import Control.Retry (fullJitterBackoff, limitRetries, retrying, RetryStatus (..))
|
|
import Data.Maybe (isNothing)
|
|
import Control.Exception.Annotated (Exception, checkpointCallStack, throw)
|
|
import Control.Monad.Catch (MonadCatch)
|
|
import Katip (Katip, Severity (WarningS), logFM, ls)
|
|
import Control.Monad (when)
|
|
import System.Metrics.Counter (Counter)
|
|
import qualified System.Metrics.Counter as Counter
|
|
import System.Metrics (Store, createCounter)
|
|
|
|
data RateLimitMetrics = RateLimitMetrics
|
|
{ rlRetries :: Counter
|
|
, rlRejects :: Counter
|
|
}
|
|
|
|
|
|
registerRateLimitMetrics :: MonadIO m => Store -> m RateLimitMetrics
|
|
registerRateLimitMetrics store = liftIO $ RateLimitMetrics
|
|
<$> createCounter "ratelimit.retry" store
|
|
<*> createCounter "ratelimit.exhausted" store
|
|
|
|
newtype RateLimited m a = RateLimited (m a)
|
|
deriving (Functor, Applicative, Monad, MonadIO, KatipContext, Katip) via m
|
|
|
|
|
|
data RateLimiter = RateLimiter
|
|
{ metrics :: RateLimitMetrics
|
|
, limiter :: forall m a. (KatipContext m, MonadIO m) => RateLimited m a -> m (Maybe a)
|
|
}
|
|
|
|
|
|
slidingWindowLimiter
|
|
:: RateLimitMetrics
|
|
-> NominalDiffTime -- ^ Within how many seconds
|
|
-> Int -- ^ How many requests we can make
|
|
-> IO RateLimiter
|
|
slidingWindowLimiter metrics timeWindow size = do
|
|
window <- newTVarIO []
|
|
return $ RateLimiter metrics $ \(RateLimited action) -> do
|
|
now <- liftIO getPOSIXTime
|
|
allowed <- liftIO $ atomically $ do
|
|
elements <- takeWhile (\x -> x > now - timeWindow) <$> readTVar window
|
|
if length elements >= size
|
|
then pure False
|
|
else True <$ writeTVar window (now : elements)
|
|
if allowed
|
|
then Just <$> action
|
|
else pure Nothing
|
|
|
|
data RateLimitException = RateLimitExhausted
|
|
deriving Show
|
|
|
|
instance Exception RateLimitException
|
|
|
|
runRateLimited :: (MonadIO m, KatipContext m, MonadCatch m) => RateLimiter -> RateLimited m a -> m a
|
|
runRateLimited (RateLimiter metrics run) action = checkpointCallStack $ do
|
|
x <- retrying retryPolicy (const $ pure . isNothing) $ \RetryStatus{rsIterNumber} -> do
|
|
when (rsIterNumber > 0) $ do
|
|
logFM WarningS $ ls $ "Rate limited, attempt no. " <> show rsIterNumber
|
|
incRetry
|
|
run action
|
|
maybe (incReject >> throw RateLimitExhausted) pure x
|
|
where
|
|
retryPolicy = fullJitterBackoff 500_000 <> limitRetries 5
|
|
incRetry = liftIO $ Counter.inc (rlRetries metrics)
|
|
incReject = liftIO $ Counter.inc (rlRejects metrics)
|