Play Scala 开发技巧 - 请求限速

在系统开发中，我们经常需要保护一些安全性较高的接口，限制这些接口每秒处理的请求数量。例如对于一个计算密集型接口，假设压测值是100rps, 如果实际情况长期高于这个值，则会引起滚雪球效应，最终导致系统崩溃。下面我们一起来看看如何在 Play 中实现一个完全异步非阻塞的请求限速 ？本文代码已提交至 play-community 项目，详情请参考 controllers.demo.ThrottleDemoController 。

1 实现思路

当 Controller 接收到请求时，为该请求建立一个“开关”，并且把该“开关”发送给“限速器”，"限速器"通过“开关”控制请求的处理速度。本文采用 Promise 实现“开关”，使用 Akka Stream 的 SouceQueue 实现“限速器”。

2 实现代码

2.1 开关

创建 ThrottledRequest 类，用于存放“开关”和请求到达时间，

case class ThrottledRequest(promise: Promise[Boolean], time: Long)

2.2 限速器

使用 SourceQueue 创建“限速器”，处理请求限速和请求超时，

val sourceQueue =
  Source
    .queue[ThrottledRequest](100, OverflowStrategy.backpressure)
    .throttle(1, 1.second, 1, ThrottleMode.shaping)
    .toMat(Sink.foreach{ r =>
      // 处理未超时请求
      if (System.currentTimeMillis() - r.time <= 15000) {
        r.promise.success(false)
      } else {
        r.promise.success(true)
      }
    })(Keep.left).run()

Source.queue 方法的声明如下,

def queue[T](bufferSize: Int, overflowStrategy: OverflowStrategy): Source[T, SourceQueueWithComplete[T]]

这里我们设置缓冲区大小为 100， 设置缓冲区溢出时的处理策略为 backpressure ，以防止请求丢失。通过 throttle 方法设置请求处理速度为 1 个每秒。 Sink 负责处理请求的放行和超时。

2.3 请求拦截

请求拦截 Action 负责拦截所有发往目标 Action 的请求，为每个请求创建“开关”并发送给“限速器”，然后只放行被“限速器”打开开关的请求，

// 只有通过限速器(sourceQueue)的请求才会被执行
def throttle[A](action: Action[A]) = Action.async(action.parser) { request =>
  val promise = Promise[Boolean]()
  sourceQueue.offer(ThrottledRequest(promise, System.currentTimeMillis()))
  promise.future.flatMap{ isTimeout =>
    if (!isTimeout) {
      action(request)
    } else {
      Future.successful(Forbidden("Timeout."))
    }
  }
}

2.4 目标 Action

目标 Action 组合了拦截 Action，代码如下，

def throttledAction = throttle { Action { implicit request: Request[AnyContent] =>
  Ok("Finish.")
}}

2.5 定义路由

GET  /demo/throttle  controllers.demo.ThrottleDemoController.throttledAction

3 测试实现

上面我们实现了一个限速接口，每秒只处理 1 个请求，当请求排队超过 15 秒时， 直接返回 403 响应。 我们通过下面的测试代码， 将 100 个请求瞬间发送过去，然后异步异步打印响应信息，

val client = buildAHCClient
(1 to 100).foreach{ _ =>
  client.prepareGet(s"http://localhost:9000/demo/throttle")
    .execute(new AsyncCompletionHandler[Response] {
      override def onCompleted(response: Response): Response = {
        println(s"status: ${response.getStatusCode}, time: ${DateTime.now().toString("HH:mm:ss")}")
        response
      }
    })
}
Thread.sleep(100000)
client.close()

测试结果如下：

status: 200, time: 23:22:06
status: 200, time: 23:22:07
status: 200, time: 23:22:08
status: 200, time: 23:22:09
status: 200, time: 23:22:10
status: 200, time: 23:22:11
status: 200, time: 23:22:12
status: 200, time: 23:22:13
status: 200, time: 23:22:14
status: 200, time: 23:22:15
status: 200, time: 23:22:16
status: 200, time: 23:22:17
status: 200, time: 23:22:18
status: 200, time: 23:22:19
status: 200, time: 23:22:20
status: 200, time: 23:22:21
status: 403, time: 23:22:22
status: 403, time: 23:22:23
...

从上面可以看出，请求按照到达顺序依次被处理，从响应时间上看，目标接口确实每秒只处理 1 个请求， 并且从 23时22分22秒 开始，后面的请求均被超时处理。

4 小结

异步非阻塞代码虽然写起来有点麻烦，并且不易于调试，但是在系统性能方面收益是巨大的。在相同的系统性能指标下，异步非阻塞代码可以让硬件成本降到最低。理论上，使用异步非阻塞方式编写的系统可以在单个线程上运行，并且可以保证较高的并发性，典型例子是Node.js。Play Framework 是一个完全异步非阻塞的 Web 开发框架，相信在不久的将来在国内会越来越受欢迎。