numerus-closus 0.4.0.1 → 0.4.1.0

raw patch · 4 files changed

+231/−3 lines, 4 filesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

API changes (from Hackage documentation)

Files

• numerus-closus.cabal +1/−1
• src/Control/NumerusClosus.hs +21/−1
• src/Control/NumerusClosus/Typed.hs +114/−0
• test/Spec.hs +95/−1

@@ -1,6 +1,6 @@ cabal-version:       3.0 name:                numerus-closus-version:             0.4.0.1+version:             0.4.1.0 author:              Gautier DI FOLCO maintainer:          gautier.difolco@gmail.com category:            Scheduling

@@ -37,6 +37,8 @@     Typed.BucketSize (..),     Typed.WindowSize (..),     Typed.BucketsCount (..),+    Typed.RefillRate (..),+    Typed.DrainRate (..),      -- * Base helpers     alwaysAllow,@@ -48,6 +50,9 @@     slidingWindow,     slidingWindowCount,     slidingWindowBucketed,+    tokenBucket,+    leakyBucket,+    gcra,      -- * Combinators     (.&&),@@ -74,7 +79,7 @@ import qualified Control.NumerusClosus.Typed as Typed import qualified Data.List.NonEmpty as NE import Data.Semigroup (First (..), Last (..), Max (..), Min (..))-import Data.Time (UTCTime (..))+import Data.Time (NominalDiffTime, UTCTime (..))  -- * Main logic @@ -126,6 +131,21 @@ -- | 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++-- | Token bucket with continuous refill.+-- @maxBucket@ is the burst capacity, @refillRate@ is tokens restored per second.+tokenBucket :: Typed.BucketSize Double -> Typed.RefillRate -> UTCTime -> RateLimiter+tokenBucket bs rr = RateLimiter . Typed.tokenBucket bs rr++-- | Leaky bucket rate limiter.+-- @maxBucket@ is the queue capacity, @drainRate@ is how fast it empties.+leakyBucket :: Typed.BucketSize Double -> Typed.DrainRate -> UTCTime -> RateLimiter+leakyBucket bs dr = RateLimiter . Typed.leakyBucket bs dr++-- | Generic Cell Rate Algorithm (GCRA).+-- @limit@ requests allowed per @period@.+gcra :: Int -> NominalDiffTime -> UTCTime -> RateLimiter+gcra limit period = RateLimiter . Typed.gcra limit period  -- * Combinators 

