Provide snapshot of numTokens / duration until next refill

[bbeck]

Provide snapshot of numTokens / duration until next refill

[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. =)