Skip to Content
Course content

188: Rate Limiting Middleware in Akka HTTP

Click on the "Edit" button in the top corner of the screen to edit your slide content.

I was looking through the logs for our weather service API this morning, and something jumped out at me. One specific IP address was hitting the /forecast endpoint about 50 times a second. It's not a DDoS attack, but it's enough to make the response times spike for everyone else. We can't just block them entirely—they might be a legitimate partner—but we definitely need to throttle them.

The Naive Guardrail

My first instinct was to just throw a counter inside the route handler. I figured I'd keep a map of IP addresses and their request counts, and if the count is too high, I'll just return a 429. It looked something like this in my head:

val requestCounts = scala.collection.concurrent.TrieMap.empty[String, Int]

val route = path("forecast") {
  extractClientIP { ip =>
    val count = requestCounts.getOrElse(ip, 0) + 1
    requestCounts.put(ip, count)
    
    if (count > 100) {
      complete(StatusCodes.TooManyRequests -> "Slow down there!")
    } else {
      complete(s"The weather is sunny for $ip")
    }
  }
}

I ran this for a few minutes, and while it technically worked, I hated it. I've effectively polluted my business logic—the actual forecasting code—with infrastructure concerns. If I decide to add rate limiting to the /alerts or /history endpoints, I'm going to be copy-pasting this boilerplate everywhere. That's a recipe for a maintenance nightmare.

Extracting the Logic into a Directive

In Akka HTTP, whenever I find myself repeating "cross-cutting" logic like this, I start thinking about Directives. I want something I can just wrap around any route I care about. I'll try to build a custom Directive that handles the tracking and the rejection automatically.

Here is where I started experimenting. I wanted a rateLimit directive that I could just drop in. I realized I needed a way to track the "window" of time, otherwise, a user would just be blocked forever once they hit 100 requests total.

import akka.http.scaladsl.server._
import akka.http.scaladsl.model._
import scala.collection.concurrent.TrieMap
import scala.concurrent.duration._

class RateLimiter(limit: Int, window: scala.concurrent.duration.FiniteDuration) {
  // Map of IP to (Count, LastResetTimestamp)
  private val clients = TrieMap.empty[String, (Int, Long)]

  def apply(inner: Route) = {
    extractClientIP { ip =>
      val now = System.currentTimeMillis()
      val (count, lastReset) = clients.getOrElse(ip, (0, now))

      if (now - lastReset > window.toMillis) {
        // Window expired, reset count
        clients.put(ip, (1, now))
        inner // Allow the request
      } else if (count < limit) {
        // Still within limit
        clients.put(ip, (count + 1, lastReset))
        inner // Allow the request
      } else {
        // Limit exceeded
        complete(StatusCodes.TooManyRequests -> "Rate limit exceeded. Try again later.")
      }
    }
  }
}

Refining the Implementation

I tried running the above, but I hit a snag. The apply method I wrote isn't a true Akka HTTP Directive; it's just a function that returns a Route. While that works for simple cases, it doesn't compose well with other directives. To do this "the right way," I should extend Directive.

I also realized that my TrieMap is going to grow indefinitely. Every single unique IP that hits the server stays in memory forever. In a real production environment, I'd probably use Redis with a TTL (Time To Live), but for our Scala-centric exploration, I'll stick to a local map while acknowledging the memory leak risk. To make it a proper directive, I'll wrap the logic in a Directive class:

def rateLimit(limit: Int, window: scala.concurrent.duration.FiniteDuration): Directive0 = {
  val clients = TrieMap.empty[String, (Int, Long)]
  
  extractClientIP { ip =>
    val now = System.currentTimeMillis()
    val (count, lastReset) = clients.getOrElse(ip, (0, now))

    if (now - lastReset > window.toMillis) {
      clients.put(ip, (1, now))
      pass // This lets the request continue to the inner route
    } else if (count < limit) {
      clients.put(ip, (count + 1, lastReset))
      pass
    } else {
      complete(StatusCodes.TooManyRequests -> "Too many requests!")
    }
  }
}

Now, look how much cleaner the actual API definition becomes. I can just wrap the sensitive endpoints:

val route = 
  path("forecast") {
    rateLimit(10, 1.minute) {
      get {
        complete("Here is your 7-day forecast!")
      }
    }
  } ~ 
  path("public-info") {
    get {
      complete("This is free for everyone!")
    }
  }

By moving this into a directive, the /forecast route doesn't even know the rate limiter exists. It just receives requests that have already been "cleansed" by the middleware. I've separated the "how many requests are allowed" concern from the "what is the weather" concern.




📋 Practical Task

Exercise: Implementing a Tiered Rate Limiter

Our weather service now has "Free" and "Premium" users. You need to modify the rate-limiting logic to support different limits based on an API key provided in the request header.

  • Requirement 1: Create a custom directive called tieredRateLimit.
  • Requirement 2: The directive should extract a header named X-API-Key.
  • Requirement 3: If the key is "premium-secret", allow 100 requests per minute.
  • Requirement 4: If the key is anything else (or missing), allow only 5 requests per minute.
  • Requirement 5: If the limit is exceeded, return a 429 Too Many Requests status with a message indicating which tier they are on.

Implement this logic as a reusable directive and apply it to a route /premium-forecast.

Rating
0 0

There are no comments for now.

to be the first to leave a comment.