{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
module Main (main) where
import Control.NumerusClosus ((.&&), (.||))
import qualified Control.NumerusClosus as NC
import Data.List.NonEmpty (NonEmpty ((:|)))
import Data.Semigroup (Max (..), Min (..))
import Data.Time (NominalDiffTime, UTCTime (..), addUTCTime, fromGregorian, secondsToNominalDiffTime)
import qualified Hedgehog
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range
import Test.Hspec (Expectation, Spec, describe, hspec, it, shouldBe)
import Test.Hspec.Hedgehog (forAll, hedgehog, (===))
main :: IO ()
main = hspec spec
baseTime :: UTCTime
baseTime = UTCTime (fromGregorian 2024 1 1) 0
offsetTime :: Int -> UTCTime
offsetTime n = addUTCTime (secondsToNominalDiffTime (fromIntegral n)) baseTime
nextTimeUnit :: UTCTime -> UTCTime
nextTimeUnit = addUTCTime 0.000001
genUTCTime :: Hedgehog.Gen UTCTime
genUTCTime = do
s <- Gen.integral (Range.linear @Integer 0 86400)
pure $ addUTCTime (secondsToNominalDiffTime (fromIntegral s)) baseTime
debitN :: Int -> UTCTime -> NC.RateLimiter -> Either NC.NextDebitable NC.RateLimiter
debitN 0 _ rl = Right rl
debitN n t rl = case NC.debit rl t of
Right rl' -> debitN (n - 1) t rl'
Left nd -> Left nd
debitSequence :: [UTCTime] -> NC.RateLimiter -> Either NC.NextDebitable NC.RateLimiter
debitSequence [] rl = Right rl
debitSequence (t : ts) rl = case NC.debit rl t of
Right rl' -> debitSequence ts rl'
Left nd -> Left nd
assertRight :: Either NC.NextDebitable NC.RateLimiter -> Expectation
assertRight (Right _) = pure ()
assertRight (Left nd) = fail $ "Expected Right, got Left " ++ show nd
assertLeftNever :: Either NC.NextDebitable NC.RateLimiter -> Expectation
assertLeftNever (Left NC.Never) = pure ()
assertLeftNever (Left nd) = fail $ "Expected Left Never, got Left " ++ show nd
assertLeftNever (Right _) = fail "Expected Left Never, got Right"
assertLeftDebitableFrom :: UTCTime -> Either NC.NextDebitable NC.RateLimiter -> Expectation
assertLeftDebitableFrom t (Left (NC.DebitableFrom t')) | t == t' = pure ()
assertLeftDebitableFrom t (Left nd) = fail $ "Expected Left (DebitableFrom " ++ show t ++ "), got Left " ++ show nd
assertLeftDebitableFrom _ (Right _) = fail "Expected Left DebitableFrom, got Right"
assertLeftAny :: Either NC.NextDebitable NC.RateLimiter -> Expectation
assertLeftAny (Left _) = pure ()
assertLeftAny (Right _) = fail "Expected Left, got Right"
assertLeftNeverH :: (Hedgehog.MonadTest m) => Either NC.NextDebitable NC.RateLimiter -> m ()
assertLeftNeverH (Left NC.Never) = pure ()
assertLeftNeverH _ = Hedgehog.failure
assertLeftDebitableFromH :: (Hedgehog.MonadTest m) => UTCTime -> Either NC.NextDebitable NC.RateLimiter -> m ()
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
describe "alwaysAllow" do
it "debiting always returns Right" do
assertRight $ NC.debit NC.alwaysAllow baseTime
assertRight $ debitN 100 baseTime NC.alwaysAllow
describe "alwaysDeny" do
it "debiting always returns Left Never" do
assertLeftNever $ NC.debit NC.alwaysDeny baseTime
describe "finiteBucket" do
it "Allows exactly n debits, then denies with Never" do
let rl = NC.finiteBucket 3
case debitN 3 baseTime rl of
Right rl' -> assertLeftNever $ NC.debit rl' baseTime
Left _ -> fail "should allow 3 debits"
it "finiteBucket 0 denies immediately" do
assertLeftNever $ NC.debit (NC.finiteBucket 0) baseTime
it "finiteBucket 1 allows once then denies" do
let rl = NC.finiteBucket 1
case NC.debit rl baseTime of
Right rl' -> assertLeftNever $ NC.debit rl' baseTime
Left _ -> fail "should allow 1 debit"
describe "fixedWindow" do
let windowSize = 60
maxBucket = 5 :: Int
rl = NC.fixedWindow windowSize (fromIntegral maxBucket) baseTime
it "Allows maxBucket debits within window" do
assertRight $ debitN maxBucket baseTime rl
it "Denies after maxBucket debits within same window" do
case debitN maxBucket baseTime rl of
Right rl' -> assertLeftDebitableFrom (nextTimeUnit (offsetTime 60)) $ NC.debit rl' baseTime
Left _ -> fail "should allow maxBucket debits"
it "Resets after window expires" do
case debitN maxBucket baseTime rl of
Right rl' -> assertRight $ NC.debit rl' (offsetTime 61)
Left _ -> fail "should allow maxBucket debits"
it "Returns DebitableFrom with correct next window time" do
case debitN maxBucket baseTime rl of
Right rl' -> assertLeftDebitableFrom (nextTimeUnit (offsetTime 60)) $ NC.debit rl' (offsetTime 10)
Left _ -> fail "should allow maxBucket debits"
describe "slidingWindow" do
let windowSize = 60
maxBucket = 3 :: Int
rl = NC.slidingWindow windowSize (fromIntegral maxBucket)
it "Allows maxBucket debits within window" do
assertRight $ debitSequence [offsetTime 1, offsetTime 2, offsetTime 3] rl
it "Denies when full within window" do
case debitSequence [offsetTime 1, offsetTime 2, offsetTime 3] rl of
Right rl' -> assertLeftDebitableFrom (nextTimeUnit (offsetTime 61)) $ NC.debit rl' (offsetTime 10)
Left _ -> fail "should allow maxBucket debits"
it "Allows again after oldest entry expires from window" do
case debitSequence [offsetTime 1, offsetTime 2, offsetTime 3] rl of
Right rl3 -> do
assertLeftDebitableFrom (nextTimeUnit (offsetTime 61)) $ NC.debit rl3 (offsetTime 30)
assertRight $ NC.debit rl3 (offsetTime 61)
Left _ -> fail "debits failed"
describe "slidingWindowCount" do
let windowSize = 60
maxBucket = 10 :: Int
rl = NC.slidingWindowCount windowSize (fromIntegral maxBucket) baseTime
it "Allows within budget" do
assertRight $ debitN maxBucket baseTime rl
it "Denies when weighted estimate exceeds budget" do
case debitN maxBucket baseTime rl of
Right rl' -> assertLeftAny $ NC.debit rl' baseTime
Left _ -> fail "should allow maxBucket debits"
it "Resets after full window elapses" do
case debitN maxBucket baseTime rl of
Right rl' -> assertRight $ NC.debit rl' (nextTimeUnit (offsetTime 60))
Left _ -> fail "should allow maxBucket debits"
describe "slidingWindowBucketed" do
let windowSize = 60
windowCount = 6
maxBucket = 10 :: Int
rl = NC.slidingWindowBucketed windowSize windowCount (fromIntegral maxBucket)
it "Allows within budget" do
assertRight $ debitN maxBucket baseTime rl
it "Denies when sum exceeds budget" do
case debitN maxBucket baseTime rl of
Right rl' -> assertLeftAny $ NC.debit rl' baseTime
Left _ -> fail "should allow maxBucket debits"
describe "(.&&) combinator" do
let rlAllow = NC.alwaysAllow
rlDeny = NC.alwaysDeny
rlFixed = NC.fixedWindow 60 1 baseTime
it "Both allow -> allows" do
assertRight $ NC.debit (rlAllow .&& rlAllow) baseTime
it "One denies -> denies" do
assertLeftNever $ NC.debit (rlAllow .&& rlDeny) baseTime
assertLeftNever $ NC.debit (rlDeny .&& rlAllow) baseTime
it "Both deny -> merges NextDebitable with max (semigroup)" do
case NC.debit rlFixed baseTime of
Right rlFixed' -> do
let rl1 = rlFixed'
rl2 = NC.alwaysDeny
assertLeftNever $ NC.debit (rl1 .&& rl2) baseTime
Left _ -> fail "should allow 1 debit"
describe "(.||) combinator" do
let rlAllow = NC.alwaysAllow
rlDeny = NC.alwaysDeny
rlFixed = NC.fixedWindow 60 1 baseTime
it "Both allow -> allows" do
assertRight $ NC.debit (rlAllow .|| rlAllow) baseTime
it "One allows -> allows" do
assertRight $ NC.debit (rlAllow .|| rlDeny) baseTime
assertRight $ NC.debit (rlDeny .|| rlAllow) baseTime
it "Both deny -> merges with min of DebitableFrom times" do
case NC.debit rlFixed baseTime of
Right rlFixed' -> do
let rlFixed2 = NC.fixedWindow 120 1 baseTime
case NC.debit rlFixed2 baseTime of
Right rlFixed2' ->
assertLeftDebitableFrom (nextTimeUnit (offsetTime 60)) $ NC.debit (rlFixed' .|| rlFixed2') baseTime
Left _ -> fail "should allow 1 debit"
Left _ -> fail "should allow 1 debit"
it "Both Never -> Never" do
assertLeftNever $ NC.debit (rlDeny .|| rlDeny) baseTime
describe "allOf / anyOf" do
let rlAllow = NC.alwaysAllow
rlDeny = NC.alwaysDeny
it "allOf works like folded .&&" do
assertRight $ NC.debit (NC.allOf (rlAllow :| [rlAllow])) baseTime
assertLeftNever $ NC.debit (NC.allOf (rlAllow :| [rlDeny])) baseTime
it "anyOf works like folded .||" do
assertLeftNever $ NC.debit (NC.anyOf (rlDeny :| [rlDeny])) baseTime
assertRight $ NC.debit (NC.anyOf (rlDeny :| [rlAllow])) baseTime
describe "NextDebitable Min / Max" do
it "Max Never (DebitableFrom t) = Max Never" do
getMax (Max NC.Never <> Max (NC.DebitableFrom baseTime)) `shouldBe` NC.Never
getMax (Max (NC.DebitableFrom baseTime) <> Max NC.Never) `shouldBe` NC.Never
it "Max (DebitableFrom x) (DebitableFrom y) = DebitableFrom (max x y)" do
let t1 = offsetTime 10
t2 = offsetTime 20
getMax (Max (NC.DebitableFrom t1) <> Max (NC.DebitableFrom t2)) `shouldBe` NC.DebitableFrom t2
it "Min Never (DebitableFrom t) = Min (DebitableFrom t)" do
getMin (Min NC.Never <> Min (NC.DebitableFrom baseTime)) `shouldBe` NC.DebitableFrom baseTime
getMin (Min (NC.DebitableFrom baseTime) <> Min NC.Never) `shouldBe` NC.DebitableFrom baseTime
it "Min (DebitableFrom x) (DebitableFrom y) = DebitableFrom (min x y)" do
let t1 = offsetTime 10
t2 = offsetTime 20
getMin (Min (NC.DebitableFrom t1) <> Min (NC.DebitableFrom t2)) `shouldBe` NC.DebitableFrom t1
describe "scheduleWith" do
it "mock PositionInTime pure test" do
let mockPit = NC.PositionInTime {NC.getTime = pure baseTime, NC.delayUntil = \_ -> pure ()}
rl = NC.finiteBucket 1
x <- NC.scheduleWith mockPit rl (pure True)
case x of
Right (_, v) -> v `shouldBe` True
Left _ -> fail "should have worked"
describe "Properties" do
it "finiteBucket allows exactly n then denies" $ hedgehog do
n <- forAll (Gen.integral (Range.linear @Integer 0 100))
let rl = NC.finiteBucket (NC.BucketSize n)
case debitN (fromIntegral n) baseTime rl of
Right rl' -> assertLeftNeverH $ NC.debit rl' baseTime
Left _ -> Hedgehog.failure
it "alwaysAllow never denies" $ hedgehog do
n <- forAll (Gen.integral (Range.linear @Integer 0 1000))
case debitN (fromIntegral n) baseTime NC.alwaysAllow of
Right _ -> pure ()
Left _ -> Hedgehog.failure
it "NextDebitable Max semigroup associativity" $ hedgehog do
a <- forAll genNextDebitable
b <- forAll genNextDebitable
c <- forAll genNextDebitable
Max a <> (Max b <> Max c) === (Max a <> Max b) <> Max c
it "NextDebitable Max Never absorption" $ hedgehog do
x <- forAll genNextDebitable
Max NC.Never <> Max x === Max NC.Never
Max x <> Max NC.Never === Max NC.Never
it "fixedWindow allows exactly maxBucket per window" $ hedgehog do
maxBucket <- forAll (Gen.integral (Range.linear 1 50))
let rl = NC.fixedWindow 60 (NC.BucketSize maxBucket) baseTime
case debitN (fromIntegral maxBucket) baseTime rl of
Right rl' -> assertLeftDebitableFromH (nextTimeUnit (offsetTime 60)) $ NC.debit rl' baseTime
Left _ -> Hedgehog.failure
it "(.&&) is at least as restrictive as either operand" $ hedgehog do
t <- forAll genUTCTime
let rl1 = NC.finiteBucket 0
rl2 = NC.alwaysAllow
assertLeftNeverH $ NC.debit (rl1 .&& rl2) t
assertLeftNeverH $ NC.debit (rl2 .&& rl1) t
it "(.||) is at least as permissive as either operand" $ hedgehog do
t <- forAll genUTCTime
let rl1 = NC.finiteBucket 0
rl2 = NC.alwaysAllow
case NC.debit (rl1 .|| rl2) t of
Right _ -> pure ()
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
genNextDebitable =
Gen.choice
[ pure NC.Never,
NC.DebitableFrom <$> genUTCTime
]