This week I had to rate limit the calls I was making to a web service. In the past I’ve done this successfully with a Timer, but it’s still a bit cumbersome to schedule/unschedule the timer to keep it from ticking when I’m not making requests, and control everything from its callback. Also, these calls were being done on background threads and I did not want the dependency on the RunLoop.

I could hack my way out of it by carefully placing a Thread.sleep() but… that smells.

Token Bucket to the rescue

There’s a simple rate limiting algorithm usually used in routers for packet switching called a token bucket. I’ve actually never used it before but it was so simple I still remembered it from my Computer Networks classes.

The idea is you have a bucket which gets filled with tokens at a predefined time interval. To do some work, each actor in the system has to take tokens from the bucket. If there are enough tokens for the desired action the actor removes the tokens and proceeds, if not he waits. Simple right?

It’s a simple concept to grasp and surprisingly powerful. You define your rate limit by controlling the rate at which the bucket is refilled which can change at runtime due to some other factors. And you can assign weights to each of the actions by requiring them to consume more tokens.

To build this there’s a lesser known lock in Foundation I’ve grown to love that would be a great fit.

Enter NSCondition

An NSCondition lock can be used to wait until a certain condition is met. In this case, until the desired tokens are available.

It works a bit differently from a traditional lock, but the documentation does a good job of explaining it. So I won’t reproduce it here.

You’ll also see how it works in the next section.

Building the token bucket

First off we need a class to represent the bucket. And we need some instance variables to hold the bucket capacity, number of tokens in the bucket, replenishing interval… and of course, the NSCondition object.

public class TokenBucket: NSObject {
    public let capacity: Int
    public private(set) var replenishingInterval: TimeInterval
    public private(set) var tokensPerInterval: Int
    public private(set) var tokenCount: Int
    
    private let condition = NSCondition()
}

Now we’ll need two methods to actually consume tokens and to refill the bucket. Let’s start with the method that refills the bucket.

We don’t want it to depend on timers or the runloop, so we can just store the date it was last replenished and use that to determine how many tokens we should add based on tokensPerInterval and replenishingInterval.

private var lastReplenished: Date
func replenish() {
    condition.lock()
    let ellapsedTime = abs(lastReplenished.timeIntervalSinceNow)
    if  ellapsedTime > replenishingInterval {
        let ellapsedIntervals = Int((floor(ellapsedTime / Double(replenishingInterval))))
        tokenCount = min(tokenCount + (ellapsedIntervals * tokensPerInterval), capacity)
        lastReplenished = Date()
        condition.signal()
    }
    condition.unlock()
}

Easy peasy. Whenever we need to get the number of current tokens we need to make sure this method gets called. So we change our public API to take this into account.

private var _tokenCount: Int
public var tokenCount: Int {
    replenish()
    return _tokenCount
}

Finally, the method that consumes the tokens also needs to call replenish before testing the number of existing tokens.

func wait(until limitDate: Date, for tokens: Int) -> Bool {
    replenish()
    
    condition.lock()
    defer {
        condition.unlock()
    }
    while _tokenCount < tokens {
        if limitDate < Date() {
            return false
        }
        DispatchQueue.global().async {
            self.replenish()
        }
        condition.wait(until: Date().addingTimeInterval(0.2))
    }
    _tokenCount -= tokens
    return true
}

Here we immediately call replenish() outside of the lock, then we lock the condition and test if we can remove the desired number of tokens. If the test passes we remove the tokens and return, if not we schedule a call to replenish() on another thread (we can’t call it directly since we’re holding the lock) and wait on the condition for 0.2 seconds to test again. While we’re waiting the condition is unlocked automatically for us, allowing other threads to acquire it. We do this until limitDate.

That’s it. For the entire class check the accompanying repo on GitHub.

Back to posts