Add rate limiter primitives

core/src/main/scala/ox/resilience/GenericRateLimiter.scala

@@ -0,0 +1,146 @@
+package ox.resilience
+
+import java.util.concurrent.Semaphore
+import GenericRateLimiter.*
+import ox.resilience.GenericRateLimiter.Strategy.Blocking
+import ox.*
+
+/** Rate limiter which allows to pass a configuration value to the execution. This can include both runtime and compile time information,
+  * allowing for customization of return types and runtime behavior. If the only behavior needed is to block or drop operations, the
+  * `RateLimiter` class provides a simpler interface.
+  */
+case class GenericRateLimiter[Returns[_[_]] <: Strategy[_]](
+    executor: Executor[Returns],
+    algorithm: RateLimiterAlgorithm
+):
+
+  import GenericRateLimiter.Strategy.given
+
+  /** Limits the rate of execution of the given operation with a custom Result type
+    */
+  def apply[T, Result[_]](operation: => T)(using Returns[Result]): Result[T] =
+    executor.schedule(algorithm, operation)
+    executor.execute(algorithm, operation)
+  end apply
+end GenericRateLimiter
+
+object GenericRateLimiter:
+
+  type Id[A] = A
+
+  /** Describe the execution strategy that must be used by the rate limiter in a given operation. It allows the encoding of return types and
+    * custom runtime behavior.
+    */
+  sealed trait Strategy[F[*]]:
+    def run[T](operation: => T): F[T]
+
+  object Strategy:
+    sealed trait Blocking[F[*]] extends Strategy[F]
+    sealed trait Dropping[F[*]] extends Strategy[F]
+    sealed trait BlockOrDrop[F[*]] extends Strategy[F]
+
+    case class Block() extends Blocking[Id] with BlockOrDrop[Id]:
+      def run[T](operation: => T): T = operation
+
+    case class Drop() extends Dropping[Option] with BlockOrDrop[Option]:
+      def run[T](operation: => T): Option[T] = Some(operation)
+
+    given Blocking[Id] = Block()
+    given Dropping[Option] = Drop()
+  end Strategy
+
+  /** Determines the policy to apply when the rate limiter is full. The executor is responsible of managing the inner state of the algorithm
+    * employed. In particular, it must ensure that operations are executed only if allowed and that the algorithm is updated.
+    */
+  trait Executor[Returns[_[_]] <: Strategy[_]]:
+
+    /** Performs any tasks needed to delay the operation or alter the execution mode. Usually, this will involve using `acquire` or
+      * `tryAcquire` methods from the algorithm and taking care of updating it.
+      */
+    def schedule[T, Result[*]](algorithm: RateLimiterAlgorithm, operation: => T)(using Returns[Result]): Unit
+
+    /** Executes the operation and returns the expected result depending on the strategy. It might perform scheduling tasks if they are not
+      * independent from the execution.
+      */
+    def execute[T, Result[*]](algorithm: RateLimiterAlgorithm, operation: => T)(using Returns[Result]): Result[T]
+
+    /** Runs the operation and returns the result using the given strategy.
+      */
+    def run[T, Result[_]](operation: => T)(using cfg: Returns[Result]): Result[T] =
+      cfg.run(operation).asInstanceOf[Result[T]]
+
+  end Executor
+
+  object Executor:
+    /** Block rejected operations until the rate limiter is ready to accept them.
+      */
+    case class Block() extends Executor[Strategy.Blocking]:
+
+      val updateLock = new Semaphore(0)
+
+      val schedule = new Semaphore(1)
+
+      def execute[T, Result[*]](algorithm: RateLimiterAlgorithm, operation: => T)(using cfg: Strategy.Blocking[Result]): Result[T] =
+        cfg.run(operation)
+
+      def schedule[T, Result[*]](algorithm: RateLimiterAlgorithm, operation: => T)(using Strategy.Blocking[Result[*]]): Unit =
+        if !algorithm.tryAcquire then
+          // starts scheduler if not already running
+          if schedule.tryAcquire() then
+            supervised:
+              val _ = forkUser:
+                runScheduler(algorithm)
+            ()
+          algorithm.acquire
+
+      private def runScheduler(algorithm: RateLimiterAlgorithm): Unit =
+        val waitTime = algorithm.getNextTime()
+        algorithm.update
+        if waitTime > 0 then
+          val millis = waitTime / 1000000
+          val nanos = waitTime % 1000000
+          Thread.sleep(millis, nanos.toInt)
+          runScheduler(algorithm)
+        else schedule.release()
+      end runScheduler
+
+    end Block
+
+    /** Drops rejected operations
+      */
+    case class Drop() extends Executor[Strategy.Dropping]:
+
+      def schedule[T, Result[*]](algorithm: RateLimiterAlgorithm, operation: => T)(using Strategy.Dropping[Result[*]]): Unit =
+        ()
+
+      def execute[T, Result[*]](algorithm: RateLimiterAlgorithm, operation: => T)(using cfg: Strategy.Dropping[Result[*]]): Result[T] =
+        algorithm.update
+        if algorithm.tryAcquire then cfg.run(operation)
+        else None.asInstanceOf[Result[T]]
+
+    end Drop
+
+    /** Blocks rejected operations until the rate limiter is ready to accept them or drops them depending on the choosen strategy.
+      */
+    case class BlockOrDrop() extends Executor[Strategy.BlockOrDrop]:
+
+      val blockExecutor = Block()
+      val dropExecutor = Drop()
+
+      def execute[T, Result[*]](algorithm: RateLimiterAlgorithm, operation: => T)(using cfg: Strategy.BlockOrDrop[Result]): Result[T] =
+        cfg match
+          case cfg: Strategy.Block =>
+            blockExecutor.execute(algorithm, operation)(using cfg.asInstanceOf[Strategy.Blocking[Result]])
+          case cfg: Strategy.Drop =>
+            dropExecutor.execute(algorithm, operation)(using cfg)
+
+      def schedule[T, Result[*]](algorithm: RateLimiterAlgorithm, operation: => T)(using Strategy.BlockOrDrop[Result]): Unit =
+        implicitly[Strategy.BlockOrDrop[Result]] match
+          case cfg: Strategy.Block =>
+            blockExecutor.schedule(algorithm, operation)(using cfg.asInstanceOf[Strategy.Blocking[Result]])
+          case cfg: Strategy.Drop =>
+            dropExecutor.schedule(algorithm, operation)(using cfg)
+    end BlockOrDrop
+
+  end Executor
+end GenericRateLimiter

