bbeck/token-bucket

Token bucket algorithm for rate-limiting

Introduction

This library provides an implementation of a token bucket algorithm which is useful for providing rate limited access to a portion of code. The implementation provided is that of a "leaky bucket" in the sense that the bucket has a finite capacity and any added tokens that would exceed this capacity will "overflow" out of the bucket and be lost forever.

In this implementation the rules for refilling the bucket are encapsulated in a provided RefillStrategy instance. Prior to attempting to consume any tokens the refill strategy will be consulted to see how many tokens should be added to the bucket

We use Travis CI for build verification.

See also:

• Wikipedia - Token Bucket
• Wikipedia - Leaky Bucket

Usage

Using a token bucket is incredibly easy and is best illustrated by an example. Suppose you have a piece of code that polls a website and you would only like to be able to access the site once per second:

// Create a token bucket with a capacity of 1 token that refills at a fixed interval of 1 token/sec.
TokenBucket bucket = TokenBuckets.builder()
  .withCapacity(1)
  .withFixedIntervalRefillStrategy(1, 1, TimeUnit.SECONDS)
  .build();

// ...

while (true) {
  // Consume a token from the token bucket.  If a token is not available this method will block until
  // the refill strategy adds one to the bucket.
  bucket.consume(1);

  poll();
}

As another example suppose you wanted to rate limit the size response of a server to the client to 20 kb/sec but want to allow for a periodic burst rate of 40 kb/sec:

// Create a token bucket with a capacity of 40 kb tokens that refills at a fixed interval of 20 kb tokens per second
TokenBucket bucket = TokenBuckets.builder()
  .withCapacity(40960)
  .withFixedIntervalRefillStrategy(20480, 1, TimeUnit.SECONDS)
  .build();

// ...

while (true) {
  String response = prepareResponse();

  // Consume tokens from the bucket commensurate with the size of the response
  bucket.consume(response.length());

  send(response);
}

Maven Setup

The token bucket library is distributed through maven central. Just include it as a dependency in your pom.xml.

<dependency>
    <groupId>org.isomorphism</groupId>
    <artifactId>token-bucket</artifactId>
    <version>1.6</version>
</dependency>

License

Copyright 2012-2015 Brandon Beck Licensed under the Apache Software License, Version 2.0: http://www.apache.org/licenses/LICENSE-2.0.

token-bucket's Issues

Support programmatically refilling token buckets

Hi again,

I was wondering if it would be possible to add support for programmatically refilling token buckets with X number of tokens prior to the refill reset time.

For example:

TokenBucket t = ...;
t.refill(10)

This would add at most 10 tokens to token bucket, t, or (capacity - 10), whichever is less.

Our use case for this is a bit complicated but basically each user in our system is associated with a token bucket and we'd like an administrator to be able to refund individual users just enough points to issue one off requests (which cost tokens), under certain exceptional conditions.

Thanks!

Should the yielding sleep strategy call Thread.yield()?

Would it be better for TokenBuckets.YIELDING_SLEEP_STRATEGY to call Thread.yield() here rather than sleeping for one nanosecond?

Performance problem

Hi, if use your token bucket implementation for types per second, it's low performance, for example this test will produce just about 275 TPS from the defined maximum 300 TPS. It's just 91% of the maximum performance.

package org.frymock.util;

import org.isomorphism.util.TokenBucket;
import org.isomorphism.util.TokenBuckets;
import org.junit.Before;
import org.junit.Ignore;
import org.junit.Test;

import java.util.concurrent.TimeUnit;

/**
 * @author Guram Savinov
 */
public class TokenBucketTest {

    private static final int MAX_TPS = 300;

    private TokenBucket sut;

    @Before
    public void setUp() {
        sut = TokenBuckets.newFixedIntervalRefill(MAX_TPS, 1, 1000000 / MAX_TPS, TimeUnit.MICROSECONDS);
    }

