Control.Concurrent.TokenBucket

This module provides rate-limiting facilities built on top of the lazy bucket algorithm heavily inspired by "Rate Limiting at Webscale: Lazy Leaky Buckets"

See also Wikipedia's Token Bucket article for general information about token bucket algorithms and their properties.

The TokenBucket type

data TokenBucket

Abstract type containing the token bucket state

newTokenBucket :: IO TokenBucket

Create new TokenBucket instance

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

Arguments

:: TokenBucket

-> Word64

burst-size (tokens)

-> Word64

avg. inverse rate (usec/token)

-> Word64

amount of tokens to allocate

-> IO Bool

True if allocation succeeded

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.

tokenBucketTryAlloc1

Arguments

:: TokenBucket

-> Word64

burst-size (tokens)

-> Word64

avg. inverse rate (usec/token)

-> IO Word64

retry-time (usecs)

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.

tokenBucketWait

Arguments

:: TokenBucket

-> Word64

burst-size (tokens)

-> Word64

avg. inverse rate (usec/token)

-> IO ()

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