Tuesday, February 10, 2015

A simple rate limiter using Java's DelayQueue

It is rare for me to develop at this level of use case but recently we had to manage a bit of work-load with limited resources that lead to a simplified and light-weight rate limiter using Java's concurrent additions.

Guava's rate limiter is pretty good but I didn't want to include a (fat) dependency on guava just for one class so I wrote a variant based on Java's DelayQueue:

class DelayedEntry implements Delayed {
  long expireAt;
  TimeUnit unit;
  DelayedEntry(long delay, TimeUnit tu) {
   unit = tu;
   setDelay(delay);
  }
  void setDelay(long delay) {
   this.expireAt = System.nanoTime() + unit.toNanos(delay);
  }
  int compareTo(Delayed other) {
   throw new IllegalStateException("Expected single element queue");
  }
  long getDelay(TimeUnit u) {
   return u.convert(expireAt - System.nanoTime(), NANOSECONDS);
  }
}

class RateLimiter {
 DelayQueue<DelayedEntry> queue;
 DelayedEntry token;
 TimeUnit rateUnit;
 AtomicInteger rate;
 RateLimiter(int rateLimit) {
  queue = new DelayQueue<>();
  rateUnit = NANOSECONDS;
  rate = new AtomicInteger(rateLimit);
  token = new DelayedEntry(0, NANOSECONDS);
 }
 boolean acquire(int permits) throws InterruptedException {
  long targetDelay = rateUnit.toNanos(permits) / max(1, rate.get());
  DelayedEntry nextToken = token;
  while (!queue.isEmpty()) {
   nextToken = queue.take();
  }
  assert nextToken != null;
  nextToken.setDelay(targetDelay);
  return queue.offer(token);
 }
}
This implementation isn't exactly a mathematically precise rate limiter that can shape traffic bursts in uniform distributions for volatile rate limits. However, for use-cases involving predictable timings, this works pretty well with minimal and bounded resource usage.

One thing that shouldn't go un-noticed is that we can easily use such utility implementations to instrument collection interfaces such as Iterable and Iterator to decorate existing code base, here's an example:
 public <T> Iterator decorate(final Iterator<T> delegate) {
  return new Iterator() {
   public boolean hasNext() {
    return delegate.hasNext();
   }
   public T next() {
    acquire();
    return delegate.next();
   }
   public void remove() {
    delegate.remove();
   }
  };
 }

I would love to hear reader's opinions about this approach.