Rate limiter utility class

Question

I had a problem in my system where the class that sends messages to other systems sometimes gets flooded and so I needed to add rate limiting. I threw together something that would work, but I would really like to see what I can do to improve and/or clean it up.

This class does not do the actual rate limiting, it only provides timing guidelines on when the rate limited class should perform its next action.

package edu.stsci.util

import scala.collection.mutable
/**
  * Used to rate limit a repeated action based on three parameters. This class doesn't
  * actually control the action, but tells the calling class how long to wait before
  * performing the action again by calling `nextWait`. Caller also indicates every time
  * an action is performed using `execute`.
  *
  * @param minSpacing The minimum amount of time between actions for burst throughput (in ms)
  * @param maxPer The number of actions that can happen during any one `movingPeriod` time window
  * @param movingPeriod The length of the moving time window for controlling sustained throughput (in ms)
  */
class RateLimiter(val minSpacing: Int, val maxPer: Int, val movingPeriod: Int) {
  val periodRate: Double = maxPer / RateLimiter.SUBWINDOW_COUNT.toDouble
  val periodInterval = movingPeriod / RateLimiter.SUBWINDOW_COUNT
  val sustainedInterval: Long = movingPeriod / maxPer

  var last: Long = 0

  var totalCount = 0

  var movingData = new mutable.Queue[Int]()
  var currentCount = 0
  var currentEnd = 0l

  /**
    * `execute` must be called every time an action occurs to update last
    * and statistics
    */
  def execute() = {
    val now = System.currentTimeMillis()
    last = now

    if (currentEnd == 0l) {
      initializeWindows()
    }
    else {
      while (now > currentEnd) {
        dropOnePeriod()
        queueLastPeriod()
      }

      currentCount += 1
      totalCount += 1
    }
  }

  def initializeWindows(): Unit = {
    currentEnd = System.currentTimeMillis() + periodInterval
    currentCount += 1
    totalCount += 1
  }

  def dropOnePeriod(): Unit = {
    if (movingData.size >= 20) {
      val count = movingData.dequeue()
      totalCount -= count
    }
  }

  def queueLastPeriod(): Unit = {
    movingData.enqueue(currentCount)
    currentCount = 0
    currentEnd += periodInterval
  }

  /**
    * Calculate the number of ms that the calling class should wait
    * before executing the next action
    * @return ms that should be waited
    */
  def nextWait(): Long = {
    val now = System.currentTimeMillis()
    val wait = calculateWait() - (now - last)

    if (wait < 0) 0
    else wait
  }

  def calculateWait(): Long = {
    if (minSpacing >= sustainedInterval) minSpacing
    else if (totalCount <= maxPer / 4) minSpacing
    else if (currentCount <= periodRate) (minSpacing + sustainedInterval) / 2l
    else if (totalCount <= maxPer * 9/10) sustainedInterval
    else sustainedInterval * 6 / 5
  }
}

object RateLimiter {
  val SUBWINDOW_COUNT = 20
}

And the tests.

package edu.stsci.util

import org.junit.runner.RunWith
import org.scalatest.{BeforeAndAfter, FunSuite, Matchers}
import org.scalatest.junit.JUnitRunner

@RunWith(classOf[JUnitRunner])
class RateLimiterTest extends FunSuite with BeforeAndAfter with Matchers {
  test("simple interval case") {
    val testObject = new RateLimiter(50, 100, 100)

    testObject.nextWait() should be (0)
    testObject.execute()

    for(count <- 1 to 100) {
      val next = testObject.nextWait()
      println("[RateLimiterTest.simple interval case] count " + count + ", wait ", next)
      next should be > 5l
      Thread.sleep(next)
      testObject.execute()
    }
  }

  test("simple flood case") {
    val testObject = new RateLimiter(25, 10, 1000)

    testObject.nextWait() should be (0)
    testObject.execute()

    val start = System.currentTimeMillis()
    for (count <- 0 to 20) {
      val next = testObject.nextWait()
      println("[RateLimiterTest.simple flood case] count " + count + ", wait ", next)
      Thread.sleep(next)
      testObject.execute()
    }

    val elapsed = System.currentTimeMillis() - start
    println("[RateLimiterTest.simple flood case] elapsed: " + elapsed)

    elapsed should be >= 2000l
  }

  test("burst case") {
    val testObject = new RateLimiter(20, 100, 5000)

    for (loopCount <- 1 to 5) {
      val start = System.currentTimeMillis()
      val zero: Long = testObject.nextWait()
      zero should be (0l)
      testObject.execute()

      for (count <- 1 to 4) {
        val next = testObject.nextWait()
        next should be >= 15l
        println("[RateLimiterTest.burstCase] count " + count + ", wait " + next)
        Thread.sleep(next)
        testObject.execute()
      }

      val elapsed = System.currentTimeMillis() - start
      println("[RateLimiterTest.simple flood case] elapsed: " + elapsed)
      elapsed should be > 80l

      Thread.sleep(250)
    }
  }
}

Answer 1

Slow tests

The tests run slow because of the Thread.sleep calls. It would be better to inject a "time provider" in the implementation, so that you can fake System.currentTimeMillis() and Thread.sleep(next) and keep the tests fast.

Cause and effect

The connection between the numbers used in the fixture setup step and in the assertions is not obvious. Take for example the assertion elapsed should be > 80l in test("burst case"). It's not obvious where 80 comes from and how it's related to the fixture setup. The consequence is that if something changes in the implementation, it won't be obvious how to change the test.

I'm guessing these are the lines that contribute to the threshold of 80:

    val testObject = new RateLimiter(20, 100, 5000)

    for (...) {
      // ...

      for (count <- 1 to 4) {
        // ...
      }

      // ...
      elapsed should be > 80l

I'm guessing that elapsed should be > 20 * 4. It would be good to make the connection obvious by using more local variables with descriptive names, for example:

    val minSpacing = 20
    val testObject = new RateLimiter(minSpacing, 100, 5000)

    for (...) {
      // ...

      val samples = 4
      for (count <- 1 to samples) {
        // ...
      }

      // ...
      elapsed should be > (minSpacing * samples).toLong

Cosmetic issues

A minor thing, zero is redundant here:

  val zero: Long = testObject.nextWait()
  zero should be (0l)

---

Instead of this:

  println("..., wait ", next)

I'm guessing you meant to write this:

  println("... , wait " + next)

Answer 2

I'm going to go, maybe, slightly off topic for the site and review the need for this code to be written in the first place... before you go too far down this particular rabbit hole :)

Given your use-case of sending messages to other systems, have you first considered existing solutions?

In particular, it sounds like some form of reactive streams might be a particularly good fit. Specifically:

• RxNetty: Which uses Netty and the Rx library. You'd want to use it via RxScala

• Akka Streams: Which is similar to Rx, but built on Akka, and plays well with Akka's own IO library (also backed by Netty)

Both have good stories to tell about handling back-pressure, including static or dynamic rate limiting in addition to other strategies.