{-# 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 _ _ = 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 pertime action =
  generateRateLimitedFunction (PerInvocation pertime) action 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 pertime action =
  generateRateLimitedFunction (PerExecution pertime) action 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 ratelimit action combiner = do
  chan <- atomically newTChan
  void $ forkIO $ runner Nothing 0 chan
  return $ resultFunction chan

  where
  currentMicroseconds :: IO Integer
  currentMicroseconds =
    toMicroseconds . (fromIntegral :: Int -> Picosecond) . fromEnum <$>
      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 mLastRun lastAllowance chan = do
    (req, respMV) <- atomically $ readTChan chan
    let baseHandler resp = putMVar respMV resp

    -- should we wait for some amount of time before running?
    beforeWait <- currentMicroseconds
    let targetPeriod     = toMicroseconds $ getRate ratelimit
        timeSinceLastRun = case mLastRun of
          Just lastRun -> beforeWait - lastRun
          Nothing -> negate targetPeriod
        targetDelay      = targetPeriod - timeSinceLastRun - 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     = afterWait - beforeWait
            overslept = slept - targetDelay
        return overslept

    -- before running, can we combine this with any other requests on the pipe?
    (req', finalHandler) <- updateRequestWithFollowers chan req baseHandler
    let run = action req' >>= 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 chan req handler = do
    isEmpty <- atomically $ isEmptyTChan chan
    if isEmpty
      then return (req, handler)
      else do mCombinedAndMV <- atomically $ do
                tup@(next, nextRespMV) <- readTChan chan
                case combiner req next of
                  Nothing -> unGetTChan chan tup $> Nothing
                  Just combined -> return $ Just (combined, nextRespMV)

              case mCombinedAndMV of
                Nothing ->
                  return (req, handler)
                Just ((req', splitResponse), nextRespMV) ->
                  updateRequestWithFollowers chan req' $ \resp -> do
                    let (theirs, mine) = splitResponse resp
                    putMVar nextRespMV mine
                    handler theirs

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

  -- getRate: what is the rate of this action?
  getRate :: RateLimit t -> t
  getRate (PerInvocation x) = x
  getRate (PerExecution  x) = 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 chan req = do
    respMV <- newEmptyMVar
    atomically $ writeTChan chan (req, respMV)
    takeMVar respMV