{-# LANGUAGE Trustworthy #-}

 {-|
  Description : Loop over an action but throttle it to a certain rate
  Copyright   : Robert Fischer, 2016
  License     : BSD3
  Portability : POSIX

  This provides a function, 'loop', which loops on the 'Action' endlessly,
  but no more than a 'Limit' number of times within a given 'TimePeriod'. This is useful
  for preventing your server from accepting too many connections and exhausting your
  available ports, or a file processor from exhausting your available file handles,
  or a client that polls a message queue from abusing it through excessive requests,
  or throttling any other number of processes that would be otherwise out of control.
  -}
module Control.Concurrent.ThrottledLoop
    (
      loop,
      Limit,
      TimePeriod,
      Action,
      ThreadHandle
    ) where

import Control.Concurrent
import Data.Natural
import Data.Time.Clock
import Data.Time.Clock.POSIX

type Limit = Natural
-- ^The number of times that the action will be executed within each period.

type TimePeriod = NominalDiffTime
-- ^The length of time that defines each period.

type Action = IO ()
-- ^An action that can be performed.

type ThreadHandle = IO ThreadId
-- ^A handle on the thread which is performing the loop

loop :: Action -> TimePeriod -> Limit -> ThreadHandle
-- ^Takes an 'Action' that is processed repeatedly in a loop,
-- but no more than the 'Limit' number of times each 'TimePeriod'.
-- Use the 'ThreadHandle' to act on the processing thread.
--
-- Having a 'TimePeriod' of 0 or less will disable the timeout
-- functionality, and instead the loop will just 'yield' in between
-- each execution.
--
-- Having a 'Limit' of 0 is equivalent to having a 'Limit' of 1.
loop io period 0 = loop io period 1
loop io period limit = forkIO $ initLoop io period limit

type Count = Natural
-- ^The count of the number of times we have executed the loop during this period

type Timeout = NominalDiffTime
-- ^The time when the current period expires

initLoop :: Action -> TimePeriod -> Limit -> Action
-- ^Starts a period, and begins looping over the action
initLoop io period 0 = initLoop io period 1
initLoop io period times = do
  now <- getPOSIXTime
  let timeout = now + period
  doLoop io period times 0 timeout

doLoop :: Action -> TimePeriod -> Limit -> Count -> Timeout -> Action
-- ^Implementation of the looping over an action throughout the period
doLoop io period limit cnt timeout | cnt < limit = io >> yield >> (doLoop io period limit (cnt + 1) timeout)
doLoop io period limit _ timeout = delayUntil timeout >> initLoop io period limit

delayUntil :: Timeout -> Action
-- ^Delays the current thread until the timeout passes, yielding if the time has already passed
delayUntil timeout = do
  untilTime <- (timeout -) <$> getPOSIXTime
  let untilMillis = untilTime * 1000
  let untilMicros = untilMillis * 1000
  doDelay $ round untilMicros
    where
      doDelay :: Int -> IO ()
      doDelay x | x > 0 = threadDelay x
      doDelay _         = yield