NumerusClosus.hs

{-# LANGUAGE StrictData #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE NoFieldSelectors #-}

-- |
-- Module        : Control.NumerusClosus
-- Copyright     : Gautier DI FOLCO
-- License       : ISC
--
-- Maintainer    : Gautier DI FOLCO <gautier.difolco@gmail.com>
-- Stability     : Stable
-- Portability   : Portable
--
-- Simple, composable, pure rate-limiting primitives supporting finite buckets,
-- fixed windows, sliding windows, and sliding window counters. Includes AND/OR
-- combinators and scheduling helpers.
--
-- This module wraps the typed rate-limiter representations from
-- "Control.NumerusClosus.Typed" behind an existential 'RateLimiter' data type,
-- which is itself an instance of the 'Typed.RateLimiter' class.
--
-- > import Control.NumerusClosus
-- > import Data.Time (getCurrentTime)
-- >
-- > main :: IO ()
-- > main = do
-- >   now <- getCurrentTime
-- >   let rl = fixedWindow 60 5 now  -- 5 requests per 60 seconds
-- >   result <- schedule rl (putStrLn "Request allowed")
-- >   print result
module Control.NumerusClosus
  ( RateLimiter (..),
    debit,

    -- * Re-exported data types
    Typed.NextDebitable (..),
    Typed.BucketSize (..),
    Typed.WindowSize (..),
    Typed.BucketsCount (..),

    -- * Base helpers
    alwaysAllow,
    alwaysDeny,

    -- * Strategies
    finiteBucket,
    fixedWindow,
    slidingWindow,
    slidingWindowCount,
    slidingWindowBucketed,

    -- * Combinators
    (.&&),
    (.||),
    allOf,
    anyOf,

    -- * Scheduling
    Typed.PositionInTime (..),
    Typed.ioPositionInTime,
    schedule,
    scheduleWith,
    loopSchedule,
    loopScheduleWith,
  )
where

import qualified Control.NumerusClosus.Typed as Typed
import qualified Data.List.NonEmpty as NE
import Data.Time (UTCTime (..))

-- * Main logic

-- | The core rate limiter type.
-- An existential wrapper around any typed rate limiter, providing a uniform
-- interface while delegating to the typed implementations from
-- "Control.NumerusClosus.Typed".
data RateLimiter = forall a. (Typed.RateLimiter a) => RateLimiter {wrapped :: a}

instance Typed.RateLimiter RateLimiter where
  debit (RateLimiter rl) t =
    case Typed.debit rl t of
      Right rl' -> Right $ RateLimiter {wrapped = rl'}
      Left nd -> Left nd

-- | Determine if a request is allowed at a given time.
-- Returns either the next debitable time or a new rate limiter state.
debit :: RateLimiter -> UTCTime -> Either Typed.NextDebitable RateLimiter
debit = Typed.debit

-- * Base helpers

-- | A rate limiter that always allows requests.
alwaysAllow :: RateLimiter
alwaysAllow = RateLimiter {wrapped = Typed.alwaysAllow}

-- | A rate limiter that always denies requests.
alwaysDeny :: RateLimiter
alwaysDeny = RateLimiter {wrapped = Typed.alwaysDeny}

-- * Strategies

-- | A simple finite bucket that allows exactly @n@ requests.
finiteBucket :: Typed.BucketSize Integer -> RateLimiter
finiteBucket = RateLimiter . Typed.finiteBucket

-- | A fixed window rate limiter that allows a given number of requests per window.
fixedWindow :: Typed.WindowSize -> Typed.BucketSize Integer -> UTCTime -> RateLimiter
fixedWindow ws bs = RateLimiter . Typed.fixedWindow ws bs

-- | A sliding window rate limiter that allows a given number of requests within the sliding window.
slidingWindow :: Typed.WindowSize -> Typed.BucketSize Int -> RateLimiter
slidingWindow ws bs = RateLimiter $ Typed.slidingWindow ws bs

-- | Sliding window counter using two-window weighted interpolation.
slidingWindowCount :: Typed.WindowSize -> Typed.BucketSize Integer -> UTCTime -> RateLimiter
slidingWindowCount ws (Typed.BucketSize bs) = RateLimiter . Typed.slidingWindowCount ws (Typed.BucketSize $ fromIntegral bs)

-- | Sliding window rate limiter using sub-bucket counters.
slidingWindowBucketed :: Typed.WindowSize -> Typed.BucketsCount Int -> Typed.BucketSize Integer -> RateLimiter
slidingWindowBucketed ws bc bs = RateLimiter $ Typed.slidingWindowBucketed ws bc bs

-- * Combinators

-- | AND combinator: allows a request if both rate limiters allow it.
(.&&) :: RateLimiter -> RateLimiter -> RateLimiter
x .&& y =
  RateLimiter
    { wrapped = x Typed.:&& y
    }

infixr 3 .&&

-- | OR combinator: allows a request if either rate limiter allows it.
(.||) :: RateLimiter -> RateLimiter -> RateLimiter
x .|| y =
  RateLimiter
    { wrapped = x Typed.:|| y
    }

infixr 3 .||

-- | Require all of the given rate limiters to allow the request.
allOf :: NE.NonEmpty RateLimiter -> RateLimiter
allOf = foldl1 (.&&)

-- | Allow the request if any of the given rate limiters allow it.
anyOf :: NE.NonEmpty RateLimiter -> RateLimiter
anyOf = foldl1 (.||)

-- * Scheduling

-- | Run an IO action if the rate limiter allows it, using the current time.
schedule :: RateLimiter -> IO a -> IO (Either Typed.NextDebitable (RateLimiter, a))
schedule = scheduleWith Typed.ioPositionInTime

-- | Run a monadic action if the rate limiter allows it, using a custom time strategy.
scheduleWith :: (Monad m) => Typed.PositionInTime m -> RateLimiter -> m a -> m (Either Typed.NextDebitable (RateLimiter, a))
scheduleWith pit rl action = do
  now <- pit.getTime
  case debit rl now of
    Right rl' -> Right . (rl',) <$> action
    Left nd -> return $ Left nd

-- | Repeatedly run an IO action, respecting the rate limiter. Will sleep until the next available slot if rate limited.
loopSchedule :: RateLimiter -> IO a -> IO ()
loopSchedule = loopScheduleWith Typed.ioPositionInTime

-- | Repeatedly run a monadic action, respecting the rate limiter and using a custom time strategy. Will sleep until the next available slot if rate limited.
loopScheduleWith :: (Monad m) => Typed.PositionInTime m -> RateLimiter -> m a -> m ()
loopScheduleWith pit rl action = go rl
  where
    go rl' = do
      result <- scheduleWith pit rl' action
      case result of
        Right (rl'', _) -> go rl''
        Left Typed.Never -> return ()
        Left (Typed.DebitableFrom at) -> pit.delayUntil at >> go rl'