{-# 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 <- 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

-- | 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
..} = IORef Int -> Int -> IO ()
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
..} = IORef Int -> IO Int
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 <- 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
      }

-- | 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 :: 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