57 lines
2.0 KiB
Haskell
57 lines
2.0 KiB
Haskell
module Utils.PersistentTokenBucket
|
|
( TokenBucketSettings(..)
|
|
, persistentTokenBucketTryAlloc'
|
|
, persistentTokenBucketTryAlloc
|
|
) where
|
|
|
|
import Import.NoFoundation
|
|
|
|
|
|
data TokenBucketSettings = TokenBucketSettings
|
|
{ tbsIdent :: TokenBucketIdent
|
|
, tbsDepth :: Word64
|
|
, tbsInvRate :: NominalDiffTime
|
|
, tbsInitialValue :: Int64
|
|
}
|
|
|
|
|
|
persistentTokenBucketTryAlloc' :: (MonadHandler m, HasAppSettings (HandlerSite m), Integral a)
|
|
=> TokenBucketIdent
|
|
-> a
|
|
-> SqlPersistT m Bool
|
|
persistentTokenBucketTryAlloc' tbsIdent tokens = do
|
|
TokenBucketConf{..} <- getsYesod $ views _appPersistentTokenBuckets ($ tbsIdent)
|
|
flip persistentTokenBucketTryAlloc tokens TokenBucketSettings
|
|
{ tbsIdent
|
|
, tbsDepth = tokenBucketDepth
|
|
, tbsInvRate = tokenBucketInvRate
|
|
, tbsInitialValue = tokenBucketInitialValue
|
|
}
|
|
|
|
|
|
persistentTokenBucketTryAlloc :: (MonadIO m, Integral a) => TokenBucketSettings -> a -> SqlPersistT m Bool
|
|
persistentTokenBucketTryAlloc TokenBucketSettings{..} (fromIntegral -> tokens) = do
|
|
now <- liftIO getCurrentTime
|
|
TokenBucket{..} <- do
|
|
existingBucket <- get $ TokenBucketKey tbsIdent
|
|
case existingBucket of
|
|
Just bkt -> return bkt
|
|
Nothing -> do
|
|
let bkt = TokenBucket
|
|
{ tokenBucketIdent = tbsIdent
|
|
, tokenBucketLastValue = tbsInitialValue
|
|
, tokenBucketLastAccess = now
|
|
}
|
|
insert_ bkt
|
|
return bkt
|
|
let currentValue = fromIntegral tbsDepth `min` tokenBucketLastValue + tokenIncrease
|
|
deltaT = now `diffUTCTime` tokenBucketLastAccess
|
|
(tokenIncrease, deltaT')
|
|
| n < 0 = (pred n, (1 + f) * tbsInvRate)
|
|
| otherwise = (n, f * tbsInvRate)
|
|
where (n, f) = properFraction $ deltaT / tbsInvRate
|
|
if | currentValue < 0 -> return False
|
|
| otherwise -> do
|
|
update (TokenBucketKey tbsIdent) [ TokenBucketLastValue =. currentValue - tokens, TokenBucketLastAccess =. addUTCTime (- deltaT') now ]
|
|
return True
|