module Irc.RateLimit
( RateLimit
, newRateLimit
, tickRateLimit
) where
import Control.Concurrent
import Control.Monad
import Data.Time
data RateLimit = RateLimit
{ RateLimit -> MVar UTCTime
rateStamp :: !(MVar UTCTime)
, RateLimit -> NominalDiffTime
rateThreshold :: !NominalDiffTime
, RateLimit -> NominalDiffTime
ratePenalty :: !NominalDiffTime
}
newRateLimit ::
Rational ->
Rational ->
IO RateLimit
newRateLimit :: Rational -> Rational -> IO RateLimit
newRateLimit Rational
penalty Rational
threshold =
do UTCTime
now <- IO UTCTime
getCurrentTime
MVar UTCTime
ref <- UTCTime -> IO (MVar UTCTime)
forall a. a -> IO (MVar a)
newMVar UTCTime
now
RateLimit -> IO RateLimit
forall (m :: * -> *) a. Monad m => a -> m a
return RateLimit :: MVar UTCTime -> NominalDiffTime -> NominalDiffTime -> RateLimit
RateLimit
{ rateStamp :: MVar UTCTime
rateStamp = MVar UTCTime
ref
, rateThreshold :: NominalDiffTime
rateThreshold = Rational -> NominalDiffTime
forall a b. (Real a, Fractional b) => a -> b
realToFrac (Rational -> Rational -> Rational
forall a. Ord a => a -> a -> a
max Rational
0 Rational
threshold)
, ratePenalty :: NominalDiffTime
ratePenalty = Rational -> NominalDiffTime
forall a b. (Real a, Fractional b) => a -> b
realToFrac (Rational -> Rational -> Rational
forall a. Ord a => a -> a -> a
max Rational
0 Rational
penalty)
}
tickRateLimit :: RateLimit -> IO ()
tickRateLimit :: RateLimit -> IO ()
tickRateLimit RateLimit
r = MVar UTCTime -> (UTCTime -> IO UTCTime) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ (RateLimit -> MVar UTCTime
rateStamp RateLimit
r) ((UTCTime -> IO UTCTime) -> IO ())
-> (UTCTime -> IO UTCTime) -> IO ()
forall a b. (a -> b) -> a -> b
$ \UTCTime
stamp ->
do UTCTime
now <- IO UTCTime
getCurrentTime
let stamp' :: UTCTime
stamp' = RateLimit -> NominalDiffTime
ratePenalty RateLimit
r NominalDiffTime -> UTCTime -> UTCTime
`addUTCTime` UTCTime -> UTCTime -> UTCTime
forall a. Ord a => a -> a -> a
max UTCTime
stamp UTCTime
now
diff :: NominalDiffTime
diff = UTCTime -> UTCTime -> NominalDiffTime
diffUTCTime UTCTime
stamp' UTCTime
now
excess :: NominalDiffTime
excess = NominalDiffTime
diff NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Num a => a -> a -> a
- RateLimit -> NominalDiffTime
rateThreshold RateLimit
r
Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (NominalDiffTime
excess NominalDiffTime -> NominalDiffTime -> Bool
forall a. Ord a => a -> a -> Bool
> NominalDiffTime
0) (Int -> IO ()
threadDelay (NominalDiffTime -> Int
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (NominalDiffTime
1000000 NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Num a => a -> a -> a
* NominalDiffTime
excess)))
UTCTime -> IO UTCTime
forall (m :: * -> *) a. Monad m => a -> m a
return UTCTime
stamp'