{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE GADTs #-}
module Keter.RateLimiter.Cache
(
Algorithm(..)
, Cache(..)
, CacheStore(..)
, InMemoryStore(..)
, ResettableStore(..)
, CreateStore(..)
, readCache
, writeCache
, deleteCache
, incrementCache
, newCache
, createInMemoryStore
, clearInMemoryStore
, cacheReset
, algorithmPrefix
, makeCacheKey
, secondsToTimeSpec
, startAutoPurge
, startCustomPurgeTokenBucket
, startCustomPurgeLeakyBucket
, startTokenBucketWorker
, startLeakyBucketWorker
, createTokenBucketEntry
, createLeakyBucketEntry
, algoToText
, parseAlgoText
) where
import Control.Concurrent.STM
import Control.Concurrent.MVar (putMVar, takeMVar, newMVar, readMVar, MVar)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad (void, forever)
import Control.Concurrent (forkIO, threadDelay)
import Data.Aeson (ToJSON, FromJSON, decodeStrict, encode)
import Data.Aeson.Types (Parser)
import qualified Data.ByteString.Lazy as LBS
import Data.Text (Text)
import qualified Data.Text as Tx (unpack, toLower)
import qualified Data.TinyLRU as TinyLRU
import Data.Text.Encoding (encodeUtf8, decodeUtf8')
import qualified Data.Cache as C
import System.Clock (TimeSpec(..), Clock(Monotonic), getTime, toNanoSecs)
import Data.Time.Clock.POSIX (getPOSIXTime)
import qualified StmContainers.Map as StmMap
import qualified Control.Concurrent.STM.TQueue as TQueue
import qualified Focus
import Keter.RateLimiter.Types (TokenBucketState(..), LeakyBucketState(..))
import Keter.RateLimiter.AutoPurge
import Keter.RateLimiter.TokenBucketWorker (startTokenBucketWorker)
import Data.Maybe (fromMaybe)
data Algorithm = FixedWindow | SlidingWindow | TokenBucket | LeakyBucket | TinyLRU
deriving (Int -> Algorithm -> ShowS
[Algorithm] -> ShowS
Algorithm -> String
(Int -> Algorithm -> ShowS)
-> (Algorithm -> String)
-> ([Algorithm] -> ShowS)
-> Show Algorithm
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Algorithm -> ShowS
showsPrec :: Int -> Algorithm -> ShowS
$cshow :: Algorithm -> String
show :: Algorithm -> String
$cshowList :: [Algorithm] -> ShowS
showList :: [Algorithm] -> ShowS
Show, Algorithm -> Algorithm -> Bool
(Algorithm -> Algorithm -> Bool)
-> (Algorithm -> Algorithm -> Bool) -> Eq Algorithm
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Algorithm -> Algorithm -> Bool
== :: Algorithm -> Algorithm -> Bool
$c/= :: Algorithm -> Algorithm -> Bool
/= :: Algorithm -> Algorithm -> Bool
Eq)
algorithmPrefix :: Algorithm
-> Text
algorithmPrefix :: Algorithm -> Text
algorithmPrefix Algorithm
FixedWindow = Text
"rate_limiter"
algorithmPrefix Algorithm
SlidingWindow = Text
"timestamps"
algorithmPrefix Algorithm
TokenBucket = Text
"token_bucket"
algorithmPrefix Algorithm
LeakyBucket = Text
"leaky_bucket"
algorithmPrefix Algorithm
TinyLRU = Text
"tiny_lru"
data Cache store = Cache
{ forall store. Cache store -> Algorithm
cacheAlgorithm :: Algorithm
, forall store. Cache store -> store
cacheStore :: store
}
class MonadIO m => CacheStore store v m | store -> v where
readStore :: store
-> Text
-> Text
-> m (Maybe v)
writeStore :: store
-> Text
-> Text
-> v
-> Int
-> m ()
deleteStore :: store
-> Text
-> Text
-> m ()
incStore :: (FromJSON v, ToJSON v, Ord v, Num v)
=> store
-> Text
-> Text
-> Int
-> m v
incStore store
store Text
prefix Text
key Int
expiresIn = do
Maybe v
mval <- store -> Text -> Text -> m (Maybe v)
forall store v (m :: * -> *).
CacheStore store v m =>
store -> Text -> Text -> m (Maybe v)
readStore store
store Text
prefix Text
key
let newVal :: v
newVal = case Maybe v
mval of
Maybe v
Nothing -> v
1
Just v
v -> if v
v v -> v -> Bool
forall a. Ord a => a -> a -> Bool
<= v
0 then v
1 else v
v v -> v -> v
forall a. Num a => a -> a -> a
+ v
1
store -> Text -> Text -> v -> Int -> m ()
forall store v (m :: * -> *).
CacheStore store v m =>
store -> Text -> Text -> v -> Int -> m ()
writeStore store
store Text
prefix Text
key v
newVal Int
expiresIn
v -> m v
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return v
newVal
class ResettableStore store where
resetStore :: store -> IO ()
data InMemoryStore (a :: Algorithm) where
CounterStore :: TVar (C.Cache Text Text) -> InMemoryStore 'FixedWindow
TimestampStore :: TVar (StmMap.Map Text [Double]) -> InMemoryStore 'SlidingWindow
TokenBucketStore :: TVar (StmMap.Map Text TokenBucketEntry) -> InMemoryStore 'TokenBucket
LeakyBucketStore :: TVar (StmMap.Map Text LeakyBucketEntry) -> InMemoryStore 'LeakyBucket
TinyLRUStore :: TVar (TinyLRU.TinyLRUCache s) -> InMemoryStore 'TinyLRU
class CreateStore (a :: Algorithm) where
createStore :: IO (InMemoryStore a)
secondsToTimeSpec :: Int
-> IO TimeSpec
secondsToTimeSpec :: Int -> IO TimeSpec
secondsToTimeSpec Int
seconds = do
TimeSpec
now <- Clock -> IO TimeSpec
getTime Clock
Monotonic
TimeSpec -> IO TimeSpec
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (TimeSpec -> IO TimeSpec) -> TimeSpec -> IO TimeSpec
forall a b. (a -> b) -> a -> b
$ TimeSpec
now TimeSpec -> TimeSpec -> TimeSpec
forall a. Num a => a -> a -> a
+ Int64 -> Int64 -> TimeSpec
TimeSpec (Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
seconds) Int64
0
instance CreateStore 'FixedWindow where
createStore :: IO (InMemoryStore 'FixedWindow)
createStore = (TVar (Cache Text Text) -> InMemoryStore 'FixedWindow)
-> IO (InMemoryStore 'FixedWindow)
forall (a :: Algorithm).
(TVar (Cache Text Text) -> InMemoryStore a) -> IO (InMemoryStore a)
createStoreWith TVar (Cache Text Text) -> InMemoryStore 'FixedWindow
CounterStore
instance CreateStore 'SlidingWindow where
createStore :: IO (InMemoryStore 'SlidingWindow)
createStore = do
Map Text [Double]
emptyMap <- STM (Map Text [Double]) -> IO (Map Text [Double])
forall a. STM a -> IO a
atomically (STM (Map Text [Double])
forall key value. STM (Map key value)
StmMap.new :: STM (StmMap.Map Text [Double]))
TVar (Map Text [Double])
tvar <- Map Text [Double] -> IO (TVar (Map Text [Double]))
forall a. a -> IO (TVar a)
newTVarIO Map Text [Double]
emptyMap
InMemoryStore 'SlidingWindow -> IO (InMemoryStore 'SlidingWindow)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InMemoryStore 'SlidingWindow -> IO (InMemoryStore 'SlidingWindow))
-> InMemoryStore 'SlidingWindow
-> IO (InMemoryStore 'SlidingWindow)
forall a b. (a -> b) -> a -> b
$ TVar (Map Text [Double]) -> InMemoryStore 'SlidingWindow
TimestampStore TVar (Map Text [Double])
tvar
instance CreateStore 'TokenBucket where
createStore :: IO (InMemoryStore 'TokenBucket)
createStore = do
Map Text TokenBucketEntry
emptyMap <- STM (Map Text TokenBucketEntry) -> IO (Map Text TokenBucketEntry)
forall a. STM a -> IO a
atomically (STM (Map Text TokenBucketEntry)
forall key value. STM (Map key value)
StmMap.new :: STM (StmMap.Map Text TokenBucketEntry))
TVar (Map Text TokenBucketEntry)
tvar <- Map Text TokenBucketEntry -> IO (TVar (Map Text TokenBucketEntry))
forall a. a -> IO (TVar a)
newTVarIO Map Text TokenBucketEntry
emptyMap
IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> IO ThreadId -> IO ()
forall a b. (a -> b) -> a -> b
$ Map Text TokenBucketEntry -> Integer -> Integer -> IO ThreadId
startCustomPurgeTokenBucket Map Text TokenBucketEntry
emptyMap (Integer
60 :: Integer) (Integer
3600 :: Integer)
InMemoryStore 'TokenBucket -> IO (InMemoryStore 'TokenBucket)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InMemoryStore 'TokenBucket -> IO (InMemoryStore 'TokenBucket))
-> InMemoryStore 'TokenBucket -> IO (InMemoryStore 'TokenBucket)
forall a b. (a -> b) -> a -> b
$ TVar (Map Text TokenBucketEntry) -> InMemoryStore 'TokenBucket
TokenBucketStore TVar (Map Text TokenBucketEntry)
tvar
instance CreateStore 'LeakyBucket where
createStore :: IO (InMemoryStore 'LeakyBucket)
createStore = do
Map Text LeakyBucketEntry
emptyMap <- STM (Map Text LeakyBucketEntry) -> IO (Map Text LeakyBucketEntry)
forall a. STM a -> IO a
atomically (STM (Map Text LeakyBucketEntry)
forall key value. STM (Map key value)
StmMap.new :: STM (StmMap.Map Text LeakyBucketEntry))
TVar (Map Text LeakyBucketEntry)
tvar <- Map Text LeakyBucketEntry -> IO (TVar (Map Text LeakyBucketEntry))
forall a. a -> IO (TVar a)
newTVarIO Map Text LeakyBucketEntry
emptyMap
IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> IO ThreadId -> IO ()
forall a b. (a -> b) -> a -> b
$ Map Text LeakyBucketEntry -> Integer -> Integer -> IO ThreadId
startCustomPurgeLeakyBucket Map Text LeakyBucketEntry
emptyMap (Integer
60 :: Integer) (Integer
3600 :: Integer)
InMemoryStore 'LeakyBucket -> IO (InMemoryStore 'LeakyBucket)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (InMemoryStore 'LeakyBucket -> IO (InMemoryStore 'LeakyBucket))
-> InMemoryStore 'LeakyBucket -> IO (InMemoryStore 'LeakyBucket)
forall a b. (a -> b) -> a -> b
$ TVar (Map Text LeakyBucketEntry) -> InMemoryStore 'LeakyBucket
LeakyBucketStore TVar (Map Text LeakyBucketEntry)
tvar
instance CreateStore 'TinyLRU where
createStore :: IO (InMemoryStore 'TinyLRU)
createStore = TVar (TinyLRUCache Any) -> InMemoryStore 'TinyLRU
forall s. TVar (TinyLRUCache s) -> InMemoryStore 'TinyLRU
TinyLRUStore (TVar (TinyLRUCache Any) -> InMemoryStore 'TinyLRU)
-> IO (TVar (TinyLRUCache Any)) -> IO (InMemoryStore 'TinyLRU)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (STM (TVar (TinyLRUCache Any)) -> IO (TVar (TinyLRUCache Any))
forall a. STM a -> IO a
atomically (STM (TVar (TinyLRUCache Any)) -> IO (TVar (TinyLRUCache Any)))
-> STM (TVar (TinyLRUCache Any)) -> IO (TVar (TinyLRUCache Any))
forall a b. (a -> b) -> a -> b
$ TinyLRUCache Any -> STM (TVar (TinyLRUCache Any))
forall a. a -> STM (TVar a)
newTVar (TinyLRUCache Any -> STM (TVar (TinyLRUCache Any)))
-> STM (TinyLRUCache Any) -> STM (TVar (TinyLRUCache Any))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> STM (TinyLRUCache Any)
forall s. Int -> STM (TinyLRUCache s)
TinyLRU.initTinyLRU Int
100)
createStoreWith :: (TVar (C.Cache Text Text) -> InMemoryStore a) -> IO (InMemoryStore a)
createStoreWith :: forall (a :: Algorithm).
(TVar (Cache Text Text) -> InMemoryStore a) -> IO (InMemoryStore a)
createStoreWith TVar (Cache Text Text) -> InMemoryStore a
mkStore = do
Cache Text Text
rawCache <- Maybe TimeSpec -> IO (Cache Text Text)
forall k v. Maybe TimeSpec -> IO (Cache k v)
C.newCache Maybe TimeSpec
forall a. Maybe a
Nothing
MVar Integer
purgeInterval <- Integer -> IO (MVar Integer)
forall a. a -> IO (MVar a)
newMVar (Integer
60 :: Integer)
MVar ()
purgeSignal <- () -> IO (MVar ())
forall a. a -> IO (MVar a)
newMVar ()
IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> IO ThreadId -> IO ()
forall a b. (a -> b) -> a -> b
$ IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ IO () -> IO ()
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar MVar ()
purgeSignal
Integer
interval <- MVar Integer -> IO Integer
forall a. MVar a -> IO a
readMVar MVar Integer
purgeInterval
TimeSpec
startTime <- Clock -> IO TimeSpec
getTime Clock
Monotonic
Cache Text Text -> IO ()
forall k v. (Eq k, Hashable k) => Cache k v -> IO ()
C.purgeExpired Cache Text Text
rawCache
TimeSpec
endTime <- Clock -> IO TimeSpec
getTime Clock
Monotonic
let elapsedMicros :: Integer
elapsedMicros = (TimeSpec -> Integer
toNanoSecs TimeSpec
endTime Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- TimeSpec -> Integer
toNanoSecs TimeSpec
startTime) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
1000
remainingMicros :: Integer
remainingMicros = Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
max (Integer
0 :: Integer) (Integer
interval Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
1000000 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
elapsedMicros)
TimeSpec -> Integer -> MVar () -> IO ()
waitUntilNextPurge TimeSpec
startTime Integer
remainingMicros MVar ()
purgeSignal
TVar (Cache Text Text)
tvar <- Cache Text Text -> IO (TVar (Cache Text Text))
forall a. a -> IO (TVar a)
newTVarIO Cache Text Text
rawCache
InMemoryStore a -> IO (InMemoryStore a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (InMemoryStore a -> IO (InMemoryStore a))
-> InMemoryStore a -> IO (InMemoryStore a)
forall a b. (a -> b) -> a -> b
$ TVar (Cache Text Text) -> InMemoryStore a
mkStore TVar (Cache Text Text)
tvar
waitUntilNextPurge :: TimeSpec -> Integer -> MVar () -> IO ()
waitUntilNextPurge :: TimeSpec -> Integer -> MVar () -> IO ()
waitUntilNextPurge TimeSpec
startTime Integer
remainingMicros MVar ()
purgeSignal = do
TimeSpec
currentTime <- Clock -> IO TimeSpec
getTime Clock
Monotonic
let elapsedMicros :: Integer
elapsedMicros = (TimeSpec -> Integer
toNanoSecs TimeSpec
currentTime Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- TimeSpec -> Integer
toNanoSecs TimeSpec
startTime) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
1000
if Integer
elapsedMicros Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
remainingMicros
then MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ()
purgeSignal ()
else do
let sleepMicros :: Int
sleepMicros = Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
min Integer
remainingMicros (Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int))) :: Int
Int -> IO ()
threadDelay Int
sleepMicros
MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ()
purgeSignal ()
createInMemoryStore :: forall (a :: Algorithm). CreateStore a => IO (InMemoryStore a)
createInMemoryStore :: forall (a :: Algorithm). CreateStore a => IO (InMemoryStore a)
createInMemoryStore = forall (a :: Algorithm). CreateStore a => IO (InMemoryStore a)
createStore @a
newCache :: Algorithm
-> store
-> Cache store
newCache :: forall store. Algorithm -> store -> Cache store
newCache Algorithm
algo store
store = Cache
{ cacheAlgorithm :: Algorithm
cacheAlgorithm = Algorithm
algo
, cacheStore :: store
cacheStore = store
store
}
readCache :: (CacheStore store v IO)
=> Cache store
-> Text
-> IO (Maybe v)
readCache :: forall store v.
CacheStore store v IO =>
Cache store -> Text -> IO (Maybe v)
readCache Cache store
cache Text
unprefixedKey =
store -> Text -> Text -> IO (Maybe v)
forall store v (m :: * -> *).
CacheStore store v m =>
store -> Text -> Text -> m (Maybe v)
readStore (Cache store -> store
forall store. Cache store -> store
cacheStore Cache store
cache) (Algorithm -> Text
algorithmPrefix (Algorithm -> Text) -> Algorithm -> Text
forall a b. (a -> b) -> a -> b
$ Cache store -> Algorithm
forall store. Cache store -> Algorithm
cacheAlgorithm Cache store
cache)
(Algorithm -> Text
algorithmPrefix (Cache store -> Algorithm
forall store. Cache store -> Algorithm
cacheAlgorithm Cache store
cache) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
":" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
unprefixedKey)
writeCache :: (CacheStore store v IO)
=> Cache store
-> Text
-> v
-> Int
-> IO ()
writeCache :: forall store v.
CacheStore store v IO =>
Cache store -> Text -> v -> Int -> IO ()
writeCache Cache store
cache Text
unprefixedKey v
val Int
expiresIn =
store -> Text -> Text -> v -> Int -> IO ()
forall store v (m :: * -> *).
CacheStore store v m =>
store -> Text -> Text -> v -> Int -> m ()
writeStore (Cache store -> store
forall store. Cache store -> store
cacheStore Cache store
cache) (Algorithm -> Text
algorithmPrefix (Algorithm -> Text) -> Algorithm -> Text
forall a b. (a -> b) -> a -> b
$ Cache store -> Algorithm
forall store. Cache store -> Algorithm
cacheAlgorithm Cache store
cache)
(Algorithm -> Text
algorithmPrefix (Cache store -> Algorithm
forall store. Cache store -> Algorithm
cacheAlgorithm Cache store
cache) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
":" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
unprefixedKey) v
val Int
expiresIn
deleteCache :: (CacheStore store v IO)
=> Cache store
-> Text
-> IO ()
deleteCache :: forall store v.
CacheStore store v IO =>
Cache store -> Text -> IO ()
deleteCache Cache store
cache Text
unprefixedKey =
store -> Text -> Text -> IO ()
forall store v (m :: * -> *).
CacheStore store v m =>
store -> Text -> Text -> m ()
deleteStore (Cache store -> store
forall store. Cache store -> store
cacheStore Cache store
cache) (Algorithm -> Text
algorithmPrefix (Algorithm -> Text) -> Algorithm -> Text
forall a b. (a -> b) -> a -> b
$ Cache store -> Algorithm
forall store. Cache store -> Algorithm
cacheAlgorithm Cache store
cache)
(Algorithm -> Text
algorithmPrefix (Cache store -> Algorithm
forall store. Cache store -> Algorithm
cacheAlgorithm Cache store
cache) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
":" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
unprefixedKey)
incrementCache :: (CacheStore store v IO, FromJSON v, ToJSON v, Ord v, Num v)
=> Cache store
-> Text
-> Int
-> IO v
incrementCache :: forall store v.
(CacheStore store v IO, FromJSON v, ToJSON v, Ord v, Num v) =>
Cache store -> Text -> Int -> IO v
incrementCache Cache store
cache Text
unprefixedKey Int
expiresIn = do
let fullKey :: Text
fullKey = Algorithm -> Text
algorithmPrefix (Cache store -> Algorithm
forall store. Cache store -> Algorithm
cacheAlgorithm Cache store
cache) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
":" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
unprefixedKey
prefix :: Text
prefix = Algorithm -> Text
algorithmPrefix (Cache store -> Algorithm
forall store. Cache store -> Algorithm
cacheAlgorithm Cache store
cache)
store -> Text -> Text -> Int -> IO v
forall store v (m :: * -> *).
(CacheStore store v m, FromJSON v, ToJSON v, Ord v, Num v) =>
store -> Text -> Text -> Int -> m v
incStore (Cache store -> store
forall store. Cache store -> store
cacheStore Cache store
cache) Text
prefix Text
fullKey Int
expiresIn
clearInMemoryStore :: ResettableStore store
=> store
-> IO ()
clearInMemoryStore :: forall store. ResettableStore store => store -> IO ()
clearInMemoryStore = store -> IO ()
forall store. ResettableStore store => store -> IO ()
resetStore
cacheReset :: ResettableStore store
=> Cache store
-> IO ()
cacheReset :: forall store. ResettableStore store => Cache store -> IO ()
cacheReset (Cache Algorithm
_ store
store) = store -> IO ()
forall store. ResettableStore store => store -> IO ()
resetStore store
store
createTokenBucketEntry :: TokenBucketState
-> IO TokenBucketEntry
createTokenBucketEntry :: TokenBucketState -> IO TokenBucketEntry
createTokenBucketEntry TokenBucketState
state = do
TVar TokenBucketState
stateVar <- TokenBucketState -> IO (TVar TokenBucketState)
forall a. a -> IO (TVar a)
newTVarIO TokenBucketState
state
TQueue (MVar Bool)
queue <- STM (TQueue (MVar Bool)) -> IO (TQueue (MVar Bool))
forall a. STM a -> IO a
atomically STM (TQueue (MVar Bool))
forall a. STM (TQueue a)
TQueue.newTQueue
TMVar ()
workerLock <- STM (TMVar ()) -> IO (TMVar ())
forall a. STM a -> IO a
atomically STM (TMVar ())
forall a. STM (TMVar a)
newEmptyTMVar
TokenBucketEntry -> IO TokenBucketEntry
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (TokenBucketEntry -> IO TokenBucketEntry)
-> TokenBucketEntry -> IO TokenBucketEntry
forall a b. (a -> b) -> a -> b
$ TVar TokenBucketState
-> TQueue (MVar Bool) -> TMVar () -> TokenBucketEntry
TokenBucketEntry TVar TokenBucketState
stateVar TQueue (MVar Bool)
queue TMVar ()
workerLock
createLeakyBucketEntry :: LeakyBucketState
-> IO LeakyBucketEntry
createLeakyBucketEntry :: LeakyBucketState -> IO LeakyBucketEntry
createLeakyBucketEntry LeakyBucketState
state = do
TVar LeakyBucketState
stateVar <- LeakyBucketState -> IO (TVar LeakyBucketState)
forall a. a -> IO (TVar a)
newTVarIO LeakyBucketState
state
TQueue (TMVar Bool)
queue <- STM (TQueue (TMVar Bool)) -> IO (TQueue (TMVar Bool))
forall a. STM a -> IO a
atomically STM (TQueue (TMVar Bool))
forall a. STM (TQueue a)
TQueue.newTQueue
TMVar ()
workerLock <- STM (TMVar ()) -> IO (TMVar ())
forall a. STM a -> IO a
atomically STM (TMVar ())
forall a. STM (TMVar a)
newEmptyTMVar
LeakyBucketEntry -> IO LeakyBucketEntry
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (LeakyBucketEntry -> IO LeakyBucketEntry)
-> LeakyBucketEntry -> IO LeakyBucketEntry
forall a b. (a -> b) -> a -> b
$ TVar LeakyBucketState
-> TQueue (TMVar Bool) -> TMVar () -> LeakyBucketEntry
LeakyBucketEntry TVar LeakyBucketState
stateVar TQueue (TMVar Bool)
queue TMVar ()
workerLock
instance CacheStore (InMemoryStore 'FixedWindow) Int IO where
readStore :: InMemoryStore 'FixedWindow -> Text -> Text -> IO (Maybe Int)
readStore (CounterStore TVar (Cache Text Text)
tvar) Text
_prefix Text
key = do
Cache Text Text
cache <- TVar (Cache Text Text) -> IO (Cache Text Text)
forall a. TVar a -> IO a
readTVarIO TVar (Cache Text Text)
tvar
Maybe Text
mval <- Cache Text Text -> Text -> IO (Maybe Text)
forall k v. (Eq k, Hashable k) => Cache k v -> k -> IO (Maybe v)
C.lookup Cache Text Text
cache Text
key
case Maybe Text
mval of
Maybe Text
Nothing -> Maybe Int -> IO (Maybe Int)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Int
forall a. Maybe a
Nothing
Just Text
txt -> Maybe Int -> IO (Maybe Int)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Int -> IO (Maybe Int)) -> Maybe Int -> IO (Maybe Int)
forall a b. (a -> b) -> a -> b
$ ByteString -> Maybe Int
forall a. FromJSON a => ByteString -> Maybe a
decodeStrict (Text -> ByteString
encodeUtf8 Text
txt)
writeStore :: InMemoryStore 'FixedWindow -> Text -> Text -> Int -> Int -> IO ()
writeStore (CounterStore TVar (Cache Text Text)
tvar) Text
_prefix Text
key Int
val Int
expiresIn = do
let bs :: ByteString
bs = Int -> ByteString
forall a. ToJSON a => a -> ByteString
encode Int
val
strictBs :: ByteString
strictBs = ByteString -> ByteString
LBS.toStrict ByteString
bs
jsonTxt :: Text
jsonTxt = case ByteString -> Either UnicodeException Text
decodeUtf8' ByteString
strictBs of
Left UnicodeException
_ -> Text
""
Right Text
txt -> Text
txt
TimeSpec
expiryTimeSpec <- Int -> IO TimeSpec
secondsToTimeSpec Int
expiresIn
STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
Cache Text Text
cache <- TVar (Cache Text Text) -> STM (Cache Text Text)
forall a. TVar a -> STM a
readTVar TVar (Cache Text Text)
tvar
Text -> Text -> Cache Text Text -> Maybe TimeSpec -> STM ()
forall k v.
(Eq k, Hashable k) =>
k -> v -> Cache k v -> Maybe TimeSpec -> STM ()
C.insertSTM Text
key Text
jsonTxt Cache Text Text
cache (TimeSpec -> Maybe TimeSpec
forall a. a -> Maybe a
Just TimeSpec
expiryTimeSpec)
deleteStore :: InMemoryStore 'FixedWindow -> Text -> Text -> IO ()
deleteStore (CounterStore TVar (Cache Text Text)
tvar) Text
_prefix Text
key = do
Cache Text Text
cache <- TVar (Cache Text Text) -> IO (Cache Text Text)
forall a. TVar a -> IO a
readTVarIO TVar (Cache Text Text)
tvar
Cache Text Text -> Text -> IO ()
forall k v. (Eq k, Hashable k) => Cache k v -> k -> IO ()
C.delete Cache Text Text
cache Text
key
incStore :: (FromJSON Int, ToJSON Int, Ord Int, Num Int) =>
InMemoryStore 'FixedWindow -> Text -> Text -> Int -> IO Int
incStore (CounterStore TVar (Cache Text Text)
tvar) Text
_prefix Text
key Int
expiresIn = do
TimeSpec
now <- Clock -> IO TimeSpec
getTime Clock
Monotonic
TimeSpec
expiryTimeSpec <- Int -> IO TimeSpec
secondsToTimeSpec Int
expiresIn
STM Int -> IO Int
forall a. STM a -> IO a
atomically (STM Int -> IO Int) -> STM Int -> IO Int
forall a b. (a -> b) -> a -> b
$ do
Cache Text Text
cache <- TVar (Cache Text Text) -> STM (Cache Text Text)
forall a. TVar a -> STM a
readTVar TVar (Cache Text Text)
tvar
Maybe Text
mval <- Bool -> Text -> Cache Text Text -> TimeSpec -> STM (Maybe Text)
forall k v.
(Eq k, Hashable k) =>
Bool -> k -> Cache k v -> TimeSpec -> STM (Maybe v)
C.lookupSTM Bool
False Text
key Cache Text Text
cache TimeSpec
now
let currentVal :: Int
currentVal = case Maybe Text
mval of
Maybe Text
Nothing -> Int
0
Just Text
txt -> Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe Int
0 (ByteString -> Maybe Int
forall a. FromJSON a => ByteString -> Maybe a
decodeStrict (Text -> ByteString
encodeUtf8 Text
txt))
let newVal :: Int
newVal = Int
currentVal Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
let bs :: ByteString
bs = Int -> ByteString
forall a. ToJSON a => a -> ByteString
encode Int
newVal
strictBs :: ByteString
strictBs = ByteString -> ByteString
LBS.toStrict ByteString
bs
jsonTxt :: Text
jsonTxt = case ByteString -> Either UnicodeException Text
decodeUtf8' ByteString
strictBs of
Left UnicodeException
_ -> Text
""
Right Text
txt -> Text
txt
Text -> Text -> Cache Text Text -> Maybe TimeSpec -> STM ()
forall k v.
(Eq k, Hashable k) =>
k -> v -> Cache k v -> Maybe TimeSpec -> STM ()
C.insertSTM Text
key Text
jsonTxt Cache Text Text
cache (TimeSpec -> Maybe TimeSpec
forall a. a -> Maybe a
Just TimeSpec
expiryTimeSpec)
Int -> STM Int
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
newVal
instance CacheStore (InMemoryStore 'SlidingWindow) [Double] IO where
readStore :: InMemoryStore 'SlidingWindow -> Text -> Text -> IO (Maybe [Double])
readStore (TimestampStore TVar (Map Text [Double])
tvar) Text
_prefix Text
key = STM (Maybe [Double]) -> IO (Maybe [Double])
forall a. STM a -> IO a
atomically (STM (Maybe [Double]) -> IO (Maybe [Double]))
-> STM (Maybe [Double]) -> IO (Maybe [Double])
forall a b. (a -> b) -> a -> b
$ do
Map Text [Double]
stmMap <- TVar (Map Text [Double]) -> STM (Map Text [Double])
forall a. TVar a -> STM a
readTVar TVar (Map Text [Double])
tvar
Text -> Map Text [Double] -> STM (Maybe [Double])
forall key value.
Hashable key =>
key -> Map key value -> STM (Maybe value)
StmMap.lookup Text
key Map Text [Double]
stmMap
writeStore :: InMemoryStore 'SlidingWindow
-> Text -> Text -> [Double] -> Int -> IO ()
writeStore (TimestampStore TVar (Map Text [Double])
tvar) Text
_prefix Text
key [Double]
val Int
_expiresIn = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
Map Text [Double]
stmMap <- TVar (Map Text [Double]) -> STM (Map Text [Double])
forall a. TVar a -> STM a
readTVar TVar (Map Text [Double])
tvar
[Double] -> Text -> Map Text [Double] -> STM ()
forall key value.
Hashable key =>
value -> key -> Map key value -> STM ()
StmMap.insert [Double]
val Text
key Map Text [Double]
stmMap
deleteStore :: InMemoryStore 'SlidingWindow -> Text -> Text -> IO ()
deleteStore (TimestampStore TVar (Map Text [Double])
tvar) Text
_prefix Text
key = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
Map Text [Double]
stmMap <- TVar (Map Text [Double]) -> STM (Map Text [Double])
forall a. TVar a -> STM a
readTVar TVar (Map Text [Double])
tvar
Text -> Map Text [Double] -> STM ()
forall key value. Hashable key => key -> Map key value -> STM ()
StmMap.delete Text
key Map Text [Double]
stmMap
instance CacheStore (InMemoryStore 'TokenBucket) TokenBucketState IO where
readStore :: InMemoryStore 'TokenBucket
-> Text -> Text -> IO (Maybe TokenBucketState)
readStore (TokenBucketStore TVar (Map Text TokenBucketEntry)
tvar) Text
_prefix Text
key = STM (Maybe TokenBucketState) -> IO (Maybe TokenBucketState)
forall a. STM a -> IO a
atomically (STM (Maybe TokenBucketState) -> IO (Maybe TokenBucketState))
-> STM (Maybe TokenBucketState) -> IO (Maybe TokenBucketState)
forall a b. (a -> b) -> a -> b
$ do
Map Text TokenBucketEntry
stmMap <- TVar (Map Text TokenBucketEntry) -> STM (Map Text TokenBucketEntry)
forall a. TVar a -> STM a
readTVar TVar (Map Text TokenBucketEntry)
tvar
Maybe TokenBucketEntry
mval <- Text -> Map Text TokenBucketEntry -> STM (Maybe TokenBucketEntry)
forall key value.
Hashable key =>
key -> Map key value -> STM (Maybe value)
StmMap.lookup Text
key Map Text TokenBucketEntry
stmMap
case Maybe TokenBucketEntry
mval of
Maybe TokenBucketEntry
Nothing -> Maybe TokenBucketState -> STM (Maybe TokenBucketState)
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe TokenBucketState
forall a. Maybe a
Nothing
Just TokenBucketEntry
entry -> TokenBucketState -> Maybe TokenBucketState
forall a. a -> Maybe a
Just (TokenBucketState -> Maybe TokenBucketState)
-> STM TokenBucketState -> STM (Maybe TokenBucketState)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar TokenBucketState -> STM TokenBucketState
forall a. TVar a -> STM a
readTVar (TokenBucketEntry -> TVar TokenBucketState
tbeState TokenBucketEntry
entry)
writeStore :: InMemoryStore 'TokenBucket
-> Text -> Text -> TokenBucketState -> Int -> IO ()
writeStore (TokenBucketStore TVar (Map Text TokenBucketEntry)
tvar) Text
_prefix Text
key TokenBucketState
val Int
_expiresIn = do
TokenBucketEntry
entry <- TokenBucketState -> IO TokenBucketEntry
createTokenBucketEntry TokenBucketState
val
STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
Map Text TokenBucketEntry
stmMap <- TVar (Map Text TokenBucketEntry) -> STM (Map Text TokenBucketEntry)
forall a. TVar a -> STM a
readTVar TVar (Map Text TokenBucketEntry)
tvar
Focus TokenBucketEntry STM ()
-> Text -> Map Text TokenBucketEntry -> STM ()
forall key value result.
Hashable key =>
Focus value STM result -> key -> Map key value -> STM result
StmMap.focus ((TokenBucketEntry -> TVar TokenBucketState)
-> TokenBucketEntry
-> TokenBucketState
-> Focus TokenBucketEntry STM ()
forall entry v.
(entry -> TVar v) -> entry -> v -> Focus entry STM ()
focusInsertOrUpdate TokenBucketEntry -> TVar TokenBucketState
tbeState TokenBucketEntry
entry TokenBucketState
val) Text
key Map Text TokenBucketEntry
stmMap
deleteStore :: InMemoryStore 'TokenBucket -> Text -> Text -> IO ()
deleteStore (TokenBucketStore TVar (Map Text TokenBucketEntry)
tvar) Text
_prefix Text
key = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
Map Text TokenBucketEntry
stmMap <- TVar (Map Text TokenBucketEntry) -> STM (Map Text TokenBucketEntry)
forall a. TVar a -> STM a
readTVar TVar (Map Text TokenBucketEntry)
tvar
Text -> Map Text TokenBucketEntry -> STM ()
forall key value. Hashable key => key -> Map key value -> STM ()
StmMap.delete Text
key Map Text TokenBucketEntry
stmMap
instance CacheStore (InMemoryStore 'LeakyBucket) LeakyBucketState IO where
readStore :: InMemoryStore 'LeakyBucket
-> Text -> Text -> IO (Maybe LeakyBucketState)
readStore (LeakyBucketStore TVar (Map Text LeakyBucketEntry)
tvar) Text
_prefix Text
key = STM (Maybe LeakyBucketState) -> IO (Maybe LeakyBucketState)
forall a. STM a -> IO a
atomically (STM (Maybe LeakyBucketState) -> IO (Maybe LeakyBucketState))
-> STM (Maybe LeakyBucketState) -> IO (Maybe LeakyBucketState)
forall a b. (a -> b) -> a -> b
$ do
Map Text LeakyBucketEntry
stmMap <- TVar (Map Text LeakyBucketEntry) -> STM (Map Text LeakyBucketEntry)
forall a. TVar a -> STM a
readTVar TVar (Map Text LeakyBucketEntry)
tvar
Maybe LeakyBucketEntry
mval <- Text -> Map Text LeakyBucketEntry -> STM (Maybe LeakyBucketEntry)
forall key value.
Hashable key =>
key -> Map key value -> STM (Maybe value)
StmMap.lookup Text
key Map Text LeakyBucketEntry
stmMap
case Maybe LeakyBucketEntry
mval of
Maybe LeakyBucketEntry
Nothing -> Maybe LeakyBucketState -> STM (Maybe LeakyBucketState)
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe LeakyBucketState
forall a. Maybe a
Nothing
Just LeakyBucketEntry
entry -> LeakyBucketState -> Maybe LeakyBucketState
forall a. a -> Maybe a
Just (LeakyBucketState -> Maybe LeakyBucketState)
-> STM LeakyBucketState -> STM (Maybe LeakyBucketState)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar LeakyBucketState -> STM LeakyBucketState
forall a. TVar a -> STM a
readTVar (LeakyBucketEntry -> TVar LeakyBucketState
lbeState LeakyBucketEntry
entry)
writeStore :: InMemoryStore 'LeakyBucket
-> Text -> Text -> LeakyBucketState -> Int -> IO ()
writeStore (LeakyBucketStore TVar (Map Text LeakyBucketEntry)
tvar) Text
_prefix Text
key LeakyBucketState
val Int
_expiresIn = do
LeakyBucketEntry
entry <- LeakyBucketState -> IO LeakyBucketEntry
createLeakyBucketEntry LeakyBucketState
val
STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
Map Text LeakyBucketEntry
stmMap <- TVar (Map Text LeakyBucketEntry) -> STM (Map Text LeakyBucketEntry)
forall a. TVar a -> STM a
readTVar TVar (Map Text LeakyBucketEntry)
tvar
Focus LeakyBucketEntry STM ()
-> Text -> Map Text LeakyBucketEntry -> STM ()
forall key value result.
Hashable key =>
Focus value STM result -> key -> Map key value -> STM result
StmMap.focus ((LeakyBucketEntry -> TVar LeakyBucketState)
-> LeakyBucketEntry
-> LeakyBucketState
-> Focus LeakyBucketEntry STM ()
forall entry v.
(entry -> TVar v) -> entry -> v -> Focus entry STM ()
focusInsertOrUpdate LeakyBucketEntry -> TVar LeakyBucketState
lbeState LeakyBucketEntry
entry LeakyBucketState
val) Text
key Map Text LeakyBucketEntry
stmMap
deleteStore :: InMemoryStore 'LeakyBucket -> Text -> Text -> IO ()
deleteStore (LeakyBucketStore TVar (Map Text LeakyBucketEntry)
tvar) Text
_prefix Text
key = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
Map Text LeakyBucketEntry
stmMap <- TVar (Map Text LeakyBucketEntry) -> STM (Map Text LeakyBucketEntry)
forall a. TVar a -> STM a
readTVar TVar (Map Text LeakyBucketEntry)
tvar
Text -> Map Text LeakyBucketEntry -> STM ()
forall key value. Hashable key => key -> Map key value -> STM ()
StmMap.delete Text
key Map Text LeakyBucketEntry
stmMap
focusInsertOrUpdate
:: (entry -> TVar v)
-> entry
-> v
-> Focus.Focus entry STM ()
focusInsertOrUpdate :: forall entry v.
(entry -> TVar v) -> entry -> v -> Focus entry STM ()
focusInsertOrUpdate entry -> TVar v
getState entry
entry v
newVal = STM ((), Change entry)
-> (entry -> STM ((), Change entry)) -> Focus entry STM ()
forall element (m :: * -> *) result.
m (result, Change element)
-> (element -> m (result, Change element))
-> Focus element m result
Focus.Focus
(((), Change entry) -> STM ((), Change entry)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((), entry -> Change entry
forall a. a -> Change a
Focus.Set entry
entry))
(\entry
existing -> do
TVar v -> v -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (entry -> TVar v
getState entry
existing) v
newVal
((), Change entry) -> STM ((), Change entry)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((), Change entry
forall a. Change a
Focus.Leave))
startLeakyBucketWorker
:: TVar LeakyBucketState
-> TQueue.TQueue (TMVar Bool)
-> Int
-> Double
-> IO ()
startLeakyBucketWorker :: TVar LeakyBucketState
-> TQueue (TMVar Bool) -> Int -> Double -> IO ()
startLeakyBucketWorker TVar LeakyBucketState
stateVar TQueue (TMVar Bool)
queue Int
capacity Double
leakRate = IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> (IO () -> IO ThreadId) -> IO () -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO () -> IO ThreadId
forkIO (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
IO () -> IO ()
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
TMVar Bool
replyVar <- STM (TMVar Bool) -> IO (TMVar Bool)
forall a. STM a -> IO a
atomically (STM (TMVar Bool) -> IO (TMVar Bool))
-> STM (TMVar Bool) -> IO (TMVar Bool)
forall a b. (a -> b) -> a -> b
$ TQueue (TMVar Bool) -> STM (TMVar Bool)
forall a. TQueue a -> STM a
readTQueue TQueue (TMVar Bool)
queue
Double
now <- POSIXTime -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac (POSIXTime -> Double) -> IO POSIXTime -> IO Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO POSIXTime
getPOSIXTime
Bool
result <- STM Bool -> IO Bool
forall a. STM a -> IO a
atomically (STM Bool -> IO Bool) -> STM Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ do
LeakyBucketState{Double
level :: Double
level :: LeakyBucketState -> Double
level,Double
lastTime :: Double
lastTime :: LeakyBucketState -> Double
lastTime} <- TVar LeakyBucketState -> STM LeakyBucketState
forall a. TVar a -> STM a
readTVar TVar LeakyBucketState
stateVar
let elapsed :: Double
elapsed = Double
now Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
lastTime
leakedLevel :: Double
leakedLevel = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
0 (Double
level Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
elapsed Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
leakRate)
nextLevel :: Double
nextLevel = Double
leakedLevel Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1
allowed :: Bool
allowed = Double
nextLevel Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
capacity
finalLevel :: Double
finalLevel = if Bool
allowed then Double
nextLevel else Double
leakedLevel
TVar LeakyBucketState -> LeakyBucketState -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar LeakyBucketState
stateVar LeakyBucketState{ level :: Double
level = Double
finalLevel, lastTime :: Double
lastTime = Double
now }
Bool -> STM Bool
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
allowed
STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TMVar Bool -> Bool -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar TMVar Bool
replyVar Bool
result
instance CacheStore (InMemoryStore 'TinyLRU) Int IO where
readStore :: InMemoryStore 'TinyLRU -> Text -> Text -> IO (Maybe Int)
readStore (TinyLRUStore TVar (TinyLRUCache s)
ref) Text
_prefix Text
key = do
TimeSpec
now <- IO TimeSpec -> IO TimeSpec
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO TimeSpec -> IO TimeSpec) -> IO TimeSpec -> IO TimeSpec
forall a b. (a -> b) -> a -> b
$ Clock -> IO TimeSpec
getTime Clock
Monotonic
STM (Maybe Int) -> IO (Maybe Int)
forall a. STM a -> IO a
atomically (STM (Maybe Int) -> IO (Maybe Int))
-> STM (Maybe Int) -> IO (Maybe Int)
forall a b. (a -> b) -> a -> b
$ do
TinyLRUCache s
cache <- TVar (TinyLRUCache s) -> STM (TinyLRUCache s)
forall a. TVar a -> STM a
readTVar TVar (TinyLRUCache s)
ref
Maybe (TVar (LRUNode s))
maybeNodeRef <- Text
-> Map Text (TVar (LRUNode s)) -> STM (Maybe (TVar (LRUNode s)))
forall key value.
Hashable key =>
key -> Map key value -> STM (Maybe value)
StmMap.lookup Text
key (TinyLRUCache s -> Map Text (TVar (LRUNode s))
forall s. TinyLRUCache s -> Map Text (TVar (LRUNode s))
TinyLRU.lruCache TinyLRUCache s
cache)
case Maybe (TVar (LRUNode s))
maybeNodeRef of
Just TVar (LRUNode s)
nodeRef -> do
LRUNode s
node <- TVar (LRUNode s) -> STM (LRUNode s)
forall a. TVar a -> STM a
readTVar TVar (LRUNode s)
nodeRef
let expired :: Bool
expired = TimeSpec -> LRUNode s -> Bool
forall s. TimeSpec -> LRUNode s -> Bool
TinyLRU.isExpired TimeSpec
now LRUNode s
node
if Bool
expired
then do
Text -> TinyLRUCache s -> STM ()
forall s. Text -> TinyLRUCache s -> STM ()
TinyLRU.deleteKey Text
key TinyLRUCache s
cache
Maybe Int -> STM (Maybe Int)
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Int
forall a. Maybe a
Nothing
else do
let Maybe Int
decoded :: Maybe Int = ByteString -> Maybe Int
forall a. FromJSON a => ByteString -> Maybe a
decodeStrict (LRUNode s -> ByteString
forall s. LRUNode s -> ByteString
TinyLRU.nodeValue LRUNode s
node)
TinyLRUCache s -> TVar (LRUNode s) -> STM ()
forall s. TinyLRUCache s -> TVar (LRUNode s) -> STM ()
TinyLRU.moveToFrontInCache TinyLRUCache s
cache TVar (LRUNode s)
nodeRef
Maybe Int -> STM (Maybe Int)
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Int
decoded
Maybe (TVar (LRUNode s))
Nothing -> Maybe Int -> STM (Maybe Int)
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Int
forall a. Maybe a
Nothing
writeStore :: InMemoryStore 'TinyLRU -> Text -> Text -> Int -> Int -> IO ()
writeStore (TinyLRUStore TVar (TinyLRUCache s)
ref) Text
_prefix Text
key Int
val Int
expiresIn = do
TimeSpec
now <- IO TimeSpec -> IO TimeSpec
forall a. IO a -> IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO TimeSpec -> IO TimeSpec) -> IO TimeSpec -> IO TimeSpec
forall a b. (a -> b) -> a -> b
$ Clock -> IO TimeSpec
getTime Clock
Monotonic
STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
TinyLRUCache s
cache <- TVar (TinyLRUCache s) -> STM (TinyLRUCache s)
forall a. TVar a -> STM a
readTVar TVar (TinyLRUCache s)
ref
Maybe Int
_ <- TimeSpec -> Text -> Int -> Int -> TinyLRUCache s -> STM (Maybe Int)
forall a s.
(FromJSON a, ToJSON a) =>
TimeSpec -> Text -> a -> Int -> TinyLRUCache s -> STM (Maybe a)
TinyLRU.updateValue TimeSpec
now Text
key Int
val Int
expiresIn TinyLRUCache s
cache
() -> STM ()
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
deleteStore :: InMemoryStore 'TinyLRU -> Text -> Text -> IO ()
deleteStore (TinyLRUStore TVar (TinyLRUCache s)
ref) Text
_prefix Text
key = do
STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
TinyLRUCache s
cache <- TVar (TinyLRUCache s) -> STM (TinyLRUCache s)
forall a. TVar a -> STM a
readTVar TVar (TinyLRUCache s)
ref
Text -> TinyLRUCache s -> STM ()
forall s. Text -> TinyLRUCache s -> STM ()
TinyLRU.deleteKey Text
key TinyLRUCache s
cache
instance ResettableStore (InMemoryStore a) where
resetStore :: InMemoryStore a -> IO ()
resetStore (CounterStore TVar (Cache Text Text)
tvar) = TVar (Cache Text Text) -> IO ()
resetStoreWith TVar (Cache Text Text)
tvar
resetStore (TimestampStore TVar (Map Text [Double])
tvar) = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
Map Text [Double]
stmMap <- TVar (Map Text [Double]) -> STM (Map Text [Double])
forall a. TVar a -> STM a
readTVar TVar (Map Text [Double])
tvar
Map Text [Double] -> STM ()
forall key value. Map key value -> STM ()
StmMap.reset Map Text [Double]
stmMap
resetStore (TokenBucketStore TVar (Map Text TokenBucketEntry)
tvar) = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
Map Text TokenBucketEntry
stmMap <- TVar (Map Text TokenBucketEntry) -> STM (Map Text TokenBucketEntry)
forall a. TVar a -> STM a
readTVar TVar (Map Text TokenBucketEntry)
tvar
Map Text TokenBucketEntry -> STM ()
forall key value. Map key value -> STM ()
StmMap.reset Map Text TokenBucketEntry
stmMap
resetStore (LeakyBucketStore TVar (Map Text LeakyBucketEntry)
tvar) = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
Map Text LeakyBucketEntry
stmMap <- TVar (Map Text LeakyBucketEntry) -> STM (Map Text LeakyBucketEntry)
forall a. TVar a -> STM a
readTVar TVar (Map Text LeakyBucketEntry)
tvar
Map Text LeakyBucketEntry -> STM ()
forall key value. Map key value -> STM ()
StmMap.reset Map Text LeakyBucketEntry
stmMap
resetStore (TinyLRUStore TVar (TinyLRUCache s)
ref) = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
TinyLRUCache s
cache <- TVar (TinyLRUCache s) -> STM (TinyLRUCache s)
forall a. TVar a -> STM a
readTVar TVar (TinyLRUCache s)
ref
TinyLRUCache s -> STM ()
forall s. TinyLRUCache s -> STM ()
TinyLRU.resetTinyLRU TinyLRUCache s
cache
resetStoreWith :: TVar (C.Cache Text Text) -> IO ()
resetStoreWith :: TVar (Cache Text Text) -> IO ()
resetStoreWith TVar (Cache Text Text)
tvar = do
Cache Text Text
newCache <- Maybe TimeSpec -> IO (Cache Text Text)
forall k v. Maybe TimeSpec -> IO (Cache k v)
C.newCache Maybe TimeSpec
forall a. Maybe a
Nothing
STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar (Cache Text Text) -> Cache Text Text -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Cache Text Text)
tvar Cache Text Text
newCache
makeCacheKey :: Text
-> Algorithm
-> Text
-> Text
-> Text
makeCacheKey :: Text -> Algorithm -> Text -> Text -> Text
makeCacheKey Text
throttleName Algorithm
algo Text
ipZone Text
userKey =
Algorithm -> Text
algoToText Algorithm
algo Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
":" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
throttleName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
":" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
ipZone Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
":" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
userKey
algoToText :: Algorithm
-> Text
algoToText :: Algorithm -> Text
algoToText Algorithm
FixedWindow = Text
"FixedWindow"
algoToText Algorithm
SlidingWindow = Text
"SlidingWindow"
algoToText Algorithm
TokenBucket = Text
"TokenBucket"
algoToText Algorithm
LeakyBucket = Text
"LeakyBucket"
algoToText Algorithm
TinyLRU = Text
"TinyLRU"
parseAlgoText :: Text
-> Parser Algorithm
parseAlgoText :: Text -> Parser Algorithm
parseAlgoText Text
t =
case Text -> Text
Tx.toLower Text
t of
Text
"fixedwindow" -> Algorithm -> Parser Algorithm
forall a. a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Algorithm
FixedWindow
Text
"fixed-window" -> Algorithm -> Parser Algorithm
forall a. a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Algorithm
FixedWindow
Text
"slidingwindow" -> Algorithm -> Parser Algorithm
forall a. a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Algorithm
SlidingWindow
Text
"sliding-window" -> Algorithm -> Parser Algorithm
forall a. a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Algorithm
SlidingWindow
Text
"tokenbucket" -> Algorithm -> Parser Algorithm
forall a. a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Algorithm
TokenBucket
Text
"token-bucket" -> Algorithm -> Parser Algorithm
forall a. a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Algorithm
TokenBucket
Text
"leakybucket" -> Algorithm -> Parser Algorithm
forall a. a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Algorithm
LeakyBucket
Text
"leaky-bucket" -> Algorithm -> Parser Algorithm
forall a. a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Algorithm
LeakyBucket
Text
"tinylru" -> Algorithm -> Parser Algorithm
forall a. a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Algorithm
TinyLRU
Text
"tiny-lru" -> Algorithm -> Parser Algorithm
forall a. a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Algorithm
TinyLRU
Text
_ -> String -> Parser Algorithm
forall a. String -> Parser a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"Unknown algorithm: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Text -> String
Tx.unpack Text
t)