@@ -49,6 +49,14 @@     BucketsCount (..),     SlidingWindowBucketed (..),     slidingWindowBucketed,+    RefillRate (..),+    TokenBucket (..),+    tokenBucket,+    DrainRate (..),+    LeakyBucket (..),+    leakyBucket,+    GCRA (..),+    gcra,      -- * Combinators     type (:&&) (..),@@ -274,6 +282,112 @@ slidingWindowCount :: WindowSize -> BucketSize Int -> UTCTime -> SlidingWindowCount slidingWindowCount windowSize@(WindowSize window) maxBucket windowStart =   SlidingWindowCount {prevCount = 0, currentCount = 0, windowEnd = addUTCTime window windowStart, ..}++-- | The rate at which tokens are refilled (tokens per second).+newtype RefillRate = RefillRate+  { unRefillRate :: Double+  }+  deriving stock (Eq, Ord, Show)++-- | Token bucket with continuous refill.+-- Tokens are consumed on each request and refilled at a steady rate over time.+-- Allows bursts up to the bucket capacity while maintaining a long-term average rate.+data TokenBucket = TokenBucket+  { maxBucket :: BucketSize Double,+    refillRate :: RefillRate,+    -- \* Iteration fields+    tokens :: Double,+    lastRefill :: UTCTime+  }+  deriving stock (Eq, Show, Generic)++instance RateLimiter TokenBucket where+  debit TokenBucket {..} now =+    if refreshedTokens >= 1+      then Right $ TokenBucket {tokens = refreshedTokens - 1, lastRefill = now, ..}+      else+        let tokensNeeded = 1 - refreshedTokens+            secondsToWait = tokensNeeded / refillRate.unRefillRate+         in Left $ DebitableFrom $ nextTimeUnit $ addUTCTime (realToFrac secondsToWait) now+    where+      elapsed = realToFrac (diffUTCTime now lastRefill) :: Double+      refreshedTokens = min maxBucket.unBucketSize (tokens + elapsed * refillRate.unRefillRate)++-- | Create a token bucket rate limiter.+-- @maxBucket@ is the burst capacity, @refillRate@ is tokens restored per second.+tokenBucket :: BucketSize Double -> RefillRate -> UTCTime -> TokenBucket+tokenBucket maxBucket refillRate lastRefill =+  TokenBucket {tokens = maxBucket.unBucketSize, ..}++-- | The rate at which the bucket drains (requests per second).+newtype DrainRate = DrainRate+  { unDrainRate :: Double+  }+  deriving stock (Eq, Ord, Show)++-- | Leaky bucket rate limiter.+-- Requests fill a bucket that drains at a constant rate. If the bucket overflows+-- the request is denied. This produces a perfectly smooth output rate.+data LeakyBucket = LeakyBucket+  { maxBucket :: BucketSize Double,+    drainRate :: DrainRate,+    -- \* Iteration fields+    level :: Double,+    lastDrain :: UTCTime+  }+  deriving stock (Eq, Show, Generic)++instance RateLimiter LeakyBucket where+  debit LeakyBucket {..} now =+    let newLevel = refreshedLevel + 1+     in if newLevel <= maxBucket.unBucketSize+          then Right $ LeakyBucket {level = newLevel, lastDrain = now, ..}+          else+            let excess = newLevel - maxBucket.unBucketSize+                secondsToWait = excess / drainRate.unDrainRate+             in Left $ DebitableFrom $ nextTimeUnit $ addUTCTime (realToFrac secondsToWait) now+    where+      elapsed = realToFrac (diffUTCTime now lastDrain) :: Double+      refreshedLevel = max 0 (level - elapsed * drainRate.unDrainRate)++-- | Create a leaky bucket rate limiter.+-- @maxBucket@ is the queue capacity, @drainRate@ is how fast it empties.+leakyBucket :: BucketSize Double -> DrainRate -> UTCTime -> LeakyBucket+leakyBucket maxBucket drainRate lastDrain =+  LeakyBucket {level = 0, ..}++-- | Generic Cell Rate Algorithm (GCRA).+-- A memoryless variant of leaky bucket that tracks only a single \"theoretical+-- arrival time\" (TAT). Each request advances the TAT by the emission interval+-- (@1 / rate@). A request is allowed if the TAT is not too far in the future+-- (bounded by the burst tolerance, derived from @limit / rate@).+data GCRA = GCRA+  { emissionInterval :: NominalDiffTime,+    burstTolerance :: NominalDiffTime,+    -- \* Iteration field+    tat :: UTCTime+  }+  deriving stock (Eq, Show, Generic)++instance RateLimiter GCRA where+  debit GCRA {..} now =+    let newTat = addUTCTime emissionInterval (max now tat)+        allowAt = addUTCTime (negate burstTolerance) newTat+     in if allowAt <= now+          then Right $ GCRA {tat = newTat, ..}+          else Left $ DebitableFrom $ nextTimeUnit allowAt++-- | Create a GCRA rate limiter.+-- @limit@ is the number of requests allowed in the period, @period@ is the+-- time window. The emission interval is @period / limit@ and the burst+-- tolerance is @period@.+gcra :: Int -> NominalDiffTime -> UTCTime -> GCRA+gcra limit period now =+  GCRA+    { emissionInterval = period / fromIntegral limit,+      burstTolerance = period,+      tat = now+    }  -- | Helper to get the smallest next time unit for debiting. nextTimeUnit :: UTCTime -> UTCTime