core/src/main/scala/ox/resilience/RateLimiter.scala

@@ -0,0 +1,57 @@
+package ox.resilience
+
+import scala.concurrent.duration.*
+import ox.*
+
+/** Rate Limiter with customizable algorithm. It allows to choose between blocking or dropping an operation.
+  */
+case class RateLimiter(
+    algorithm: RateLimiterAlgorithm
+):
+  import GenericRateLimiter.*
+
+  private val rateLimiter =
+    supervised:
+      GenericRateLimiter(Executor.BlockOrDrop(), algorithm)
+
+  /** Blocks the operation until the rate limiter allows it.
+    */
+  def runBlocking[T](operation: => T): T = rateLimiter(operation)(using Strategy.Block())
+
+  /** Drops the operation if not allowed by the rate limiter returning `None`.
+    */
+  def runOrDrop[T](operation: => T): Option[T] = rateLimiter(operation)(using Strategy.Drop())
+
+end RateLimiter
+
+object RateLimiter:
+
+  def leakyBucket(
+      capacity: Int,
+      leakInterval: FiniteDuration
+  ): RateLimiter =
+    RateLimiter(RateLimiterAlgorithm.LeakyBucket(capacity, leakInterval))
+  end leakyBucket
+
+  def tokenBucket(
+      maxTokens: Int,
+      refillInterval: FiniteDuration
+  ): RateLimiter =
+    RateLimiter(RateLimiterAlgorithm.TokenBucket(maxTokens, refillInterval))
+  end tokenBucket
+
+  def fixedRate(
+      maxRequests: Int,
+      windowSize: FiniteDuration
+  ): RateLimiter =
+    RateLimiter(RateLimiterAlgorithm.FixedRate(maxRequests, windowSize))
+  end fixedRate
+
+  def slidingWindow(
+      maxRequests: Int,
+      windowSize: FiniteDuration
+  ): RateLimiter =
+    RateLimiter(RateLimiterAlgorithm.SlidingWindow(maxRequests, windowSize))
+  end slidingWindow
+
+end RateLimiter

core/src/main/scala/ox/resilience/RateLimiterAlgorithm.scala

