Example for rate limiting of client-side api calls / runnables

Vyshibala.java

/*
 *  Copyright 2017 Wladislaw Mitzel
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

import java.util.Date;
import java.util.Random;
import java.util.concurrent.DelayQueue;
import java.util.concurrent.Delayed;
import java.util.concurrent.TimeUnit;

import lombok.Data;
import lombok.extern.slf4j.Slf4j;

/**
 * This is an example for rate limiting of api calls as a client. Vyshibala is
 * the russian word for bouncer. It will only let x requests be executed in y
 * seconds. The other executions have to wait until they can be executed.
 * 
 * We use a {@link DelayQueue} which is initially populated with x tokens.
 * Before execution of the {@link #run(Runnable)} one token is consumed. The
 * {@link DelayQueue} will make us wait if ther eis no consumable token, yet.
 * After the runnable has been executed a new token is generated with a delay of
 * y seconds.
 * 
 */
@Slf4j
public class Vyshibala {

	@Data
	private class Token implements Delayed {

		/** the point in time at which this token can be consumed */
		private Long earliestExecution;

		/**
		 * Create a new token that will be consumable in
		 * <code>earlistesExecutionInMs</code> milliseconds.
		 * 
		 * @param earlistesExecutionInMs
		 *            the time in milliseconds one has to wait until this token
		 *            can be consumed.
		 */
		public Token(Long earlistesExecutionInMs) {
			this.earliestExecution = new Date().getTime() + earlistesExecutionInMs;
		}

		@Override
		public int compareTo(Delayed o) {
			log.trace("comparing {} with {}", this, o);
			return this.earliestExecution.compareTo(((Token) o).earliestExecution);
		}

		@Override
		public long getDelay(TimeUnit unit) {
			long delay = this.earliestExecution - new Date().getTime();
			log.trace("delay until next execution: {}", delay);
			return unit.convert(delay, TimeUnit.MILLISECONDS);
		}
	}

	/** our queue holds the tokens - one for each execution that is allowed */
	private DelayQueue<Token> queue = new DelayQueue<>();
	/** new tokens use this delay until they are consumable */
	private int perXSeconds;

	public Vyshibala(int numberOfRequests, int perXSeconds) {
		super();
		this.perXSeconds = perXSeconds;
		/*
		 * initially fill the queue with numberOfRequests tokens and use a delay
		 * of 0 so that they are consumable instantly.
		 */
		for (int i = 0; i < numberOfRequests; i++) {
			Token e = new Token(0L);
			this.queue.put(e);
		}
	}

	/**
	 * Call this method for rate limiting your calls to your {@link Runnable}.
	 * It will wait until a consumable token is available in the {@link #queue}.
	 * 
	 * @param r
	 *            The runable you want to execute.
	 */
	public void run(Runnable r) {
		/*
		 * poll the queue for a consumable token. If there is none we wait until
		 * one is availabe.
		 */
		while (this.queue.poll() == null) {
			try {
				Thread.sleep(100); // this period could be configurable...
			} catch (InterruptedException e) {
				e.printStackTrace();
			}
		}
		/*
		 * now execute the runnable
		 */
		try {
			r.run();
		} finally {
			/*
			 * and after execution insert a new token that will be consumable in
			 * this.perXSeconds seconds.
			 */
			Token e = new Token((long) this.perXSeconds * 1_000);
			this.queue.put(e);
		}
	}

	/**
	 * A small test runner that creates a {@link #Vyshibala(int, int)}, runs
	 * indefinetly and and generates some log output.
	 */
	public static void main(String[] args) throws InterruptedException {
		Vyshibala v = new Vyshibala(15, 20);
		Random r = new Random();
		while (true) {
			v.run(() -> {
				log.info("calling api");
			});
			Thread.sleep(r.nextInt(2_000));
		}
	}
}