{-# LANGUAGE ScopedTypeVariables #-}

-- | This module implements rate-limiting functionality for Haskell programs.
-- Rate-limiting is useful when trying to control / limit access to a
-- particular resource over time. For example, you might want to limit the
-- rate at which you make requests to a server, as an administrator may block
-- your access if you make too many requests too quickly. Similarly, one may
-- wish to rate-limit certain communication actions, in order to avoid
-- accidentally performing a denial-of-service attack on a critical resource.
--
-- The fundamental idea of this library is that given some basic information
-- about the requests you wante rate limited, it will return you a function
-- that hides all the rate-limiting detail. In short, you make a call to one
-- of the function generators in this file, and you will be returned a function
-- to use. For example:
--
-- @
--   do f <- generateRateLimitedFunction ...
--      ...
--      res1 <- f a
--      ...
--      res2 <- f b
--      ...
-- @
--
-- The calls to the generated function (f) will be rate limited based on the
-- parameters given to 'generateRateLimitedFunction'.
--
-- 'generateRateLimitedFunction' is the most general version of the rate
-- limiting functionality, but specialized versions of it are also exported
-- for convenience.
--
module Control.RateLimit (
    generateRateLimitedFunction
  , RateLimit(..)
  , ResultsCombiner
  , dontCombine
  , rateLimitInvocation
  , rateLimitExecution
  ) where

import Control.Concurrent
import Control.Concurrent.STM
import Control.Monad (void)
import Data.Functor (($>))
import Data.Time.Clock.POSIX (getPOSIXTime)
import Data.Time.Units

-- | The rate at which to limit an action.
data RateLimit a
  = PerInvocation a -- ^ Rate limit the action to invocation once per time
                    -- unit. With this option, the time it takes for the
                    -- action to take place is not taken into consideration
                    -- when computing the rate, only the time between
                    -- invocations of the action. This may cause the action
                    -- to execute concurrently, as an invocation may occur
                    -- while an action is still running.
  | PerExecution a  -- ^ Rate limit the action to execution once per time
                    -- unit. With this option, the time it takes for the
                    -- action to take plase is taken into account, and all
                    -- actions will necessarily occur sequentially. However,
                    -- if your action takes longer than the time unit given,
                    -- then the rate of execution will be slower than the
                    -- given unit of time.

-- | In some cases, if two requests are waiting to be run, it may be possible
-- to combine them into a single request and thus increase the overall
-- bandwidth. The rate limit system supports this, but requires a little
-- additional information to make everything work out right. You may also
-- need to do something a bit wonky with your types to make this work ...
-- sorry.
--
-- The basic idea is this: Given two requests, you can either return Nothing
-- (signalling that the two requests can be combined), or a Just with a new
-- request representing the combination of the two requests. In addition, you
-- will need to provide a function that can turn the response to this single
-- request into two responses, one for each of the original requests.
--
-- I hope this description helps you work through the type, which I'll admit
-- is a bit opaque.
type ResultsCombiner req resp = req -> req -> Maybe (req, resp -> (resp, resp))

dontCombine :: ResultsCombiner a b
dontCombine :: forall a b. ResultsCombiner a b
dontCombine a
_ a
_ = Maybe (a, b -> (b, b))
forall a. Maybe a
Nothing

-- | Rate limit the invocation of a given action. This is equivalent to calling
-- 'generateRateLimitedFunction' with a 'PerInvocation' rate limit and the
-- 'dontCombine' combining function.
rateLimitInvocation :: TimeUnit t
                    => t
                    -> (req -> IO resp)
                    -> IO (req -> IO resp)
rateLimitInvocation :: forall t req resp.
TimeUnit t =>
t -> (req -> IO resp) -> IO (req -> IO resp)
rateLimitInvocation t
pertime req -> IO resp
action =
  RateLimit t
-> (req -> IO resp)
-> ResultsCombiner req resp
-> IO (req -> IO resp)
forall req resp t.
TimeUnit t =>
RateLimit t
-> (req -> IO resp)
-> ResultsCombiner req resp
-> IO (req -> IO resp)
generateRateLimitedFunction (t -> RateLimit t
forall a. a -> RateLimit a
PerInvocation t
pertime) req -> IO resp
action ResultsCombiner req resp
forall a b. ResultsCombiner a b
dontCombine

-- | Rate limit the execution of a given action. This is equivalent to calling
-- 'generateRateLimitedFunction' with a 'PerExecution' rate limit and the
-- 'dontCombine' combining function.
rateLimitExecution :: TimeUnit t
                   => t
                   -> (req -> IO resp)
                   -> IO (req -> IO resp)
rateLimitExecution :: forall t req resp.
TimeUnit t =>
t -> (req -> IO resp) -> IO (req -> IO resp)
rateLimitExecution t
pertime req -> IO resp
action =
  RateLimit t
-> (req -> IO resp)
-> ResultsCombiner req resp
-> IO (req -> IO resp)
forall req resp t.
TimeUnit t =>
RateLimit t
-> (req -> IO resp)
-> ResultsCombiner req resp
-> IO (req -> IO resp)
generateRateLimitedFunction (t -> RateLimit t
forall a. a -> RateLimit a
PerExecution t
pertime) req -> IO resp
action ResultsCombiner req resp
forall a b. ResultsCombiner a b
dontCombine

-- | The most generic way to rate limit an invocation.
generateRateLimitedFunction :: forall req resp t
                             . TimeUnit t
                            => RateLimit t
                               -- ^ What is the rate limit for this action
                            -> (req -> IO resp)
                               -- ^ What is the action you want to rate limit,
                               -- given as an a MonadIO function from requests
                               -- to responses?
                            -> ResultsCombiner req resp
                               -- ^ A function that can combine requests if
                               -- rate limiting happens. If you cannot combine
                               -- two requests into one request, we suggest
                               -- using 'dontCombine'.
                            -> IO (req -> IO resp)
generateRateLimitedFunction :: forall req resp t.
TimeUnit t =>
RateLimit t
-> (req -> IO resp)
-> ResultsCombiner req resp
-> IO (req -> IO resp)
generateRateLimitedFunction RateLimit t
ratelimit req -> IO resp
action ResultsCombiner req resp
combiner = do
  chan <- STM (TChan (req, MVar resp)) -> IO (TChan (req, MVar resp))
forall a. STM a -> IO a
atomically STM (TChan (req, MVar resp))
forall a. STM (TChan a)
newTChan
  void $ forkIO $ runner Nothing 0 chan
  return $ resultFunction chan

  where
  currentMicroseconds :: IO Integer
  currentMicroseconds :: IO Integer
currentMicroseconds =
    Picosecond -> Integer
forall a. TimeUnit a => a -> Integer
toMicroseconds (Picosecond -> Integer)
-> (POSIXTime -> Picosecond) -> POSIXTime -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Picosecond
forall a b. (Integral a, Num b) => a -> b
fromIntegral :: Int -> Picosecond) (Int -> Picosecond)
-> (POSIXTime -> Int) -> POSIXTime -> Picosecond
forall b c a. (b -> c) -> (a -> b) -> a -> c
. POSIXTime -> Int
forall a. Enum a => a -> Int
fromEnum (POSIXTime -> Integer) -> IO POSIXTime -> IO Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
      IO POSIXTime
