Adaptive REST api rate limiter

Question

I'm building a java client interacting with a server that returns responses with rate limit info in the headers:

• rate limit per period (for instance 10 per second)
• remaining limit left until next period
• start of next period (epoch time in seconds)

So I built a rate limiter that doesn't follow a fixed rate, but adapts to the info from the responses. The code looks as follows:

import java.util.concurrent.Semaphore;
import java.util.concurrent.atomic.AtomicBoolean;

import javax.ws.rs.core.Response;

public final class BolRateLimiter {

  private static final String HEADER_KEY_LIMIT = "X-RateLimit-Limit"; //$NON-NLS-1$
  private static final String HEADER_KEY_REMAINING = "X-RateLimit-Remaining"; //$NON-NLS-1$
  private static final String HEADER_KEY_RESET = "X-RateLimit-Reset"; //$NON-NLS-1$


  private final Semaphore limiter = new Semaphore(1, true);
  private final AtomicBoolean init = new AtomicBoolean(true);
  private final AtomicBoolean resetting = new AtomicBoolean(false);


  public void limit() {
    try {
      limiter.acquire();

    } catch (final InterruptedException e) {
      Thread.currentThread().interrupt();
    }
  }


  public void read(final Response r) {
    if (init.compareAndSet(true, false)) {
      final String remainstr = r.getHeaderString(HEADER_KEY_REMAINING);
      final int remain = parseInt(remainstr);
      limiter.release(remain);
    }
    if (resetting.compareAndSet(false, true)) {
      final String limitstr = r.getHeaderString(HEADER_KEY_LIMIT);
      final String resetstr = r.getHeaderString(HEADER_KEY_RESET);
      final int limit = parseInt(limitstr);
      final long reset = SECONDS.toMillis(parseLong(resetstr));
      new Thread(() -> {
        try {
          final long now = currentTimeMillis();
          final long delta = reset - now + 100l;
          MILLISECONDS.sleep(delta);
          limiter.drainPermits();
          limiter.release(limit);
          resetting.set(false);

        } catch (final InterruptedException e) {
          currentThread().interrupt();
        }
      }).start();
    }
  }
}

And it's used in for instance the following context:

  @Override
  public BolProcess checkProcess(final String processid) {
--> limiter1.limit(); //
    final Invocation request = endpoints //
        .path("process-status") //$NON-NLS-1$
        .path(processid) //
        .request() //
        .header(AUTHORIZATION, tokenprov.getToken()) //
        .header(ACCEPT, HeaderAcceptValueJson) //
        .buildGet(); //
    try (final Response response = request.invoke()) {
-->   limiter1.read(response);
      if (response.getStatusInfo().getFamily().equals(Response.Status.Family.SUCCESSFUL))
        return response.readEntity(BolProcess.class);
      throw response.readEntity(BolEndpointException.class);
    }
  }

The semaphore starts with 1 permit, so the client can initially send just 1 message before the rate limiter blocks. When the rate limiter reads its first response, the semaphore is instructed to release more permits.

Besides that, the limiter reads from the response header the rate limit as well as the time to reset. It then starts a thread that will reset the semaphore permits once the time to reset passed. The rate limiter then ignores rate limit info from subsequent responses, until thread reset the semaphore.

I need this rate limiter to work in a parallel/concurrent context, like parallel streams. In that context it does not work to read the remaining limit left from a response (except from the first). By the time the rate limiter would read a response and the remaining limit left, that limit is already outdated.

So am I overlooking anything here? I'd be happy to get your feedback. Thanks in advance!

Answer 1

There are a few holes that I see:

• The first thread can get into the init.compareAndSet() block and then stall out, causing later threads to not progress until the first thread makes it to the release line. If the first request fails for any reason, no other requests will ever be sent out.

• Any thread can stall out while in the resetting.compareAndSet() block. If it sleeps past a time limit boundary, threads will be blocked until it wakes back up.

• Responses can come back out of order, so the response that goes into the resetting.compareAndSet() block may not have the correct number of permitted calls remaining.

• If a thread acquires a permit and then sleeps until the semaphore is resized, it may cause you to unexpectedly hit the rate limit.

