Ratelimit with cats-effect

ratelimited-cats.scala

package be.icteam.sample

import cats.effect._
import cats.implicits._
import cats.effect.concurrent.Semaphore

import java.util.concurrent.Executors
import scala.concurrent.ExecutionContext


object Main extends IOApp {

  val es = Executors.newFixedThreadPool(5)
  val ec = ExecutionContext.fromExecutor(es)
  implicit val cs: ContextShift[IO] = IO.contextShift(ec)

  override def contextShift: ContextShift[IO] = cs

  val rnd = new scala.util.Random()

  def work(x: Int): IO[Int] = IO {
    val cost = rnd.nextInt(4000)
    println(s"$x sleeping for ${cost}ms.")
    Thread.sleep(cost)
    println(s"$x done")
    x
  }

  def rateLimited[A, B](limit: Int)(fn: A => IO[B])(items: List[A]): IO[List[B]] = {
    for {
      s <- Semaphore[IO](limit)
      results <- items.parTraverse(x => s.withPermit(fn(x)))
    } yield results
  }


  override def run(args: List[String]): IO[ExitCode] = {

    val items: List[Int] = Seq.range(1, 10).toList
    val resultsIO: IO[List[Int]] = rateLimited(20)(work)(items)
    println(resultsIO.unsafeRunSync())

    IO(ExitCode.Success)
  }


}