NumerusClosus.hs

{-# LANGUAGE TupleSections #-}

-- |
-- 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.
--
-- > 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 (..),
    NextDebitable (..),

    -- * Base helpers
    alwaysAllow,
    alwaysDeny,

    -- * Strategies
    BucketSize (..),
    finiteBucket,
    fixedWindow,
    slidingWindow,
    slidingWindowCount,
    WindowsCount (..),
    slidingWindowBucketed,

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

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

import Control.Concurrent.Thread.Delay (delay)
import Control.Monad (mfilter)
import qualified Data.List.NonEmpty as NE
import qualified Data.Map.Strict as Map
import Data.Maybe (fromMaybe)
import qualified Data.Set as Set
import Data.Time (FormatTime, NominalDiffTime, ParseTime, UTCTime, addUTCTime, diffUTCTime, getCurrentTime)

-- * Main logic

-- | The core rate limiter type.
-- Given the current time, returns either the next debitable time or a new rate limiter state.
newtype RateLimiter = RateLimiter
  { -- | Function to determine if a request is allowed at a given time.
    debit :: UTCTime -> Either NextDebitable RateLimiter
  }

-- | Represents when the next request can be debited.
data NextDebitable
  = -- | The request will never be debitable again.
    Never
  | -- | The request can be debited from this time onwards.
    DebitableFrom UTCTime
  deriving stock (Eq, Show)

instance Semigroup NextDebitable where
  Never <> _ = Never
  _ <> Never = Never
  DebitableFrom x <> DebitableFrom y = DebitableFrom $ max x y

-- * Base helpers

-- | A rate limiter that always allows requests.
alwaysAllow :: RateLimiter
alwaysAllow =
  RateLimiter
    { debit = const $ Right alwaysAllow
    }

-- | A rate limiter that always denies requests.
alwaysDeny :: RateLimiter
alwaysDeny =
  RateLimiter
    { debit = const $ Left Never
    }

-- * Strategies

-- | The maximum number of requests a bucket can hold.
newtype BucketSize a
  = BucketSize a
  deriving stock (Eq, Ord, Show)
  deriving newtype (Num)

-- | A simple finite bucket that allows exactly @n@ requests.
finiteBucket :: BucketSize Integer -> RateLimiter
finiteBucket (BucketSize n) =
  RateLimiter
    { debit =
        const $
          if n > 0
            then Right $ finiteBucket $ BucketSize (n - 1)
            else Left Never
    }

-- | The size of a time window.
newtype WindowSize = WindowSize NominalDiffTime
  deriving stock (Eq, Ord, Show)
  deriving newtype (Num, Fractional, Real, RealFrac, FormatTime, ParseTime)

-- | A fixed window rate limiter that allows a given number of requests per window.
fixedWindow :: WindowSize -> BucketSize Integer -> UTCTime -> RateLimiter
fixedWindow (WindowSize window) (BucketSize maxBucket) startTime = go maxBucket $ addUTCTime window startTime
  where
    go bucket endTime =
      RateLimiter
        { debit =
            \now ->
              let (refreshedBucket, refreshedEndTime) =
                    if now > endTime
                      then (maxBucket, addUTCTime (fromIntegral (floor (diffUTCTime now endTime / window) + 1 :: Integer) * window) endTime)
                      else (bucket, endTime)
               in if refreshedBucket > 0
                    then Right $ go (refreshedBucket - 1) refreshedEndTime
                    else Left $ DebitableFrom $ nextTimeUnit refreshedEndTime
        }

-- | A sliding window rate limiter that allows a given number of requests within the sliding window.
slidingWindow :: WindowSize -> BucketSize Int -> RateLimiter
slidingWindow (WindowSize window) (BucketSize maxBucket) = go Set.empty
  where
    go bucket =
      RateLimiter
        { debit =
            \now ->
              let refreshedBucket = snd $ Set.split (addUTCTime ((-1) * window) now) bucket
               in if Set.size refreshedBucket < maxBucket
                    then Right $ go $ Set.insert now refreshedBucket
                    else Left $ DebitableFrom $ nextTimeUnit $ addUTCTime window $ fromMaybe now $ Set.lookupMin refreshedBucket
        }

-- | The number of sub-windows for bucketed sliding windows.
newtype WindowsCount a
  = WindowsCount a
  deriving stock (Eq, Ord, Show)
  deriving newtype (Num)

-- | Sliding window rate limiter using sub-bucket counters.
-- Divides the window into sub-buckets keyed by request arrival times and sums
-- their counters. This is a bucketed variant rather than the canonical
-- two-window weighted interpolation.
slidingWindowBucketed :: WindowSize -> WindowsCount Int -> BucketSize Integer -> RateLimiter
slidingWindowBucketed (WindowSize window) (WindowsCount windowsCount) (BucketSize maxBucket) = go Map.empty
  where
    bucketWindow = window / fromIntegral windowsCount
    go buckets =
      RateLimiter
        { debit =
            \now ->
              let refreshedBucket =
                    Map.restrictKeys buckets $
                      snd $
                        Set.split (addUTCTime ((-1) * window) now) $
                          Map.keysSet buckets
                  lastBucket =
                    fromMaybe now $
                      mfilter (> addUTCTime ((-1) * bucketWindow) now) $
                        fst <$> Map.lookupMax refreshedBucket
               in if sum refreshedBucket < maxBucket
                    then Right $ go $ Map.alter (Just . maybe 1 (+ 1)) lastBucket refreshedBucket
                    else Left $ DebitableFrom $ nextTimeUnit $ addUTCTime window $ maybe now fst $ Map.lookupMin refreshedBucket
        }

-- | Sliding window counter using two-window weighted interpolation.
-- Tracks counters for the current and previous fixed windows, then estimates
-- the request rate as: @prev * (1 - elapsed\/window) + current@.
-- This is the canonical sliding window counter algorithm with O(1) memory.
slidingWindowCount :: WindowSize -> BucketSize Integer -> UTCTime -> RateLimiter
slidingWindowCount (WindowSize window) (BucketSize maxBucket) startTime = go (0 :: Integer) 0 $ addUTCTime window startTime
  where
    go prevCount currentCount windowEnd =
      RateLimiter
        { debit =
            \now ->
              let (refreshedPrev, refreshedCurrent, refreshedEnd) =
                    if now > windowEnd
                      then
                        let windowsElapsed = floor (diffUTCTime now windowEnd / window) :: Integer
                            newEnd = addUTCTime (fromIntegral (windowsElapsed + 1) * window) windowEnd
                         in if windowsElapsed == 0
                              then (currentCount, 0, newEnd)
                              else (0, 0, newEnd)
                      else (prevCount, currentCount, windowEnd)
                  windowStart = addUTCTime ((-1) * window) refreshedEnd
                  fraction = realToFrac (diffUTCTime now windowStart) / realToFrac window :: Double
                  estimate = fromIntegral refreshedPrev * (1 - fraction) + fromIntegral refreshedCurrent :: Double
               in if estimate < fromIntegral maxBucket
                    then Right $ go refreshedPrev (refreshedCurrent + 1) refreshedEnd
                    else Left $ DebitableFrom $ nextTimeUnit refreshedEnd
        }

-- | Helper to get the smallest next time unit for debiting.
nextTimeUnit :: UTCTime -> UTCTime
nextTimeUnit = addUTCTime 0.000001

-- * Combinators

-- | AND combinator: allows a request if both rate limiters allow it.
(.&&) :: RateLimiter -> RateLimiter -> RateLimiter
x .&& y =
  RateLimiter
    { debit = \at ->
        case (x.debit at, y.debit at) of
          (Right x', Right y') -> Right $ x' .&& y'
          (Left x', Left y') -> Left $ x' <> y'
          (Left x', _) -> Left x'
          (_, Left y') -> Left y'
    }

infixr 3 .&&

-- | OR combinator: allows a request if either rate limiter allows it.
(.||) :: RateLimiter -> RateLimiter -> RateLimiter
x .|| y =
  RateLimiter
    { debit = \at ->
        case (x.debit at, y.debit at) of
          (Right x', Right y') -> Right $ x' .|| y'
          (Left x', Left y') ->
            Left $
              case (x', y') of
                (Never, Never) -> Never
                (DebitableFrom x'', DebitableFrom y'') -> DebitableFrom $ min x'' y''
                (DebitableFrom x'', _) -> DebitableFrom x''
                (_, DebitableFrom y'') -> DebitableFrom y''
          (Right x', _) -> Right x'
          (_, Right y') -> Right 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

-- | Fetch the time and compute the delay time.
data PositionInTime m = PositionInTime
  { -- | How to fetch the current time.
    getTime :: m UTCTime,
    -- | How to delay until a specific time.
    delayUntil :: UTCTime -> m ()
  }

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

-- | Default time strategy for IO.
ioPositionInTime :: PositionInTime IO
ioPositionInTime = PositionInTime getCurrentTime diffDelay
  where
    diffDelay to = do
      now <- getCurrentTime
      delay $ 0 `max` round (diffUTCTime to now * 1000000)

-- | Run a monadic action if the rate limiter allows it, using a custom time strategy.
scheduleWith :: (Monad m) => PositionInTime m -> RateLimiter -> m a -> m (Either NextDebitable (RateLimiter, a))
scheduleWith pit rl action = do
  now <- pit.getTime
  case rl.debit 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 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) => 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 Never -> return ()
        Left (DebitableFrom at) -> pit.delayUntil at >> go rl'