@@ -0,0 +1,189 @@
+package ox.resilience
+
+import ox.*
+import ox.resilience.RateLimiterAlgorithm.*
+import scala.concurrent.duration.*
+import java.util.concurrent.atomic.AtomicLong
+import scala.concurrent.*
+import java.util.concurrent.atomic.AtomicReference
+import java.util.concurrent.Semaphore
+import java.util.{LinkedList, Queue}
+
+/** Determines the algorithm to use for the rate limiter
+  */
+trait RateLimiterAlgorithm:
+
+  /** Acquire a permit to execute the operation. This method should block until a permit is available.
+    */
+  def acquire: Unit
+
+  /** Try to acquire a permit to execute the operation. This method should not block.
+    */
+  def tryAcquire: Boolean
+
+  /** Updates the internal state of the rate limiter to check whether new operations can be accepted.
+    */
+  def update: Unit
+
+  /** Returns the time until the next operation can be accepted to be used by the `GenericRateLimiter.Executor`. It should return 0 only if
+    * there is no need of rescheduling an update in the future. It should not modify internal state.
+    */
+  def getNextTime(): Long
+
+end RateLimiterAlgorithm
+
+object RateLimiterAlgorithm:
+  /** Fixed rate algorithm
+    */
+  case class FixedRate(rate: Int, per: FiniteDuration) extends RateLimiterAlgorithm:
+    private val lastUpdate = new AtomicLong(System.nanoTime())
+    private val semaphore = new Semaphore(rate)
+    val lock = new java.util.concurrent.locks.ReentrantLock()
+
+    def acquire: Unit =
+      semaphore.acquire()
+
+    def tryAcquire: Boolean =
+      semaphore.tryAcquire()
+
+    def getNextTime(): Long =
+      val waitTime = lastUpdate.get() + per.toNanos - System.nanoTime()
+      val q = semaphore.getQueueLength()
+      if waitTime > 0 then waitTime
+      else if q > 0 then per.toNanos
+      else 0L
+
+    def update: Unit =
+      val now = System.nanoTime()
+      lastUpdate.updateAndGet { time =>
+        if time + per.toNanos < now then
+          semaphore.drainPermits()
+          semaphore.release(rate)
+          now
+        else time
+      }
+      ()
+    end update
+
+  end FixedRate
+
+  /** Sliding window algorithm
+    */
+  case class SlidingWindow(rate: Int, per: FiniteDuration) extends RateLimiterAlgorithm:
+    private val log = new AtomicReference[Queue[Long]](new LinkedList[Long]())
+    private val semaphore = new Semaphore(rate)
+
+    def acquire: Unit =
+      semaphore.acquire()
+      val now = System.nanoTime()
+      log.updateAndGet { q =>
+        q.add(now)
+        q
+      }
+      ()
+    end acquire
+
+    def tryAcquire: Boolean =
+      if semaphore.tryAcquire() then
+        val now = System.nanoTime()
+        log.updateAndGet { q =>
+          q.add(now)
+          q
+        }
+        true
+      else false
+
+    def getNextTime(): Long =
+      val furtherLog = log.get().peek()
+      if null eq furtherLog then
+        if semaphore.getQueueLength() > 0 then per.toNanos
+        else 0L
+      else
+        val waitTime = log.get().peek() + per.toNanos - System.nanoTime()
+        val q = semaphore.getQueueLength()
+        if waitTime > 0 then waitTime
+        else if q > 0 then
+          update
+          getNextTime()
+        else 0L
+      end if
+    end getNextTime
+
+    def update: Unit =
+      val now = System.nanoTime()
+      while semaphore.availablePermits() < rate && log
+          .updateAndGet { q =>
+            if q.peek() < now - per.toNanos then
+              q.poll()
+              semaphore.release()
+              q
+            else q
+          }
+          .peek() < now - per.toNanos
+      do ()
+      end while
+    end update
+
+  end SlidingWindow
+
+  /** Token bucket algorithm
+    */
+  case class TokenBucket(rate: Int, per: FiniteDuration) extends RateLimiterAlgorithm:
+    private val refillInterval = per.toNanos
+    private val lastRefillTime = new AtomicLong(System.nanoTime())
+    private val semaphore = new Semaphore(1)
+
+    def acquire: Unit =
+      semaphore.acquire()
+
+    def tryAcquire: Boolean =
+      semaphore.tryAcquire()
+
+    def getNextTime(): Long =
+      val waitTime = lastRefillTime.get() + refillInterval - System.nanoTime()
+      val q = semaphore.getQueueLength()
+      if waitTime > 0 then waitTime
+      else if q > 0 then refillInterval
+      else 0L
+
+    def update: Unit =
+      val now = System.nanoTime()
+      val elapsed = now - lastRefillTime.get()
+      val newTokens = elapsed / refillInterval
+      lastRefillTime.set(newTokens * refillInterval + lastRefillTime.get())
+      semaphore.release(newTokens.toInt)
+
+  end TokenBucket
+
+  /** Leaky bucket algorithm
+    */
+  case class LeakyBucket(capacity: Int, leakRate: FiniteDuration) extends RateLimiterAlgorithm:
+    private val leakInterval = leakRate.toNanos
+    private val lastLeakTime = new AtomicLong(System.nanoTime())
+    private val semaphore = new Semaphore(capacity)
+
+    def acquire: Unit =
+      semaphore.acquire()
+
+    def tryAcquire: Boolean =
+      semaphore.tryAcquire()
+
+    def getNextTime(): Long =
+      val waitTime = lastLeakTime.get() + leakInterval - System.nanoTime()
+      val q = semaphore.getQueueLength()
+      if waitTime > 0 then waitTime
+      else if q > 0 then leakInterval
+      else 0L
+
+    def update: Unit =
+      val now = System.nanoTime()
+      val lastLeak = lastLeakTime.get()
+      val elapsed = now - lastLeak
+      val leaking = elapsed / leakInterval
+      val newTime = leaking * leakInterval + lastLeak
+      semaphore.release(leaking.toInt)
+      lastLeakTime.set(newTime)
+    end update
+
+  end LeakyBucket
+end RateLimiterAlgorithm