getPOSIXTime

  -- runner: Repeatedly run requests from the channel, keeping track of the
  -- time immediately before the last request, and a "sleep discount" allowance
  -- we can spend (i.e. reduce future sleep times) based on the amount of time
  -- we've "overslept" in the past.
  runner :: Maybe Integer -> Integer -> TChan (req, MVar resp) -> IO a
  runner :: forall a.
Maybe Integer -> Integer -> TChan (req, MVar resp) -> IO a
runner Maybe Integer
mLastRun Integer
lastAllowance TChan (req, MVar resp)
chan = do
    (req, respMV) <- STM (req, MVar resp) -> IO (req, MVar resp)
forall a. STM a -> IO a
atomically (STM (req, MVar resp) -> IO (req, MVar resp))
-> STM (req, MVar resp) -> IO (req, MVar resp)
forall a b. (a -> b) -> a -> b
$ TChan (req, MVar resp) -> STM (req, MVar resp)
forall a. TChan a -> STM a
readTChan TChan (req, MVar resp)
chan
    let baseHandler resp
resp = MVar resp -> resp -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar resp
respMV resp
resp

    -- should we wait for some amount of time before running?
    beforeWait <- currentMicroseconds
    let targetPeriod     = t -> Integer
forall a. TimeUnit a => a -> Integer
toMicroseconds (t -> Integer) -> t -> Integer
forall a b. (a -> b) -> a -> b
$ RateLimit t -> t
getRate RateLimit t
ratelimit
        timeSinceLastRun = case Maybe Integer
mLastRun of
          Just Integer
lastRun -> Integer
beforeWait Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
lastRun
          Maybe Integer
Nothing -> Integer -> Integer
forall a. Num a => a -> a
negate Integer
targetPeriod
        targetDelay      = Integer
targetPeriod Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
timeSinceLastRun Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
lastAllowance

    -- sleep if necessary; determine sleep-discount allowance for next round
    nextAllowance <- if targetDelay < 0
      then pure $ abs targetDelay -- we have more allowance left
      else do
        -- sleep for *at least* our target delay time
        threadDelay $ fromIntegral targetDelay
        afterWait <- currentMicroseconds
        let slept     = Integer
afterWait Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
beforeWait
            overslept = Integer