    @Test
    public void checkTime() {
        int cnt = 0;
        long startTime = System.currentTimeMillis();
        while(true) {
            cnt++;
            if (cnt % MAX_TPS == 0) {
                sut.consume(MAX_TPS);
                long time = System.currentTimeMillis() - startTime;
                System.out.println("current time: " + time + ", counter = " + (cnt / MAX_TPS) + ", average TPS = " + (cnt * 1000.0 / time));
            }
        }
    }

}

I think that timer precision cause this low performance.
TokenBucket#consume() calls Uninterruptibles.sleepUninterruptibly(1, TimeUnit.NANOSECONDS) but it sleep much more than 1 nanoseond.
Tested on the GNU/Linux, kernel: 3.11.0-20

Add accessors for next refill timestamp, number tokens remaining in bucket and capacity

Hi there,

I was wondering if you could add the following functionality:

// returns Unix UTC timestamp for when the next refill will occur
long getNextRefillTimestamp();

// returns the capacity that the token bucket was initialized with
long getCapacity();

// returns the number of tokens in the bucket (size?) for the current interval
long getNumTokensRemaining();

Thanks!

Issue should be ignored and deleted

getNumTokens() not returning capacity() initially

TokenBuckets.builder().withCapacity(20).withFixedIntervalRefillStrategy(10, 1, TimeUnit.SECONDS).build();

then when I call tokenBucket.getNumTokens()

it returns 10, instead of 20?

FixedIntervalRefillStrategy assumes period to be starting at 0.

There is problem with initializing lastRefillTime to negative periodDurationInNanos. Say periodDurationInNanos is 1 (period 1, time unit nanoseconds), in that case it will initialize lastRefillTime to -1.

Now lets look at the period calculation:

public synchronized long refill()
  {
    long now = ticker.read();  // [1]
    ...
    long numPeriods = Math.max(0, (now - lastRefillTime) / periodDurationInNanos); // [2]
    ...

[1] will return a big number (say 1188589593996144), and [2] will simple add 1 to the big number making it 1188589593996145. (See the problem yet). If you divide this by periodDurationInNanos (which was 1 btw). You get huge amount of tokens prebuilt.

Why does not test catch this? MockTicket always starts at 0. But System ticker won't.

Provide snapshot of numTokens / duration until next refill

We're having some trouble generating our rate limit headers with the current TokenBucket API.

For reference, we are trying to generate headers similar to the following:

X-RateLimit-Capacity  // max number of tokens for a user during a given period
X-RateLimit-Cost // the number of tokens that the current request costs
X-RateLimit-Remaining // the number of tokens remaining for the current user
X-RateLimit-Reset // the next timestamp when the user's remaining token buckets will be reset to capacity

The problem is that I need to make multiple calls on the tokenBucket to collect all the information for the headers and between any 2 calls the state of the remaining tokens / and duration until next refill may have changed. Trying to deliver a consistent set of headers that is accurate to the user seems impossible without having the consume/tryConsume methods return a snapshot of the state of the required header information at the time the tokens were consumed:

TokenBucketInfo response = tockenBucket.tryConsume(50);
if(response.isSuccessful()) {
   response.getCapacity(); 
   response.getNumTokens(); // this would not try to refill the token bucket, instead it should report the number of tokens in the token bucket at the point in time immediately following the consume.
   response.getDurationUntilNextRefill(someTimeUnit); // this would report the duration until the next refill at the time the consume occurred.
}

Basically any method that can change the state of the token bucket consume/refill/etc. should return a snapshot of the information for generating accurate headers. Also exposing a method that returns a consistent snapshot would also be handy, should we ever decide to create a resource that simply returns a user's token information without consuming any tokens.

Thanks!

To make TokenBucketImpl serializable

Would that be possible? I'm trying to manage the buckets on a shared cache (Hazelcast). To do that TokenBucketImpl must be serializable.