Provide snapshot of numTokens / duration until next refill

[mmadson]

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!

[bbeck]

All of the mutation methods in TokenBucket are synchronized. So I believe you can just synchronize on the instance and call whatever methods you need to and be guaranteed that nothing will change in the internal state of the bucket during that block. Something like this might work...

TokenBucket tokenBucket = ...;

synchronized(tokenBucket) {
  boolean consumedTokens = tokenBucket.tryConsume(50);
  long capacity = tokenBucket.getCapacity();
  long numTokens = tokenBucket.getNumTokens();
  long duration = tokenBucket.getDurationUntilNextRefill(TimeUnit.SECONDS);

  response.setHeader("X-RateLimit-Capacity", capacity);
  response.setHeader("X-RateLimit-Cost", 50);
  response.setHeader("X-RateLimit-Remaining", numTokens);
  response.setHeader("X-RateLimit-Reset", duration);

  if (consumedTokens) {
    response.setStatusCode(200);
  } else {
    response.setStatusCode(429);
  }
}

Keep in mind that as long as you're in the synchronized block no other threads can operate on the bucket, so make sure you put as little code as possible in there.

[mmadson]

getNumTokens can refill the bucket following the tryConsume resulting in the following inconsistent headers:

Capacity: 100
Cost: 50
Remaining: 100
...

[bbeck]

Ahh, I see your point. Your use case is quite interesting -- it's different from how I've used token buckets in the past in that you're exposing a lot more information outside of the bucket. The challenge is that getting that information either complicates the API and makes it more difficult for "simple" usage, or your more complex case get inaccurate information.

I'm really hesitant to increase complexity of the API especially given that normal uses of a token bucket are in time critical sections of code where you typically don't want to be doing a memory allocation on every request. I do think there might be a way for you to achieve what you're trying to do, it's not ideal though.

Basically you can encapsulate both the token bucket and a special refill strategy into your own class. The refill strategy is special in that while it can wrap any existing refill strategy it will also offer the ability to pause and unpause itself. While it's paused, its refill method always returns 0 tokens and never delegates to the underlying refill strategy. It might look something like this.

class RateLimiter {
  private final PausableRefillStrategy refillStrategy;
  private final TokenBucket bucket;

  RateLimiter(long capacity, long initialTokens, TokenBucket.RefillStrategy refillStrategy) {
    this.refillStrategy = new PausableRefillStrategy(refillStrategy);
    this.bucket = TokenBuckets.builder()
      .withCapacity(capacity)
      .withInitialTokens(initialTokens)
      .withRefillStrategy(this.refillStrategy)
      .build();
  }

  synchronized Response consume(long numTokens) {
    // Execute the first tryConsume while the refill strategy is running.
    boolean consumedTokens = bucket.tryConsume(numTokens);
    refillStrategy.pause();
    try {
      // ... build your response object here ...
    } finally {
      refillStrategy.unpause();
    }

    return response;
  }

  private static final class PausableRefillStrategy implements TokenBucket.RefillStrategy{
    private final TokenBucket.RefillStrategy delegate;
    private volatile boolean isPaused = false;

    PausableRefillStrategy(TokenBucket.RefillStrategy delegate) {
      this.delegate = delegate;
    }

    public long refill() {
      if (isPaused) return 0;
      return delegate.refill();
    }

    public long getDurationUntilNextRefill(TimeUnit unit) {
      return delegate.getDurationUntilNextRefill(unit);
    }

    public void pause() {
      isPaused = true;
    }

    public void resume() {
      isPaused = false;
    }
  }
}

[mmadson]

I totally see where you're coming from, not wanting to over-complicate and slow down the API to support our rather obscure use case. What do you think about adding a snapshot accessor instead, something like:

TokenInfoSnapshot getTokenInfoSnapshot()

Each time refill() is called from tryConsume(), consume() or the new refill() you can save off the time until next refill as well as the current num tokens into separate fields. Then when I call getTokenInfoSnapshot(), you can pack those fields into a new immutable TokenInfoSnapshot object. This way you avoid the overhead of object allocation for the common case, yet provide a much simpler means for me to get a consistent view of the info I need.

So my usage would become:

synchronize(tokenBucket) {
  tokenBucket.tryConsume(cost);
  TokenInfoSnapshot snapshot = tokenBucket.getTokenInfoSnapshot();
  response.setHeader("X-RateLimit-Capacity", snapshot.getCapacity()); // capacity doesn't actually need to be in snapshot
  response.setHeader("X-RateLimit-Cost", cost);
  response.setHeader("X-RateLimit-Remaining", snapshot.getNumTokens());
  response.setHeader("X-RateLimit-Reset", snapshot.getDurationUntilNextRefill(TimeUnit.SECONDS));
}

I really appreciate all the assistance you've provided. =)