{-# LANGUAGE CPP #-}

-- | This module provides rate-limiting facilities built on top of the lazy bucket algorithm heavily inspired by
-- <http://ksdlck.com/post/17418037348/rate-limiting-at-webscale-lazy-leaky-buckets "Rate Limiting at Webscale: Lazy Leaky Buckets">
--
-- See also Wikipedia's <http://en.wikipedia.org/wiki/Token_bucket Token Bucket> article for general information about token bucket algorithms and their properties.
module Control.Concurrent.TokenBucket
    ( -- * The 'TokenBucket' type
      TokenBucket
    , newTokenBucket

      -- * Operations on 'TokenBucket'
      --
      -- | The following operations take two parameters, a burst-size and an average token rate.
      --
      -- === Average token rate
      --
      -- The average rate is expressed as inverse rate in terms of
      -- microseconds-per-token (i.e. one token every
      -- @n@ microseconds). This representation exposes the time
      -- granularity of the underlying implementation using integer
      -- arithmetic.
      --
      -- So in order to convert a token-rate @r@ expressed in
      -- tokens-per-second (i.e. @Hertz@) to microseconds-per-token the
      -- simple function below can be used:
      --
      -- @
      -- toInvRate :: Double -> Word64
      -- toInvRate r = round (1e6 / r)
      -- @
      --
      -- An inverse-rate @0@ denotes an infinite average rate, which
      -- will let token allocation always succeed (regardless of the
      -- burst-size parameter).
      --
      -- === Burst size
      --
      -- The burst-size parameter denotes the depth of the token
      -- bucket, and allows for temporarily exceeding the average
      -- token rate. The burst-size parameter should be at least as
      -- large as the maximum amount of tokens that need to be
      -- allocated at once, since an allocation-size smaller than the
      -- current burst-size will always fail unless an infinite token
      -- rate is used.

    , 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)

-- | Abstract type containing the token bucket state
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

-- getTBData :: TokenBucket -> IO TBData
-- getTBData (TB lbd) = readIORef lbd

#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

-- | Create new 'TokenBucket' instance
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)

-- | Attempt to allocate a given amount of tokens from the 'TokenBucket'
--
-- This operation either succeeds in allocating the requested amount
-- of tokens (and returns 'True'), or else, if allocation fails the
-- 'TokenBucket' remains in its previous allocation state.
tokenBucketTryAlloc :: TokenBucket
                    -> Word64  -- ^ burst-size (tokens)
                    -> Word64  -- ^ avg. inverse rate (usec/token)
                    -> Word64  -- ^ amount of tokens to allocate
                    -> IO Bool -- ^ 'True' if allocation succeeded
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 -- infinitive rate, no-op
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

-- | Try to allocate a single token from the token bucket.
--
-- Returns 0 if successful (i.e. a token was successfully allocated from
-- the token bucket).
--
-- On failure, i.e. if token bucket budget was exhausted, the minimum
-- non-zero amount of microseconds to wait till allocation /may/
-- succeed is returned.
--
-- This function does not block. See 'tokenBucketWait' for wrapper
-- around this function which blocks until a token could be allocated.
tokenBucketTryAlloc1 :: TokenBucket
                     -> Word64     -- ^ burst-size (tokens)
                     -> Word64     -- ^ avg. inverse rate (usec/token)
                     -> IO Word64  -- ^ retry-time (usecs)
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 -- infinite rate, no-op
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

-- | Blocking wrapper around 'tokenBucketTryAlloc1'. Uses 'threadDelay' when blocking.
--
-- This is effectively implemented as
--
-- @
-- 'tokenBucketWait' tb burst invRate = do
--   delay <- 'tokenBucketTryAlloc1' tb burst invRate
--   unless (delay == 0) $ do
--     threadDelay (fromIntegral delay)
--     'tokenBucketWait' tb burst invRate
-- @
tokenBucketWait :: TokenBucket
                -> Word64  -- ^ burst-size (tokens)
                -> Word64  -- ^ avg. inverse rate (usec/token)
                -> 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

-- saturated arithmetic helpers
(∸), (∔) :: 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 () #-}