{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

{- |
A standalone rate limiter that accepts some fluctuations in the traffic.

QPS_new = QPS_old * exp(-alpha * t) + n * alpha
where alpha is the decay factor, t is the time elapsed and n is the number of
new requests (that are allowed).

See also www/flib/social/ma_rate_limiter/MovingAverageRateLimiter.php
and fbcode/common/datastruct/MovingAverageRateLimiter.h
-}

module Data.MovingAverageRateLimiter
  ( RateLimiterConfig(qpsLimit, decayHalfLife)
  , mkRateLimiterConfig
  , RateLimiter(qps)
  , mkRateLimiter
  , updateConfig
  , observedRate
  , allow
  ) where

-- ----------------------------------------------------------------------------

data RateLimiterConfig = RateLimiterConfig
  { qpsLimit      :: {-# UNPACK #-} !Double  -- ^ Query per second limit
  , decayHalfLife :: {-# UNPACK #-} !Double
  -- ^ Time in seconds for the weight of past query to decay to half
  -- in moving average calculation
  , alpha         :: {-# UNPACK #-} !Double
  }

mkRateLimiterConfig
  :: Double  -- ^ query per second limit
  -> Double  -- ^ decay half life
  -> RateLimiterConfig
mkRateLimiterConfig ql d = RateLimiterConfig
  { qpsLimit = ql
  , decayHalfLife = d
  , alpha = log 2.0 / d
  }

updateConfig
  :: Double  -- ^ new query per second limit
  -> Double  -- ^ new decay half life
  -> RateLimiterConfig -> RateLimiterConfig
updateConfig ql d r = r
  { qpsLimit = ql
  , decayHalfLife = d
  , alpha = log 2.0 / d
  }

-- ----------------------------------------------------------------------------

data RateLimiter = RateLimiter
  { qps           :: {-# UNPACK #-} !Double
  , lastTimestamp :: {-# UNPACK #-} !Double  -- in sec
  }

mkRateLimiter :: RateLimiter
mkRateLimiter = RateLimiter
  { qps = 0
  , lastTimestamp = 0
  }

observedRate
  :: RateLimiter
  -> Double      -- ^ queries per second
  -> Double      -- ^ timestamp of the last query
  -> RateLimiter
observedRate r q t = r
  { qps = q
  , lastTimestamp = t
  }

allow
  :: RateLimiterConfig
  -> RateLimiter
  -> Double               -- ^ current timestamp
  -> (Bool, RateLimiter)  -- ^ should the query is allowed? and the new state
allow RateLimiterConfig{..} r@RateLimiter{..} t = (a, observedRate r q2 t)
  where
  timeDiff = t - lastTimestamp
  q1 = if timeDiff > 0 then qps * exp (- alpha * timeDiff) else qps
  (a, q2) | lastTimestamp == 0 = (True , qpsLimit * (1 - alpha) + alpha)
                                   -- See Note [First Bump]
          | q1 < qpsLimit      = (True , q1 + alpha)
          | otherwise          = (False, q1)

-- Note [First Bump]
-- If the lastTimestamp is zero, this is the first time a RateLimiter has been
-- bumped. In that case, always allow, then set the qps to the qpsLimit minus
-- enough room for (qpsLimit - 1) additional bumps. This means the RateLimiter
-- will allow qpsLimit bumps at time zero, then throttle at bump qpsLimit + 1.
-- Without this, time is required for the count to approach the decay rate,
-- resulting in an initial spike of allowed bumps. For example, a RateLimiter
-- configured to 1 qps with 600 s half-life would allow at least the first
-- 865 bumps before actually limiting.