{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TupleSections #-}
-- |
-- Module : Advent.Throttle
-- Copyright : (c) Justin Le 2018
-- License : BSD3
--
-- Maintainer : justin@jle.im
-- Stability : experimental
-- Portability : non-portable
--
-- (Internal) Implement throttling of API requests.
module Advent.Throttle (
Throttler,
newThrottler,
throttling,
setLimit,
getLimit,
) where
import Control.Concurrent
import Control.Exception
import Data.IORef
data Throttler = Throt
{ _throtSem :: QSem
, _throtWaiting :: IORef Int
, _throtLim :: IORef Int
}
acquireThrottler :: Throttler -> IO Bool
acquireThrottler Throt{..} = do
currWait <- readIORef _throtWaiting
throtLim <- readIORef _throtLim
if currWait >= throtLim
then pure False
else do
atomicModifyIORef' _throtWaiting ((,()) . (+ 1))
waitQSem _throtSem `finally` atomicModifyIORef' _throtWaiting ((,()) . subtract 1)
pure True
releaseThrottler :: Throttler -> IO ()
releaseThrottler Throt{..} = signalQSem _throtSem
-- | Set the maximum capacity of a 'Throttler'
setLimit :: Throttler -> Int -> IO ()
setLimit Throt{..} = atomicWriteIORef _throtLim
-- | Get the current maximum capacity of a 'Throttler'
getLimit :: Throttler -> IO Int
getLimit Throt{..} = readIORef _throtLim
-- | Create a new 'Throttler' with a given maximum capacity.
newThrottler :: Int -> IO Throttler
newThrottler n = do
s <- newQSem 1
w <- newIORef 0
l <- newIORef n
pure
Throt
{ _throtSem = s
, _throtWaiting = w
, _throtLim = l
}
-- | Perform an IO action with the given 'Throttler' and delay. The IO
-- action will "wait in line" and be performed when the line is clear. The
-- IO action will delay the next incoming IO action by the delay amount
-- given.
throttling ::
Throttler ->
-- | delay (in milliseconds)
Int ->
IO a ->
IO (Maybe a)
throttling throt delay act = bracketOnError
(acquireThrottler throt)
(const (releaseThrottler throt))
$ \case
False -> pure Nothing
True ->
Just <$> do
res <- act
_ <- forkIO $ do
threadDelay delay
releaseThrottler throt
pure res