module Irc.RateLimit
( RateLimit
, newRateLimit
, newRateLimitDefault
, tickRateLimit
) where
import Control.Concurrent
import Control.Monad
import Data.Time
data RateLimit = RateLimit
{ rateStamp :: !(MVar UTCTime)
, rateThreshold :: !Int
, ratePenalty :: !Int
}
newRateLimitDefault :: IO RateLimit
newRateLimitDefault = newRateLimit 2 10
newRateLimit ::
Int ->
Int ->
IO RateLimit
newRateLimit penalty threshold =
do unless (penalty > 0)
(fail "newRateLimit: Penalty too small")
unless (threshold > 0)
(fail "newRateLimit: Threshold too small")
now <- getCurrentTime
ref <- newMVar now
return RateLimit
{ rateStamp = ref
, rateThreshold = threshold
, ratePenalty = penalty
}
tickRateLimit :: RateLimit -> IO ()
tickRateLimit r = modifyMVar_ (rateStamp r) $ \stamp ->
do now <- getCurrentTime
let stamp' = fromIntegral (ratePenalty r) `addUTCTime` max stamp now
diff = diffUTCTime stamp' now
excess = diff - fromIntegral (rateThreshold r)
when (excess > 0)
(threadDelay
(ceiling (1000000 * realToFrac excess :: Rational)))
return stamp'