A rate-limiter in 15 lines of code with cats-effect
By: Olivier Melois
Tags: Technology, Typelevel, Scala

Scala (and its ecosystem) makes it extremely easy to parallelise computations. However, doing so can have an impact. For instance, parallelising the migration of a live database could greatly reduce its availability for the rest of your users.

Wrapping a function in a rate-limiter to constrain the number of times it can be called each second is a simple solution to that problem, allowing you to benefit from parallelism - with the added peace of mind that you won’t accidentally DDoS yourself.

Image from typelevel cats github (Semaphore Cat)

The concept

Given a function of the shape A => F[B] , (potentially to be called many times in parallel) we want to output a function of the exact same shape that has an upper limit on the number of times it can be called each second.

F represents the effect-type in which the computation occurs. Typically, it is going to be some referentially-transparent IO datatype.

The wrapping logic goes as follows :

Ask for permission to run the function through some tracking-construct that counts all the calls currently in flight.

Once we have permission, start a timer for one second

Run the function

When the function completes, if it completed in less than a second, wait for the timer to beep

Notify the tracking-construct we’re done

Return the result

The tracking-construct ensures that only a limited number N of permits are granted, so because we’re enforcing a minimum of one-second before returning the permit, we’re effectively ensuring that the function cannot be called more than N times per second.

Shopping for ingredients

When you have to deal with concurrency/parallelism, cats-effect is one of the best libraries. It offers principled abstractions over datatypes that have the capabilities for concurrency, such as :

cats.effect.IO

monix.eval.Task

zio.ZIO

It also provides a number of concurrency primitives that can be used solve seemingly complex issues in a very clean manner, issues like rate-limiting.

For this particular logic, the components we’ll need are :

cats.effect.Timer which lets you sleep without blocking threads

which lets you without blocking threads cats.effect.concurrent.Semaphore which lets you allocate a number of permits . If a computation asks for a permit and none are available, it'll sleep until one becomes available.

which lets you allocate a number of . If a computation asks for a permit and none are available, it'll sleep until one becomes available. cats.effect.Fiber , which allows for a very simple fork/join API, letting you start computations in the background and join them back into the current "thread" when you need their result.

The implementation

For readability purposes, we’ll be fixing the effect-type to cats.effect.IO , but it could be abstracted away by declaring an abstract effect type F[_] and associated with an implicit cats.effect.Concurrent typeclass instance. This would allow the rate-limiter to work with datatypes provided by Monix and ZIO.

import cats.effect.{IO, Timer}

import cats.effect.concurrent.Semaphore



import scala.concurrent.duration._



/**

* @tparam A the input

* @tparam B the output

* @param T timer associated to the effect type, granting

* the capability to sleep

* @param CS pretty much an ExecutionContext, required for fibers

* @param semaphore: our tracking-construct: holds a number of

* permits distributes-them / gather them back. Instantiated with

* a given number

* @param function: The function we want to limit calls to

* @return a function of the same type

*/

def rateLimited[A, B](

semaphore : Semaphore[IO],

function : A => IO[B])

(implicit T: Timer[IO], CS: ContextShift[IO]) = { input : A => for { // Asking for permission. If no more permit

// is available, this computation will wait

// asynchronously (ie no thread will be blocked)

_ <- semaphore.acquire // Starting a timer counting 1 second in the

// background (start == fork)

timerFiber <- IO.sleep(1.second).start // Calling the function

result <- function(a) // Waiting for the full second to elapse.

// Because the timer is running in the

// background, if the function's execution

// takes more than one second, joining

// will yield instantly. Otherwise, it'll

// block the computation (asynchronously)

// until the full second has passed.

_ <- timerFiber.join // releasing the permit for the next function

// call to pick up.

_ <- semaphore.release } yield result }

Usage

The semaphore has to be passed as parameter of our rateLimited wrapper, because we do not want to create a new one every time we call the wrapped function. However, creating a semaphore is wrapped in an IO . The reason is that Semaphore relies on mutable atomic constructs under the hood, and the act of creating a mutable construct is itself a side effect.

This also means that our rateLimited function has to be passed as a parameter to whatever piece of logic is going to call it, and cannot be referred to globally.

Typically, it looks like this :

import cats.implicits._ // function we want to limit

def myFunction(a : Int) : IO[String] = ??? // big dataset

val myData : List[Int] = ??? val process : IO[List[String]] = for {

semaphore <- Semaphore[IO](200)

rateLimitedFunction = rateLimited(semaphore, myFunction) // Applying our rateLimitedFunction to the full data set

// parallelising as much as possible up to a limit of 200

// calls per second.

allResults <- myData.parTraverse(rateLimitedFunction)

} yield allResults process.unsafeRunSync()

One thing worth noting is that the number of permits the semaphore is initialised with is not a hard limit, and logic could be added to dynamically increase (or decrease) the total number of permits on the fly based on external factors.

Conclusion

cats-effect offers immense simplicity when it comes to solving problems that have to do with concurrency and parallelism. Creating a rate-limiter using the tools it provides is a reasonably trivial task, and is one of many examples where having the knowledge to apply cats-effect more than outweighs the initial investment.