slept Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
targetDelay
        return overslept

    -- before running, can we combine this with any other requests on the pipe?
    (req', finalHandler) <- updateRequestWithFollowers chan req baseHandler
    let run = req -> IO resp
action req
req' IO resp -> (resp -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= resp -> IO ()
finalHandler

    beforeRun <- currentMicroseconds
    if shouldFork ratelimit
      then void $ forkIO run
      else run

    runner (Just beforeRun) nextAllowance chan

  -- updateRequestWithFollowers: We have one request. Can we combine it with
  -- some other requests into a cohesive whole?
  updateRequestWithFollowers :: TChan (req, MVar resp)
                             -> req
                             -> (resp -> IO ())
                             -> IO (req, (resp -> IO ()))
  updateRequestWithFollowers :: TChan (req, MVar resp)
-> req -> (resp -> IO ()) -> IO (req, resp -> IO ())
updateRequestWithFollowers TChan (req, MVar resp)
chan req
req resp -> IO ()
handler = do
    isEmpty <- STM Bool -> IO Bool
forall a. STM a -> IO a
atomically (STM Bool -> IO Bool) -> STM Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ TChan (req, MVar resp) -> STM Bool
forall a. TChan a -> STM Bool
isEmptyTChan TChan (req, MVar resp)
chan
    if isEmpty
      then return (req, handler)
      else do mCombinedAndMV <- atomically $ do
                tup@(next, nextRespMV) <- readTChan chan
                case combiner req next of
                  Maybe (req, resp -> (resp, resp))
Nothing -> TChan (req, MVar resp) -> (req, MVar resp) -> STM ()
forall a. TChan a -> a -> STM ()
unGetTChan TChan (req, MVar resp)
chan (req, MVar resp)
tup STM ()
-> Maybe ((req, resp -> (resp, resp)), MVar resp)
-> STM (Maybe ((req, resp -> (resp, resp)), MVar resp))
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Maybe ((req, resp -> (resp, resp)), MVar resp)
forall a. Maybe a
Nothing
                  Just (req, resp -> (resp, resp))
combined -> Maybe ((req, resp -> (resp, resp)), MVar resp)
-> STM (Maybe ((req, resp -> (resp, resp)), MVar resp))
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe ((req, resp -> (resp, resp)), MVar resp)
 -> STM (Maybe ((req, resp -> (resp, resp)), MVar resp)))
-> Maybe ((req, resp -> (resp, resp)), MVar resp)
-> STM (Maybe ((req, resp -> (resp, resp)), MVar resp))
forall a b. (a -> b) -> a -> b
$ ((req, resp -> (resp, resp)), MVar resp)
-> Maybe ((req, resp -> (resp, resp)), MVar resp)
forall a. a -> Maybe a
Just ((req, resp -> (resp, resp))
combined, MVar resp
nextRespMV)

              case mCombinedAndMV of
                Maybe ((req, resp -> (resp, resp)), MVar resp)
Nothing ->
                  (req, resp -> IO ()) -> IO (req, resp -> IO ())
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (req
req, resp -> IO ()
handler)
                Just ((req
req', resp -> (resp, resp)
splitResponse), MVar resp
nextRespMV) ->
                  TChan (req, MVar resp)
-> req -> (resp -> IO ()) -> IO (req, resp -> IO ())
updateRequestWithFollowers TChan (req, MVar resp)
chan req
req' ((resp -> IO ()) -> IO (req, resp -> IO ()))
-> (resp -> IO ()) -> IO (req, resp -> IO ())
forall a b. (a -> b) -> a -> b
$ \resp
resp -> do
                    let (resp
theirs, resp
mine) = resp -> (resp, resp)
splitResponse resp
resp
                    MVar resp -> resp -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar resp
nextRespMV resp
mine
                    resp -> IO ()
handler resp
theirs

  -- shouldFork: should we fork or execute the action in place?
  shouldFork :: RateLimit t -> Bool
  shouldFork :: RateLimit t -> Bool
shouldFork (PerInvocation t
_) = Bool
True
  shouldFork (PerExecution t
_)  = Bool
False

  -- getRate: what is the rate of this action?
  getRate :: RateLimit t -> t
  getRate :: RateLimit t -> t
getRate (PerInvocation t
x) = t
x
  getRate (PerExecution  t
x) = t
x

  -- resultFunction: the function (partially applied on the channel) that will
  -- be returned from this monstrosity.
  resultFunction :: TChan (req, MVar resp) -> req -> IO resp
  resultFunction :: TChan (req, MVar resp) -> req -> IO resp
resultFunction TChan (req, MVar resp)
chan req
req = do
    respMV <- IO (MVar resp)
forall a. IO (MVar a)
newEmptyMVar
    atomically $ writeTChan chan (req, respMV)
    takeMVar respMV