{-# 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 { Throttler -> QSem
_throtSem :: QSem
, Throttler -> IORef Int
_throtWaiting :: IORef Int
, Throttler -> IORef Int
_throtLim :: IORef Int
}
acquireThrottler :: Throttler -> IO Bool
acquireThrottler :: Throttler -> IO Bool
acquireThrottler Throt{QSem
IORef Int
_throtLim :: IORef Int
_throtWaiting :: IORef Int
_throtSem :: QSem
_throtLim :: Throttler -> IORef Int
_throtWaiting :: Throttler -> IORef Int
_throtSem :: Throttler -> QSem
..} = do
Int
currWait <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
_throtWaiting
Int
throtLim <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
_throtLim
if Int
currWait Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
throtLim
then Bool -> IO Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
else do
IORef Int -> (Int -> (Int, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Int
_throtWaiting ((,()) (Int -> (Int, ())) -> (Int -> Int) -> Int -> (Int, ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1))
QSem -> IO ()
waitQSem QSem
_throtSem IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`finally` IORef Int -> (Int -> (Int, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Int
_throtWaiting ((,()) (Int -> (Int, ())) -> (Int -> Int) -> Int -> (Int, ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1)
Bool -> IO Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
releaseThrottler :: Throttler -> IO ()
releaseThrottler :: Throttler -> IO ()
releaseThrottler Throt{QSem
IORef Int
_throtLim :: IORef Int
_throtWaiting :: IORef Int
_throtSem :: QSem
_throtLim :: Throttler -> IORef Int
_throtWaiting :: Throttler -> IORef Int
_throtSem :: Throttler -> QSem
..} = QSem -> IO ()
signalQSem QSem
_throtSem
setLimit :: Throttler -> Int -> IO ()
setLimit :: Throttler -> Int -> IO ()
setLimit Throt{QSem
IORef Int
_throtLim :: IORef Int
_throtWaiting :: IORef Int
_throtSem :: QSem
_throtLim :: Throttler -> IORef Int
_throtWaiting :: Throttler -> IORef Int
_throtSem :: Throttler -> QSem
..} = IORef Int -> Int -> IO ()
forall a. IORef a -> a -> IO ()
atomicWriteIORef IORef Int
_throtLim
getLimit :: Throttler -> IO Int
getLimit :: Throttler -> IO Int
getLimit Throt{QSem
IORef Int
_throtLim :: IORef Int
_throtWaiting :: IORef Int
_throtSem :: QSem
_throtLim :: Throttler -> IORef Int
_throtWaiting :: Throttler -> IORef Int
_throtSem :: Throttler -> QSem
..} = IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
_throtLim
newThrottler :: Int -> IO Throttler
newThrottler :: Int -> IO Throttler
newThrottler Int
n = do
QSem
s <- Int -> IO QSem
newQSem Int
1
IORef Int
w <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
0
IORef Int
l <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
n
Throttler -> IO Throttler
forall (f :: * -> *) a. Applicative f => a -> f a
pure Throt :: QSem -> IORef Int -> IORef Int -> Throttler
Throt
{ _throtSem :: QSem
_throtSem = QSem
s
, _throtWaiting :: IORef Int
_throtWaiting = IORef Int
w
, _throtLim :: IORef Int
_throtLim = IORef Int
l
}
throttling
:: Throttler
-> Int
-> IO a
-> IO (Maybe a)
throttling :: Throttler -> Int -> IO a -> IO (Maybe a)
throttling Throttler
throt Int
delay IO a
act = IO Bool
-> (Bool -> IO ()) -> (Bool -> IO (Maybe a)) -> IO (Maybe a)
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracketOnError (Throttler -> IO Bool
acquireThrottler Throttler
throt)
(IO () -> Bool -> IO ()
forall a b. a -> b -> a
const (Throttler -> IO ()
releaseThrottler Throttler
throt)) ((Bool -> IO (Maybe a)) -> IO (Maybe a))
-> (Bool -> IO (Maybe a)) -> IO (Maybe a)
forall a b. (a -> b) -> a -> b
$ \case
Bool
False -> Maybe a -> IO (Maybe a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing
Bool
True -> a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> IO a -> IO (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> do
a
res <- IO a
act
ThreadId
_ <- IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ do
Int -> IO ()
threadDelay Int
delay
Throttler -> IO ()
releaseThrottler Throttler
throt
a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
res