{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TupleSections #-}
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
setLimit :: Throttler -> Int -> IO ()
setLimit Throt{..} = atomicWriteIORef _throtLim
getLimit :: Throttler -> IO Int
getLimit Throt{..} = readIORef _throtLim
newThrottler :: Int -> IO Throttler
newThrottler n = do
s <- newQSem 1
w <- newIORef 0
l <- newIORef n
pure Throt
{ _throtSem = s
, _throtWaiting = w
, _throtLim = l
}
throttling
:: Throttler
-> 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