Introducing numerus-closus

Last week, I introduced on Hackage numerus-closus, a set of rate-limiting primitives.

Usually we try to rate-limit the inbound traffic, but here the context is a bit different; we are trying to, to some extent, rate-limit the outbound traffic, respecting the quota.

It's true that in some cases, for example, HTTP, there are some attempts to normalize rate-limit in a pull fashion, for example, emitting a response header to specify the current quota, but here we are attempting to limit SNTP (email sending).

I attempted to see if there was any library that could handle it, but they were having some of the following trade-offs, which were unacceptable for my use case:

• Bound to only one algorithm * Embedding the primitive regarding the concurrency
• Coupled with a technology (servant, wai, persistent) or IO
• Bound to only a time frame

Let's start with a simple type representing the result of consuming a token, as follows:

newtype RateLimiter = RateLimiter
  { debit :: UTCTime -> Either NextDebitable RateLimiter
  }


data NextDebitable
  = DebitableFrom UTCTime
  | Never
  deriving stock (Eq, Ord, Show)

The main idea is to turn the problem into a discrete expression, asking: "At this time, can I get a token, or when will the next one be available?"

Note: in a sense, it follows antikythera design.

We can express a few trivial rate limiters as the following ones:

alwaysAllow :: RateLimiter
alwaysAllow =
  RateLimiter
    { debit = const $ Right alwaysAllow
    }

alwaysDeny :: RateLimiter
alwaysDeny =
  RateLimiter
    { debit = const $ Left Never
    }

We can also define basic algebra functions as and, and or as follows:

(.&&) :: 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 $ getLast $ Last x' <> Last y'
          (Left x', _) -> Left x'
          (_, Left y') -> Left y'
    }

infixr 3 .&&

(.||) :: 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 $ getFirst $ First x' <> First y'
          (Right x', _) -> Right $ x' .|| y
          (_, Right y') -> Right $ x .|| y'
    }

infixr 3 .||

We have leveraged First/Last Semigroup instances derived from Ord.

We also can derive foldable-based versions, as follows:

allOf :: NE.NonEmpty RateLimiter -> RateLimiter
allOf = foldl1 (.&&)

anyOf :: NE.NonEmpty RateLimiter -> RateLimiter
anyOf = foldl1 (.||)

An easy algorithm to implement is a bucket with a finite number of tokens, as follows:

newtype BucketSize a
  = BucketSize a
  deriving stock (Eq, Ord, Show)
  deriving newtype (Num)

finiteBucket :: BucketSize Integer -> RateLimiter
finiteBucket (BucketSize n) =
  RateLimiter
    { debit =
        const $
          if n > 0
            then Right $ finiteBucket $ BucketSize (n - 1)
            else Left Never
    }

The idea is simple: given a bucket size, until it is empty, re-create the rate limiter with a smaller bucket size.

The next one is fixed window, implemented as follows:

newtype WindowSize = WindowSize NominalDiffTime
  deriving stock (Eq, Ord, Show)
  deriving newtype (Num, Fractional, Real, RealFrac, FormatTime, ParseTime)

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
        }

nextTimeUnit :: UTCTime -> UTCTime
nextTimeUnit = addUTCTime 0.000001

The idea is to have a finite bucket that is reset when the current time is beyond the current window.

It is a bit more involved than the previous, as we had to use a helper function to hold the internal state.

The next one is sliding window, which registers each token, cleaning those outside of the current window, which ends with the current time, implemented as follows:

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
        }

Note: We rely on Set, with its internal tree representation; cleaning with split is only O(ln n).

Sadly, it has a space complexity of O(bucket size).

This is where sliding window bucketed is useful; it splits a window into sub-windows, only counting in each of them, as follows:

newtype WindowsCount a
  = WindowsCount a
  deriving stock (Eq, Ord, Show)
  deriving newtype (Num)

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
        }

We store each sub-window in a Map; given that keys are stored in a Set, we can reuse split to clean up old entries.

The last one of this log is the most famous one, called sliding window count; it is implemented as follows:

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
        }

It uses two windows, the current and the previous one, estimating the request rate with prev * (1 - elapsed\/window) + current.

That's it for this introduction.

The library goes beyond this:

• Support of token bucket, leaky bucket, GCRA algorithm
• Scheduling and loop functions
• A typed version, which is useful for persistence