• There is always a way you can unexpectedly get a rate limit failure. A thread acquire a valid semaphore permit, then sleep until after all the quota has been used for the cycle after the current one, and then the thread wakes up and sends its request during that cycle. There is no way to protect against this that I am aware of.

It is reasonable for the implementation to attempt to respect the rate limit by not sending extraneous requests, but the implementation should be able to handle rate limiting failures more cleanly than throwing the same exception as any other error response. Requests that hit the rate limit should be retried. There should also be monitoring to ensure you aren't backing up on your end because you're generating more requests than the service API can handle over time. That may be in the code, but would probably be better in tooling.

Other: I don't know what a "BolRateLimiter" is. Avoid acronyms where possible to enhance readability.

Strive for maximally readable code. response would read more cleanly than r.

When indicating a numeric value is a long, prefer L to l for readability reasons.

Random musing: You may consider extending Semaphore into DrainingSemaphore, which would give you access to reduce permits(). You could then try something like the code below (100% untested):

private final DrainingSemaphore semaphore = new DrainingSemaphore(Integer.MAX_VALUE, true);
private final Lock lock = new ReentrantLock();
private long startOfNextPeriod = 0;

public void read(final Response response) {
    long responseStartOfNextPeriod = ..;
    int responseRequestsRemaining = ..;

    try {
        // can't let multiple responses update the semaphore
        lock.lockInterruptibly();

        // first request of a new period
        if (responseStartOfNextPeriod > startOfNextPeriod) {
            startOfNextPeriod = responseStartOfNextPeriod;
            semaphore.drainPermits();
            semaphore.release(responseRequestsRemaining);
            return;
        }

        // request from the current period
        if (responseStartOfNextPeriod == startOfNextPeriod) {
            // case: permit from prior period, call happens this period
            // drain semaphore permits down to the correct value
            semaphore.drainToSize(responseRequestsRemaining);
        }
    } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
    } finally {
        lock.unlock();
    }

    // Requests from a prior period can be ignored
}

Answer 2

With the input from Eric Stein's answer, I updated the rate limiter class. It may not completely cover all issues Eric mentioned, but at least should be an improvement.

The main difference with the original code is replacing the two AtomicBooleans with a ConcurrentHashMap.KeySetView. In the original code the AtomicBoolean resetting would not allow starting another reset, not even for a next (or previous) period. The ConcurrentHashmap.KeySetView resetTimes allows starting multiple resets, as long as the reset for a certain period is not scheduled yet.

public final class BolRateLimiter {

  private static final String KEY_INIT_RELEASE = ""; //$NON-NLS-1$
  private static final String HEADER_KEY_LIMIT = "X-RateLimit-Limit"; //$NON-NLS-1$
  private static final String HEADER_KEY_RESET = "X-RateLimit-Reset"; //$NON-NLS-1$
  private static final String HEADER_KEY_REMAINING = "X-RateLimit-Remaining"; //$NON-NLS-1$


  private final Semaphore semaphore = new Semaphore(1, true);
  private final Set<String> resetTimes = ConcurrentHashMap.newKeySet();


  public void limit() {
    try {
      semaphore.acquire();

    } catch (final InterruptedException e) {
      currentThread().interrupt();
    }
  }


  public void read(final Response response) {
    if (resetTimes.add(KEY_INIT_RELEASE)) {
      final String remainstr = response.getHeaderString(HEADER_KEY_REMAINING);
      final int remain = parseInt(remainstr);
      semaphore.release(remain);
    }

    final String newResetTimeStr = response.getHeaderString(HEADER_KEY_RESET);
    final long newResetTime = SECONDS.toMillis(parseLong(newResetTimeStr));
    if (resetTimes.add(newResetTimeStr)) {
      final String limitStr = response.getHeaderString(HEADER_KEY_LIMIT);
      final int limit = parseInt(limitStr);
      final Thread t = new Thread(() -> {
        try {
          final long now = currentTimeMillis();
          final long delta = newResetTime - now + 100L;
          MILLISECONDS.sleep(delta);
          semaphore.drainPermits();
          semaphore.release(limit);

        } catch (final InterruptedException e) {
          currentThread().interrupt();

        } finally {
          resetTimes.remove(newResetTimeStr);
        }
      });
      t.setDaemon(true);
      t.start();
    }
  }
}
```