Scala algorithm: Token Bucket Rate Limiter
Published
Algorithm goal
While similar to LeakyBucketRateLimiter, it works a bit differently, and is specifically useful as a network layer algorithm in order rate-limit packets:
Each packet consumes a number of tokens; tokens are added to the bucket at a particular rate, up to a limit and any overflow is discarded.If there are not enough tokens for a packet to consume, then the packet is discarded; otherwise it is passed on.
Design an algorithm to model this, in pure-functional immutable Scala.
Test cases in Scala
assert(
sampleRefillRate.checkPossibleFills(java.time.Duration.ofSeconds(4)) ==
None
)
assert(
sampleRefillRate.checkPossibleFills(java.time.Duration.ofSeconds(5)) ==
Some(5)
)
assert(
sampleRefillRate.checkPossibleFills(java.time.Duration.ofSeconds(10)) ==
Some(10)
)
assert(
sampleRefillRate.checkPossibleFills(java.time.Duration.ofSeconds(11)) ==
Some(10)
)
assert(
sample.requestTokens(1).isEmpty,
"Initially there is no capacity, so we cannot request anything"
)
assert(
sample
.maybeRefill(sampleStartTime.plusSeconds(4))
.flatMap(_.requestTokens(1))
.isEmpty,
"After 4 seconds, we have not refilled, so we cannot request tokens"
)
assert(
sample
.maybeRefill(sampleStartTime.plusSeconds(5))
.flatMap(_.requestTokens(1))
.nonEmpty,
"After 5 seconds, we have refilled, so we can request tokens"
)
assert(
sample
.maybeRefill(sampleStartTime.plusSeconds(5))
.flatMap(_.requestTokens(6))
.isEmpty,
"After 5 seconds, we have only 5 tokens, so cannot request 6"
)
assert(
sample
.maybeRefill(sampleStartTime.plusSeconds(10))
.flatMap(_.requestTokens(6))
.nonEmpty,
"After 10 seconds, we have 10 tokens, so can request 6"
)
assert(
sample
.maybeRefill(sampleStartTime.plusSeconds(10))
.flatMap(_.requestTokens(6))
.flatMap(_.requestTokens(4))
.nonEmpty,
"After 10 seconds, we have 10 tokens, so can request 6, and then 4"
)
assert(
sample
.maybeRefill(sampleStartTime.plusSeconds(10))
.flatMap(_.requestTokens(6))
.flatMap(_.requestTokens(4))
.flatMap(_.requestTokens(1))
.isEmpty,
"After 10 seconds, we have 10 tokens, so can request 6, and then 4, but not 1 more"
)
Algorithm in Scala
41 lines of Scala (compatible versions 2.13 & 3.0).
Explanation
This solution was developed through test-driven development (TDD).
The token bucket is fully described by a combination of the rate, capacity and the time there was a last refill. (this is © from www.scala-algorithms.com)
Using immutable programming, it is even possible to have the rate and capacities vary across time very easily.
Scala concepts & Hints
Drop, Take, dropRight, takeRight
Scala's `drop` and `take` methods typically remove or select `n` items from a collection.
Option Type
The 'Option' type is used to describe a computation that either has a result or does not. In Scala, you can 'chain' Option processing, combine with lists and other data structures. For example, you can also turn a pattern-match into a function that return an Option, and vice-versa!