    {-# 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
    $cshowsPrec :: Int -> TBData -> ShowS
    showsPrec :: Int -> TBData -> ShowS
    $cshow :: TBData -> String
    show :: TBData -> String
    $cshowList :: [TBData] -> ShowS
    showList :: [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
        now <- IO Word64
    getPosixTimeUsecs
        lbd <- newIORef $! TBData 0 now
        evaluate (TB 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 a. a -> IO a
    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 a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return Bool
    False
    tokenBucketTryAlloc (TB IORef TBData
    lbref) Word64
    burst Word64
    invRate Word64
    alloc = do
        now <- IO Word64
    getPosixTimeUsecs
        atomicModifyIORef' lbref (go 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 a. a -> IO a
    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
        now <- IO Word64
    getPosixTimeUsecs
        atomicModifyIORef' lbref (go 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
        delay <- TokenBucket -> Word64 -> Word64 -> IO Word64
    tokenBucketTryAlloc1 TokenBucket
    tb Word64
    burst Word64
    invRate
        unless (delay == 0) $ do
            threadDelay (fromIntegral delay)
            tokenBucketWait tb burst 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 (∔) #-}