diff --git a/numerus-closus.cabal b/numerus-closus.cabal --- a/numerus-closus.cabal +++ b/numerus-closus.cabal @@ -1,6 +1,6 @@ cabal-version: 3.0 name: numerus-closus -version: 0.1.0.0 +version: 0.2.0.0 author: Gautier DI FOLCO maintainer: gautier.difolco@gmail.com category: Scheduling @@ -22,6 +22,7 @@ hs-source-dirs: src exposed-modules: Control.NumerusClosus + Control.NumerusClosus.Typed other-modules: Paths_numerus_closus autogen-modules: diff --git a/src/Control/NumerusClosus.hs b/src/Control/NumerusClosus.hs --- a/src/Control/NumerusClosus.hs +++ b/src/Control/NumerusClosus.hs @@ -1,4 +1,6 @@ +{-# LANGUAGE StrictData #-} {-# LANGUAGE TupleSections #-} +{-# LANGUAGE NoFieldSelectors #-} -- | -- Module : Control.NumerusClosus @@ -13,6 +15,10 @@ -- 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) -- > @@ -24,19 +30,23 @@ -- > print result module Control.NumerusClosus ( RateLimiter (..), - NextDebitable (..), + debit, + -- * Re-exported data types + Typed.NextDebitable (..), + Typed.BucketSize (..), + Typed.WindowSize (..), + Typed.BucketsCount (..), + -- * Base helpers alwaysAllow, alwaysDeny, -- * Strategies - BucketSize (..), finiteBucket, fixedWindow, slidingWindow, slidingWindowCount, - WindowsCount (..), slidingWindowBucketed, -- * Combinators @@ -46,8 +56,8 @@ anyOf, -- * Scheduling - PositionInTime (..), - ioPositionInTime, + Typed.PositionInTime (..), + Typed.ioPositionInTime, schedule, scheduleWith, loopSchedule, @@ -55,170 +65,62 @@ ) where -import Control.Concurrent.Thread.Delay (delay) -import Control.Monad (mfilter) +import qualified Control.NumerusClosus.Typed as Typed 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) +import Data.Time (UTCTime (..)) -- * 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 - } +-- 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} --- | 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 Typed.RateLimiter RateLimiter where + debit (RateLimiter rl) t = + case Typed.debit rl t of + Right rl' -> Right $ RateLimiter {wrapped = rl'} + Left nd -> Left nd -instance Semigroup NextDebitable where - Never <> _ = Never - _ <> Never = Never - DebitableFrom x <> DebitableFrom y = DebitableFrom $ max x y +-- | 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 - { debit = const $ Right alwaysAllow - } +alwaysAllow = RateLimiter {wrapped = Typed.alwaysAllow} -- | A rate limiter that always denies requests. alwaysDeny :: RateLimiter -alwaysDeny = - RateLimiter - { debit = const $ Left Never - } +alwaysDeny = RateLimiter {wrapped = Typed.alwaysDeny} -- * 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) +finiteBucket :: Typed.BucketSize Integer -> RateLimiter +finiteBucket = RateLimiter . Typed.finiteBucket -- | 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 - } +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 :: 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 +slidingWindow :: Typed.WindowSize -> Typed.BucketSize Int -> RateLimiter +slidingWindow ws bs = RateLimiter $ Typed.slidingWindow ws bs epoch 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 - } + epoch = UTCTime (toEnum 0) 0 -- | 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 - } +slidingWindowCount :: Typed.WindowSize -> Typed.BucketSize Integer -> UTCTime -> RateLimiter +slidingWindowCount ws (Typed.BucketSize bs) = RateLimiter . Typed.slidingWindowCount ws (Typed.BucketSize $ fromIntegral bs) --- | Helper to get the smallest next time unit for debiting. -nextTimeUnit :: UTCTime -> UTCTime -nextTimeUnit = addUTCTime 0.000001 +-- | 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 @@ -226,12 +128,7 @@ (.&&) :: 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' + { wrapped = x Typed.:&& y } infixr 3 .&& @@ -240,18 +137,7 @@ (.||) :: 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' + { wrapped = x Typed.:|| y } infixr 3 .|| @@ -266,45 +152,29 @@ -- * 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) +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) => PositionInTime m -> RateLimiter -> m a -> m (Either NextDebitable (RateLimiter, a)) +scheduleWith :: (Monad m) => Typed.PositionInTime m -> RateLimiter -> m a -> m (Either Typed.NextDebitable (RateLimiter, a)) scheduleWith pit rl action = do now <- pit.getTime - case rl.debit now of + 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 ioPositionInTime +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) => PositionInTime m -> RateLimiter -> m a -> m () +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 Never -> return () - Left (DebitableFrom at) -> pit.delayUntil at >> go rl' + Left Typed.Never -> return () + Left (Typed.DebitableFrom at) -> pit.delayUntil at >> go rl' diff --git a/src/Control/NumerusClosus/Typed.hs b/src/Control/NumerusClosus/Typed.hs new file mode 100644 --- /dev/null +++ b/src/Control/NumerusClosus/Typed.hs @@ -0,0 +1,360 @@ +{-# LANGUAGE OverloadedRecordDot #-} +{-# LANGUAGE StrictData #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE NoFieldSelectors #-} + +-- | +-- Module : Control.NumerusClosus.Typed +-- Copyright : Gautier DI FOLCO +-- License : ISC +-- +-- Maintainer : Gautier DI FOLCO +-- 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.Typed + ( RateLimiter (..), + NextDebitable (..), + + -- * Base helpers + AlwaysAllow (..), + alwaysAllow, + AlwaysDeny (..), + alwaysDeny, + + -- * Strategies + BucketSize (..), + FiniteBucket (..), + finiteBucket, + WindowSize (..), + FixedWindow (..), + fixedWindow, + SlidingWindow (..), + slidingWindow, + SlidingWindowCount (..), + slidingWindowCount, + BucketsCount (..), + SlidingWindowBucketed (..), + slidingWindowBucketed, + + -- * Combinators + type (:&&) (..), + type (:||) (..), + + -- * Scheduling + PositionInTime (..), + ioPositionInTime, + schedule, + scheduleWith, + loopSchedule, + loopScheduleWith, + ) +where + +import Control.Concurrent.Thread.Delay (delay) +import Control.Monad (mfilter) +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) +import GHC.Generics (Generic) + +-- * Main logic + +-- | The core rate limiter type. +-- Given the current time, returns either the next debitable time or a new rate limiter state. +class RateLimiter a where + -- | Function to determine if a request is allowed at a given time. + debit :: a -> UTCTime -> Either NextDebitable a + +-- | 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. +data AlwaysAllow = AlwaysAllow + deriving stock (Eq, Show, Generic) + +instance RateLimiter AlwaysAllow where + debit AlwaysAllow _ = Right AlwaysAllow + +-- | Smart constructor for 'AlwaysAllow'. +alwaysAllow :: AlwaysAllow +alwaysAllow = AlwaysAllow + +-- | A rate limiter that always denies requests. +data AlwaysDeny = AlwaysDeny + deriving stock (Eq, Show, Generic) + +instance RateLimiter AlwaysDeny where + debit AlwaysDeny _ = Left Never + +-- | Smart constructor for 'AlwaysDeny'. +alwaysDeny :: AlwaysDeny +alwaysDeny = AlwaysDeny + +-- * Strategies + +-- | The maximum number of requests a bucket can hold. +newtype BucketSize a = BucketSize + { unBucketSize :: a + } + deriving stock (Eq, Ord, Show) + deriving newtype (Num) + +-- | A simple finite bucket that allows exactly @n@ requests. +newtype FiniteBucket + = FiniteBucket {fullBucket :: BucketSize Integer} + deriving stock (Eq, Show, Generic) + +instance RateLimiter FiniteBucket where + debit (FiniteBucket (BucketSize n)) _ = + if n > 0 + then Right $ FiniteBucket $ BucketSize (n - 1) + else Left Never + +finiteBucket :: BucketSize Integer -> FiniteBucket +finiteBucket fullBucket = FiniteBucket {..} + +-- | The size of a time window. +newtype WindowSize = WindowSize + { unWindowSize :: 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. +data FixedWindow = FixedWindow + { windowSize :: WindowSize, + maxBucket :: BucketSize Integer, + -- \* Iteration fields + bucket :: BucketSize Integer, + windowStart :: UTCTime, + windowEnd :: UTCTime + } + deriving stock (Eq, Show, Generic) + +instance RateLimiter FixedWindow where + debit (FixedWindow {..}) now = + if refreshedBucket > 0 + then Right $ FixedWindow {bucket = refreshedBucket - 1, windowEnd = refreshedEndTime, ..} + else Left $ DebitableFrom $ nextTimeUnit refreshedEndTime + where + (refreshedBucket, refreshedEndTime) = + if now > windowEnd + then (maxBucket, addUTCTime (fromIntegral (floor (diffUTCTime now windowEnd / windowSize.unWindowSize) + 1 :: Integer) * windowSize.unWindowSize) windowEnd) + else (bucket, windowEnd) + +fixedWindow :: WindowSize -> BucketSize Integer -> UTCTime -> FixedWindow +fixedWindow windowSize@(WindowSize window) maxBucket windowStart = + FixedWindow {bucket = maxBucket, windowEnd = addUTCTime window windowStart, ..} + +-- | A sliding window rate limiter that allows a given number of requests within the sliding window. +data SlidingWindow = SlidingWindow + { windowSize :: WindowSize, + maxBucket :: BucketSize Int, + -- \* Iteration field + windowStart :: UTCTime, + bucket :: Set.Set UTCTime + } + deriving stock (Eq, Show, Generic) + +instance RateLimiter SlidingWindow where + debit SlidingWindow {..} now = + if Set.size refreshedBucket < maxBucket.unBucketSize + then Right $ SlidingWindow {bucket = Set.insert now refreshedBucket, ..} + else Left $ DebitableFrom $ nextTimeUnit $ addUTCTime windowSize.unWindowSize $ fromMaybe now $ Set.lookupMin refreshedBucket + where + refreshedBucket = snd $ Set.split (addUTCTime ((-1) * windowSize.unWindowSize) now) bucket + +slidingWindow :: WindowSize -> BucketSize Int -> UTCTime -> SlidingWindow +slidingWindow windowSize maxBucket windowStart = + SlidingWindow {bucket = Set.empty, ..} + +-- | The number of sub-windows for bucketed sliding windows. +newtype BucketsCount a = BucketsCount + { unBucketsCount :: 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. +data SlidingWindowBucketed = SlidingWindowBucketed + { windowSize :: WindowSize, + bucketsCount :: BucketsCount Int, + maxBucket :: BucketSize Integer, + -- \* Iteration field + buckets :: Map.Map UTCTime Integer + } + deriving stock (Eq, Show, Generic) + +instance RateLimiter SlidingWindowBucketed where + debit SlidingWindowBucketed {..} now = + if sum refreshedBucket < maxBucket.unBucketSize + then Right $ SlidingWindowBucketed {buckets = Map.alter (Just . maybe 1 (+ 1)) lastBucket refreshedBucket, ..} + else Left $ DebitableFrom $ nextTimeUnit $ addUTCTime windowSize.unWindowSize $ maybe now fst $ Map.lookupMin refreshedBucket + where + refreshedBucket = + Map.restrictKeys buckets $ + snd $ + Set.split (addUTCTime ((-1) * windowSize.unWindowSize) now) $ + Map.keysSet buckets + lastBucket = + fromMaybe now $ + mfilter (> addUTCTime ((-1) * bucketWindow) now) $ + fst <$> Map.lookupMax refreshedBucket + bucketWindow = windowSize.unWindowSize / fromIntegral bucketsCount.unBucketsCount + +slidingWindowBucketed :: WindowSize -> BucketsCount Int -> BucketSize Integer -> SlidingWindowBucketed +slidingWindowBucketed windowSize bucketsCount maxBucket = + SlidingWindowBucketed {buckets = Map.empty, ..} + +-- | 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. +data SlidingWindowCount = SlidingWindowCount + { windowSize :: WindowSize, + maxBucket :: BucketSize Int, + -- \* Iteration field + prevCount :: Integer, + currentCount :: Integer, + windowStart :: UTCTime, + windowEnd :: UTCTime + } + deriving stock (Eq, Show, Generic) + +instance RateLimiter SlidingWindowCount where + debit SlidingWindowCount {..} now = + if estimate < fromIntegral maxBucket.unBucketSize + then Right $ SlidingWindowCount {prevCount = refreshedPrev, currentCount = refreshedCurrent + 1, windowEnd = refreshedEnd, ..} + else Left $ DebitableFrom $ nextTimeUnit refreshedEnd + where + (refreshedPrev, refreshedCurrent, refreshedEnd) = + if now > windowEnd + then + let windowsElapsed = floor (diffUTCTime now windowEnd / windowSize.unWindowSize) :: Integer + newEnd = addUTCTime (fromIntegral (windowsElapsed + 1) * windowSize.unWindowSize) windowEnd + in if windowsElapsed == 0 + then (currentCount, 0, newEnd) + else (0, 0, newEnd) + else (prevCount, currentCount, windowEnd) + newWindowStart = addUTCTime ((-1) * windowSize.unWindowSize) refreshedEnd + fraction = realToFrac (diffUTCTime now newWindowStart) / realToFrac windowSize.unWindowSize :: Double + estimate = fromIntegral refreshedPrev * (1 - fraction) + fromIntegral refreshedCurrent :: Double + +slidingWindowCount :: WindowSize -> BucketSize Int -> UTCTime -> SlidingWindowCount +slidingWindowCount windowSize@(WindowSize window) maxBucket windowStart = + SlidingWindowCount {prevCount = 0, currentCount = 0, windowEnd = addUTCTime window windowStart, ..} + +-- | Helper to get the smallest next time unit for debiting. +nextTimeUnit :: UTCTime -> UTCTime +nextTimeUnit = addUTCTime 0.000001 + +-- * Combinators + +-- | AND combinator type: allows a request if both rate limiters allow it. +data a :&& b = a :&& b + deriving stock (Eq, Show, Generic) + +infixr 3 :&& + +instance (RateLimiter a, RateLimiter b) => RateLimiter (a :&& b) where + debit (x :&& y) at = + case (debit x at, debit y at) of + (Right x', Right y') -> Right $ x' :&& y' + (Left x', Left y') -> Left $ x' <> y' + (Left x', _) -> Left x' + (_, Left y') -> Left y' + +-- | OR combinator type: allows a request if either rate limiter allows it. +data a :|| b = a :|| b + deriving stock (Eq, Show, Generic) + +infixr 3 :|| + +instance (RateLimiter a, RateLimiter b) => RateLimiter (a :|| b) where + debit (x :|| y) at = + case (debit x at, debit y 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' :|| y + (_, Right y') -> Right $ x :|| y' + +-- * 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 a) => a -> (a -> IO b) -> IO (Either NextDebitable (a, b)) +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, RateLimiter a) => PositionInTime m -> a -> (a -> m b) -> m (Either NextDebitable (a, b)) +scheduleWith pit rl action = do + now <- pit.getTime + case debit rl now of + Right rl' -> Right . (rl',) <$> action rl' + 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 a) => a -> (a -> IO b) -> 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, RateLimiter a) => PositionInTime m -> a -> (a -> m b) -> 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'