Haskell token bucket rate limiter
Haskell
import Data.IORef (IORef, atomicModifyIORef', newIORef)
import GHC.Clock (getMonotonicTime)
data Bucket = Bucket { available :: Double, lastRefill :: Double }
newBucket :: Double -> IO (IORef Bucket)
newBucket capacity = do
now <- getMonotonicTime
newIORef (Bucket capacity now)
tryAcquire :: Double -> Double -> Double -> IORef Bucket -> IO Bool
tryAcquire capacity rate cost ref = do
now <- getMonotonicTime
atomicModifyIORef' ref (refill now)
where
refill now b =
let refilled = min capacity (available b + (now - lastRefill b) * rate)
in if refilled >= cost
then (Bucket (refilled - cost) now, True)
else (Bucket refilled now, False)