{-# LANGUAGE CPP #-}
module Control.Concurrent.TokenBucket
(
TokenBucket
, newTokenBucket
, tokenBucketTryAlloc
, tokenBucketTryAlloc1
, tokenBucketWait
) where
import Control.Concurrent
import Control.Exception
import Control.Monad
import Data.IORef
#if !defined(USE_CBITS)
import Data.Time.Clock.POSIX (getPOSIXTime)
#endif
import Data.Word (Word64)
newtype TokenBucket = TB (IORef TBData)
data TBData = TBData !Word64 !PosixTimeUsecs
deriving Int -> TBData -> ShowS
[TBData] -> ShowS
TBData -> String
(Int -> TBData -> ShowS)
-> (TBData -> String) -> ([TBData] -> ShowS) -> Show TBData
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [TBData] -> ShowS
$cshowList :: [TBData] -> ShowS
show :: TBData -> String
$cshow :: TBData -> String
showsPrec :: Int -> TBData -> ShowS
$cshowsPrec :: Int -> TBData -> ShowS
Show
type PosixTimeUsecs = Word64
#if defined(USE_CBITS)
foreign import ccall unsafe "hs_token_bucket_get_posix_time_usecs"
getPosixTimeUsecs :: IO PosixTimeUsecs
#else
getPosixTimeUsecs :: IO PosixTimeUsecs
getPosixTimeUsecs = fmap (floor . (*1e6)) getPOSIXTime
#endif
newTokenBucket :: IO TokenBucket
newTokenBucket :: IO TokenBucket
newTokenBucket = do
Word64
now <- IO Word64
getPosixTimeUsecs
IORef TBData
lbd <- TBData -> IO (IORef TBData)
forall a. a -> IO (IORef a)
newIORef (TBData -> IO (IORef TBData)) -> TBData -> IO (IORef TBData)
forall a b. (a -> b) -> a -> b
$! Word64 -> Word64 -> TBData
TBData Word64
0 Word64
now
TokenBucket -> IO TokenBucket
forall a. a -> IO a
evaluate (IORef TBData -> TokenBucket
TB IORef TBData
lbd)
tokenBucketTryAlloc :: TokenBucket
-> Word64
-> Word64
-> Word64
-> IO Bool
tokenBucketTryAlloc :: TokenBucket -> Word64 -> Word64 -> Word64 -> IO Bool
tokenBucketTryAlloc TokenBucket
_ Word64
_ Word64
0 Word64
_ = Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
tokenBucketTryAlloc TokenBucket
_ Word64
burst Word64
_ Word64
alloc | Word64
alloc Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
> Word64
burst = Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
tokenBucketTryAlloc (TB IORef TBData
lbref) Word64
burst Word64
invRate Word64
alloc = do
Word64
now <- IO Word64
getPosixTimeUsecs
IORef TBData -> (TBData -> (TBData, Bool)) -> IO Bool
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef TBData
lbref (Word64 -> TBData -> (TBData, Bool)
go Word64
now)
where
go :: Word64 -> TBData -> (TBData, Bool)
go Word64
now (TBData Word64
lvl Word64
ts)
| Word64
lvl'' Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
> Word64
burst = (Word64 -> Word64 -> TBData
TBData Word64
lvl' Word64
ts', Bool
False)
| Bool
otherwise = (Word64 -> Word64 -> TBData
TBData Word64
lvl'' Word64
ts', Bool
True)
where
lvl' :: Word64
lvl' = Word64
lvl Word64 -> Word64 -> Word64
∸ Word64
dl
(Word64
dl,Word64
dtRem) = Word64
dt Word64 -> Word64 -> (Word64, Word64)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Word64
invRate
dt :: Word64
dt = Word64
now Word64 -> Word64 -> Word64
∸ Word64
ts
ts' :: Word64
ts' = Word64
now Word64 -> Word64 -> Word64
∸ Word64
dtRem
lvl'' :: Word64
lvl'' = Word64
lvl' Word64 -> Word64 -> Word64
∔ Word64
alloc
tokenBucketTryAlloc1 :: TokenBucket
-> Word64
-> Word64
-> IO Word64
tokenBucketTryAlloc1 :: TokenBucket -> Word64 -> Word64 -> IO Word64
tokenBucketTryAlloc1 TokenBucket
_ Word64
_ Word64
0 = Word64 -> IO Word64
forall (m :: * -> *) a. Monad m => a -> m a
return Word64
0
tokenBucketTryAlloc1 (TB IORef TBData
lbref) Word64
burst Word64
invRate = do
Word64
now <- IO Word64
getPosixTimeUsecs
IORef TBData -> (TBData -> (TBData, Word64)) -> IO Word64
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef TBData
lbref (Word64 -> TBData -> (TBData, Word64)
go Word64
now)
where
go :: Word64 -> TBData -> (TBData, Word64)
go Word64
now (TBData Word64
lvl Word64
ts)
| Word64
lvl'' Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
> Word64
burst = (Word64 -> Word64 -> TBData
TBData Word64
lvl' Word64
ts', Word64
invRateWord64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
-Word64
dtRem)
| Bool
otherwise = (Word64 -> Word64 -> TBData
TBData Word64
lvl'' Word64
ts', Word64
0)
where
lvl' :: Word64
lvl' = Word64
lvl Word64 -> Word64 -> Word64
∸ Word64
dl
(Word64
dl,Word64
dtRem) = Word64
dt Word64 -> Word64 -> (Word64, Word64)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Word64
invRate
dt :: Word64
dt = Word64
now Word64 -> Word64 -> Word64
∸ Word64
ts
ts' :: Word64
ts' = Word64
now Word64 -> Word64 -> Word64
∸ Word64
dtRem
lvl'' :: Word64
lvl'' = Word64
lvl' Word64 -> Word64 -> Word64
∔ Word64
1
tokenBucketWait :: TokenBucket
-> Word64
-> Word64
-> IO ()
tokenBucketWait :: TokenBucket -> Word64 -> Word64 -> IO ()
tokenBucketWait TokenBucket
tb Word64
burst Word64
invRate = do
Word64
delay <- TokenBucket -> Word64 -> Word64 -> IO Word64
tokenBucketTryAlloc1 TokenBucket
tb Word64
burst Word64
invRate
Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Word64
delay Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
Int -> IO ()
threadDelay (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
delay)
TokenBucket -> Word64 -> Word64 -> IO ()
tokenBucketWait TokenBucket
tb Word64
burst Word64
invRate
(∸), (∔) :: Word64 -> Word64 -> Word64
Word64
x ∸ :: Word64 -> Word64 -> Word64
∸ Word64
y = if Word64
xWord64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
>Word64
y then Word64
xWord64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
-Word64
y else Word64
0
{-# INLINE (∸) #-}
Word64
x ∔ :: Word64 -> Word64 -> Word64
∔ Word64
y = let s :: Word64
s=Word64
xWord64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+Word64
y in if Word64
x Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word64
s then Word64
s else Word64
forall a. Bounded a => a
maxBound
{-# INLINE (∔) #-}