{-# 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 { 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 <- forall a. IORef a -> IO a
readIORef IORef Int
_throtWaiting
    Int
throtLim <- forall a. IORef a -> IO a
readIORef IORef Int
_throtLim
    if Int
currWait forall a. Ord a => a -> a -> Bool
>= Int
throtLim
      then forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
      else do
        forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Int
_throtWaiting ((,()) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Num a => a -> a -> a
+Int
1))
        QSem -> IO ()
waitQSem QSem
_throtSem forall a b. IO a -> IO b -> IO a
`finally` forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Int
_throtWaiting ((,()) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Num a => a -> a -> a
subtract Int
1)
        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

-- | Set the maximum capacity of a 'Throttler'
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
..} = forall a. IORef a -> a -> IO ()
atomicWriteIORef IORef Int
_throtLim

-- | Get the current maximum capacity of a 'Throttler'
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
..} = forall a. IORef a -> IO a
readIORef IORef Int
_throtLim

-- | Create a new 'Throttler' with a given maximum capacity.
newThrottler :: Int -> IO Throttler
newThrottler :: Int -> IO Throttler
newThrottler Int
n = do
    QSem
s <- Int -> IO QSem
newQSem Int
1
    IORef Int
w <- forall a. a -> IO (IORef a)
newIORef Int
0
    IORef Int
l <- forall a. a -> IO (IORef a)
newIORef Int
n
    forall (f :: * -> *) a. Applicative f => a -> f a
pure Throt
      { _throtSem :: QSem
_throtSem     = QSem
s
      , _throtWaiting :: IORef Int
_throtWaiting = IORef Int
w
      , _throtLim :: IORef Int
_throtLim     = IORef Int
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
    -> Int                  -- ^ delay (in milliseconds)
    -> IO a
    -> IO (Maybe a)
throttling :: forall a. Throttler -> Int -> IO a -> IO (Maybe a)
throttling Throttler
throt Int
delay IO a
act = forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracketOnError (Throttler -> IO Bool
acquireThrottler Throttler
throt)
                                            (forall a b. a -> b -> a
const (Throttler -> IO ()
releaseThrottler Throttler
throt)) forall a b. (a -> b) -> a -> b
$ \case
    Bool
False -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
    Bool
True  -> forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> do
      a
res <- IO a
act
      ThreadId
_ <- IO () -> IO ThreadId
forkIO forall a b. (a -> b) -> a -> b
$ do
        Int -> IO ()
threadDelay Int
delay
        Throttler -> IO ()
releaseThrottler Throttler
throt
      forall (f :: * -> *) a. Applicative f => a -> f a
pure a
res