{-# 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
PosixTimeUsecs
now <- IO PosixTimeUsecs
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
$! PosixTimeUsecs -> PosixTimeUsecs -> TBData
TBData 0 PosixTimeUsecs
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
-> PosixTimeUsecs -> PosixTimeUsecs -> PosixTimeUsecs -> IO Bool
tokenBucketTryAlloc _ _ 0 _ = Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
tokenBucketTryAlloc _ burst :: PosixTimeUsecs
burst _ alloc :: PosixTimeUsecs
alloc | PosixTimeUsecs
alloc PosixTimeUsecs -> PosixTimeUsecs -> Bool
forall a. Ord a => a -> a -> Bool
> PosixTimeUsecs
burst = Bool -> IO Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
tokenBucketTryAlloc (TB lbref :: IORef TBData
lbref) burst :: PosixTimeUsecs
burst invRate :: PosixTimeUsecs
invRate alloc :: PosixTimeUsecs
alloc = do
PosixTimeUsecs
now <- IO PosixTimeUsecs
getPosixTimeUsecs
IORef TBData -> (TBData -> (TBData, Bool)) -> IO Bool
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef TBData
lbref (PosixTimeUsecs -> TBData -> (TBData, Bool)
go PosixTimeUsecs
now)
where
go :: PosixTimeUsecs -> TBData -> (TBData, Bool)
go now :: PosixTimeUsecs
now (TBData lvl :: PosixTimeUsecs
lvl ts :: PosixTimeUsecs
ts)
| PosixTimeUsecs
lvl'' PosixTimeUsecs -> PosixTimeUsecs -> Bool
forall a. Ord a => a -> a -> Bool
> PosixTimeUsecs
burst = (PosixTimeUsecs -> PosixTimeUsecs -> TBData
TBData PosixTimeUsecs
lvl' PosixTimeUsecs
ts', Bool
False)
| Bool
otherwise = (PosixTimeUsecs -> PosixTimeUsecs -> TBData
TBData PosixTimeUsecs
lvl'' PosixTimeUsecs
ts', Bool
True)
where
lvl' :: PosixTimeUsecs
lvl' = PosixTimeUsecs
lvl PosixTimeUsecs -> PosixTimeUsecs -> PosixTimeUsecs
∸ PosixTimeUsecs
dl
(dl :: PosixTimeUsecs
dl,dtRem :: PosixTimeUsecs
dtRem) = PosixTimeUsecs
dt PosixTimeUsecs
-> PosixTimeUsecs -> (PosixTimeUsecs, PosixTimeUsecs)
forall a. Integral a => a -> a -> (a, a)
`quotRem` PosixTimeUsecs
invRate
dt :: PosixTimeUsecs
dt = PosixTimeUsecs
now PosixTimeUsecs -> PosixTimeUsecs -> PosixTimeUsecs
∸ PosixTimeUsecs
ts
ts' :: PosixTimeUsecs
ts' = PosixTimeUsecs
now PosixTimeUsecs -> PosixTimeUsecs -> PosixTimeUsecs
∸ PosixTimeUsecs
dtRem
lvl'' :: PosixTimeUsecs
lvl'' = PosixTimeUsecs
lvl' PosixTimeUsecs -> PosixTimeUsecs -> PosixTimeUsecs
∔ PosixTimeUsecs
alloc
tokenBucketTryAlloc1 :: TokenBucket
-> Word64
-> Word64
-> IO Word64
tokenBucketTryAlloc1 :: TokenBucket
-> PosixTimeUsecs -> PosixTimeUsecs -> IO PosixTimeUsecs
tokenBucketTryAlloc1 _ _ 0 = PosixTimeUsecs -> IO PosixTimeUsecs
forall (m :: * -> *) a. Monad m => a -> m a
return 0
tokenBucketTryAlloc1 (TB lbref :: IORef TBData
lbref) burst :: PosixTimeUsecs
burst invRate :: PosixTimeUsecs
invRate = do
PosixTimeUsecs
now <- IO PosixTimeUsecs
getPosixTimeUsecs
IORef TBData
-> (TBData -> (TBData, PosixTimeUsecs)) -> IO PosixTimeUsecs
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef TBData
lbref (PosixTimeUsecs -> TBData -> (TBData, PosixTimeUsecs)
go PosixTimeUsecs
now)
where
go :: PosixTimeUsecs -> TBData -> (TBData, PosixTimeUsecs)
go now :: PosixTimeUsecs
now (TBData lvl :: PosixTimeUsecs
lvl ts :: PosixTimeUsecs
ts)
| PosixTimeUsecs
lvl'' PosixTimeUsecs -> PosixTimeUsecs -> Bool
forall a. Ord a => a -> a -> Bool
> PosixTimeUsecs
burst = (PosixTimeUsecs -> PosixTimeUsecs -> TBData
TBData PosixTimeUsecs
lvl' PosixTimeUsecs
ts', PosixTimeUsecs
invRatePosixTimeUsecs -> PosixTimeUsecs -> PosixTimeUsecs
forall a. Num a => a -> a -> a
-PosixTimeUsecs
dtRem)
| Bool
otherwise = (PosixTimeUsecs -> PosixTimeUsecs -> TBData
TBData PosixTimeUsecs
lvl'' PosixTimeUsecs
ts', 0)
where
lvl' :: PosixTimeUsecs
lvl' = PosixTimeUsecs
lvl PosixTimeUsecs -> PosixTimeUsecs -> PosixTimeUsecs
∸ PosixTimeUsecs
dl
(dl :: PosixTimeUsecs
dl,dtRem :: PosixTimeUsecs
dtRem) = PosixTimeUsecs
dt PosixTimeUsecs
-> PosixTimeUsecs -> (PosixTimeUsecs, PosixTimeUsecs)
forall a. Integral a => a -> a -> (a, a)
`quotRem` PosixTimeUsecs
invRate
dt :: PosixTimeUsecs
dt = PosixTimeUsecs
now PosixTimeUsecs -> PosixTimeUsecs -> PosixTimeUsecs
∸ PosixTimeUsecs
ts
ts' :: PosixTimeUsecs
ts' = PosixTimeUsecs
now PosixTimeUsecs -> PosixTimeUsecs -> PosixTimeUsecs
∸ PosixTimeUsecs
dtRem
lvl'' :: PosixTimeUsecs
lvl'' = PosixTimeUsecs
lvl' PosixTimeUsecs -> PosixTimeUsecs -> PosixTimeUsecs
∔ 1
tokenBucketWait :: TokenBucket
-> Word64
-> Word64
-> IO ()
tokenBucketWait :: TokenBucket -> PosixTimeUsecs -> PosixTimeUsecs -> IO ()
tokenBucketWait tb :: TokenBucket
tb burst :: PosixTimeUsecs
burst invRate :: PosixTimeUsecs
invRate = do
PosixTimeUsecs
delay <- TokenBucket
-> PosixTimeUsecs -> PosixTimeUsecs -> IO PosixTimeUsecs
tokenBucketTryAlloc1 TokenBucket
tb PosixTimeUsecs
burst PosixTimeUsecs
invRate
Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (PosixTimeUsecs
delay PosixTimeUsecs -> PosixTimeUsecs -> Bool
forall a. Eq a => a -> a -> Bool
== 0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
Int -> IO ()
threadDelay (PosixTimeUsecs -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral PosixTimeUsecs
delay)
TokenBucket -> PosixTimeUsecs -> PosixTimeUsecs -> IO ()
tokenBucketWait TokenBucket
tb PosixTimeUsecs
burst PosixTimeUsecs
invRate
(∸), (∔) :: Word64 -> Word64 -> Word64
x :: PosixTimeUsecs
x ∸ :: PosixTimeUsecs -> PosixTimeUsecs -> PosixTimeUsecs
∸ y :: PosixTimeUsecs
y = if PosixTimeUsecs
xPosixTimeUsecs -> PosixTimeUsecs -> Bool
forall a. Ord a => a -> a -> Bool
>PosixTimeUsecs
y then PosixTimeUsecs
xPosixTimeUsecs -> PosixTimeUsecs -> PosixTimeUsecs
forall a. Num a => a -> a -> a
-PosixTimeUsecs
y else 0
{-# INLINE (∸) #-}
x :: PosixTimeUsecs
x ∔ :: PosixTimeUsecs -> PosixTimeUsecs -> PosixTimeUsecs
∔ y :: PosixTimeUsecs
y = let s :: PosixTimeUsecs
s=PosixTimeUsecs
xPosixTimeUsecs -> PosixTimeUsecs -> PosixTimeUsecs
forall a. Num a => a -> a -> a
+PosixTimeUsecs
y in if PosixTimeUsecs
x PosixTimeUsecs -> PosixTimeUsecs -> Bool
forall a. Ord a => a -> a -> Bool
<= PosixTimeUsecs
s then PosixTimeUsecs
s else PosixTimeUsecs
forall a. Bounded a => a
maxBound
{-# INLINE (∔) #-}