Fix supervisor and proper rate limiting

src/HomeAssistant/Runtime/RateLimit.hs

@@ -0,0 +1,74 @@
+{-# 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 50000 <> limitRetries 5
+    incRetry = liftIO $ Counter.inc (rlRetries metrics)
+    incReject = liftIO $ Counter.inc (rlRejects metrics)