@@ -8,7 +8,7 @@ import qualified Control.NumerusClosus as NC import Data.List.NonEmpty (NonEmpty ((:|))) import Data.Semigroup (Max (..), Min (..))-import Data.Time (UTCTime (..), addUTCTime, fromGregorian, secondsToNominalDiffTime)+import Data.Time (NominalDiffTime, UTCTime (..), addUTCTime, fromGregorian, secondsToNominalDiffTime) import qualified Hedgehog import qualified Hedgehog.Gen as Gen import qualified Hedgehog.Range as Range@@ -70,6 +70,10 @@ assertLeftDebitableFromH t (Left (NC.DebitableFrom t')) | t == t' = pure () assertLeftDebitableFromH _ _ = Hedgehog.failure +assertLeftAnyH :: (Hedgehog.MonadTest m) => Either NC.NextDebitable NC.RateLimiter -> m ()+assertLeftAnyH (Left _) = pure ()+assertLeftAnyH (Right _) = Hedgehog.failure+ spec :: Spec spec = do   describe "Control.NumerusClosus" do@@ -306,6 +310,96 @@         Left _ -> Hedgehog.failure       case NC.debit (rl2 .|| rl1) t of         Right _ -> pure ()+        Left _ -> Hedgehog.failure++    describe "tokenBucket" do+      let maxBucket = 5 :: Int+          refillRate = NC.RefillRate 1.0+          rl = NC.tokenBucket (NC.BucketSize (fromIntegral maxBucket)) refillRate baseTime++      it "Allows burst up to maxBucket" do+        assertRight $ debitN maxBucket baseTime rl++      it "Denies after burst exhausted" do+        case debitN maxBucket baseTime rl of+          Right rl' -> assertLeftAny $ NC.debit rl' baseTime+          Left _ -> fail "should allow maxBucket debits"++      it "Refills tokens over time" do+        case debitN maxBucket baseTime rl of+          Right rl' -> assertRight $ NC.debit rl' (offsetTime 1)+          Left _ -> fail "should allow maxBucket debits"++      it "Does not exceed maxBucket after long wait" do+        case debitN maxBucket (offsetTime 1000) rl of+          Right rl' -> assertLeftAny $ NC.debit rl' (offsetTime 1000)+          Left _ -> fail "should allow maxBucket debits after refill"++    describe "leakyBucket" do+      let maxBucket = 5 :: Int+          drainRate = NC.DrainRate 1.0+          rl = NC.leakyBucket (NC.BucketSize (fromIntegral maxBucket)) drainRate baseTime++      it "Allows up to maxBucket requests" do+        assertRight $ debitN maxBucket baseTime rl++      it "Denies when bucket overflows" do+        case debitN maxBucket baseTime rl of+          Right rl' -> assertLeftAny $ NC.debit rl' baseTime+          Left _ -> fail "should allow maxBucket debits"++      it "Drains over time allowing more requests" do+        case debitN maxBucket baseTime rl of+          Right rl' -> assertRight $ NC.debit rl' (offsetTime 1)+          Left _ -> fail "should allow maxBucket debits"++      it "Fully drained after sufficient time" do+        case debitN maxBucket baseTime rl of+          Right rl' -> assertRight $ debitN maxBucket (offsetTime 10) rl'+          Left _ -> fail "should allow maxBucket debits"++    describe "gcra" do+      let limit = 5 :: Int+          period = 60 :: NominalDiffTime+          rl = NC.gcra limit period baseTime++      it "Allows burst up to limit" do+        assertRight $ debitN limit baseTime rl++      it "Denies after limit exhausted within period" do+        case debitN limit baseTime rl of+          Right rl' -> assertLeftAny $ NC.debit rl' baseTime+          Left _ -> fail "should allow limit debits"++      it "Allows again after emission interval passes" do+        case debitN limit baseTime rl of+          Right rl' -> assertRight $ NC.debit rl' (offsetTime 12)+          Left _ -> fail "should allow limit debits"++      it "Fully resets after full period" do+        case debitN limit baseTime rl of+          Right rl' -> assertRight $ debitN limit (offsetTime 61) rl'+          Left _ -> fail "should allow limit debits"++    it "tokenBucket allows exactly maxBucket burst then denies" $ hedgehog do+      maxBucket <- forAll (Gen.integral (Range.linear 1 50))+      let rl = NC.tokenBucket (NC.BucketSize (fromIntegral @Integer maxBucket)) (NC.RefillRate 1.0) baseTime+      case debitN (fromIntegral maxBucket) baseTime rl of+        Right rl' -> assertLeftAnyH $ NC.debit rl' baseTime+        Left _ -> Hedgehog.failure++    it "leakyBucket allows exactly maxBucket then denies" $ hedgehog do+      maxBucket <- forAll (Gen.integral (Range.linear 1 50))+      let rl = NC.leakyBucket (NC.BucketSize (fromIntegral @Integer maxBucket)) (NC.DrainRate 1.0) baseTime+      case debitN (fromIntegral maxBucket) baseTime rl of+        Right rl' -> assertLeftAnyH $ NC.debit rl' baseTime+        Left _ -> Hedgehog.failure++    it "gcra allows exactly limit then denies" $ hedgehog do+      limit <- forAll (Gen.integral (Range.linear 1 50))+      let rl = NC.gcra (fromIntegral @Integer limit) 60 baseTime+      case debitN (fromIntegral limit) baseTime rl of+        Right rl' -> assertLeftAnyH $ NC.debit rl' baseTime         Left _ -> Hedgehog.failure  genNextDebitable :: Hedgehog.Gen NC.NextDebitable