自定义Akka Stream大小限制未抛出EntityStreamSizeException问题排查

我有一个Akka HTTP API，用户向服务器发送S3 URL，服务器从AWS拉取流并执行后续操作。由于无法使用Akka HTTP的withSizeLimit指令，自定义实现了流大小限制逻辑，但测试时发现文件大小超过4MB时，服务器并未抛出预期的EntityStreamSizeException。

自定义实现代码如下：

final case class SizeLimit(maxBytes: Long, contentLength: Option[Long] = None) extends Attributes.Attribute {
  def isDisabled = maxBytes < 0
}

object Limitable {
  def applyForByteStrings[Mat](source: Source[ByteString, Mat], limit: SizeLimit): Source[ByteString, Mat] =
    applyLimit(source, limit)(_.size)

  def applyForChunks[Mat](source: Source[ChunkStreamPart, Mat], limit: SizeLimit): Source[ChunkStreamPart, Mat] =
    applyLimit(source, limit)(_.data.size)

  def applyLimit[T, Mat](source: Source[T, Mat], limit: SizeLimit)(sizeOf: T => Int): Source[T, Mat] =
    if (limit.isDisabled) source withAttributes Attributes(limit) // no need to add stage, it's either there or not needed
    else source.via(new Limitable(sizeOf)) withAttributes Attributes(limit)

  private val limitableDefaults = Attributes.name("limitable")
}

final class Limitable[T](sizeOf: T => Int) extends GraphStage[FlowShape[T, T]] {
  val in = Inlet[T]("Limitable.in")
  val out = Outlet[T]("Limitable.out")
  var numPullCalls = 0
  var numPushCalls = 0
  override val shape = FlowShape.of(in, out)
  override protected val initialAttributes: Attributes = Limitable.limitableDefaults

  override def createLogic(_attributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) with InHandler with OutHandler {
    private var maxBytes = -1L
    private var bytesLeft = Long.MaxValue

    @nowarn("msg=deprecated") // we need getFirst semantics
    override def preStart(): Unit = {
      _attributes.getFirst[SizeLimit] match {
        case Some(limit: SizeLimit) if limit.isDisabled =>
        // "no limit"
        case Some(SizeLimit(bytes, cl @ Some(contentLength))) =>
          if (contentLength > bytes) failStage(EntityStreamSizeException(bytes, cl))
        // else we still count but never throw an error
        case Some(SizeLimit(bytes, None)) =>
          maxBytes = bytes
          bytesLeft = bytes
        case None =>
      }
    }

    override def onPush(): Unit = {
      numPushCalls += 1
      println(s"Push calls $numPushCalls")
      val elem = grab(in)
      val temp = sizeOf(elem)
      println(s"Elem size is $temp")
      bytesLeft -= temp
      if (bytesLeft >= 0) {
        push(out, elem)
      }
      else {
        println(s"EntityStreamSizeException Bytes left $bytesLeft")
        failStage(EntityStreamSizeException(maxBytes))
      }
    }

    override def onPull(): Unit = {
      numPullCalls += 1
      println(s"Pull calls $numPullCalls")
      pull(in)
    }

    setHandlers(in, out, this)
  }
}


val filePath = Paths.get("/Users/<username>/Documents/bigfile.pdf")
val fileSource: Source[ByteString, Any] = FileIO.fromPath(filePath)
val res = Limitable.applyForByteStrings(fileSource, SizeLimit(4000000L))
val sink   = StreamConverters.asInputStream()
val result = res.runWith(sink)
val tis = TikaInputStream.get(result)

问题根源分析 #

1. bytesLeft初始化不完整
   当SizeLimit传入contentLength且该值未超过限制时，代码未初始化bytesLeft，此时bytesLeft保持默认值Long.MaxValue，后续的字节计数永远不会触发负数判断，自然不会抛出异常。

2. 属性传递方式错误
   applyLimit方法中使用withAttributes替换了Limitable阶段的所有属性，覆盖了阶段自带的limitableDefaults，可能导致SizeLimit属性无法被正确获取。

3. 测试代码未完全消费流
   TikaInputStream可能仅读取部分流内容就完成解析，导致Akka Stream的限制阶段未处理到超过大小限制的部分元素，无法触发异常。

修复方案 #

1. 统一初始化bytesLeft #

修改preStart方法，确保所有分支都正确初始化bytesLeft，同时保留contentLength的预检查：

override def preStart(): Unit = {
  _attributes.getFirst[SizeLimit] match {
    case Some(limit: SizeLimit) if limit.isDisabled =>
      // 无限制，保持bytesLeft为Long.MaxValue
    case Some(SizeLimit(max, cl)) =>
      maxBytes = max
      // 预检查contentLength，如果存在且超过限制直接失败
      cl.foreach { contentLength =>
        if (contentLength > max) failStage(EntityStreamSizeException(max, cl))
      }
      // 初始化剩余字节数为限制值，开始实际计数
      bytesLeft = max
    case None =>
      // 未设置SizeLimit，不做限制
  }
}

2. 修正属性传递方式 #

将withAttributes改为addAttributes，避免替换阶段原有属性：

def applyLimit[T, Mat](source: Source[T, Mat], limit: SizeLimit)(sizeOf: T => Int): Source[T, Mat] =
  if (limit.isDisabled) source.addAttributes(Attributes(limit))
  else source.via(new Limitable(sizeOf)).addAttributes(Attributes(limit))

3. 确保测试时完全消费流 #

修改测试代码，读取整个流以触发所有阶段的执行：

val filePath = Paths.get("/Users/<username>/Documents/bigfile.pdf")
val fileSource: Source[ByteString, Any] = FileIO.fromPath(filePath)
val res = Limitable.applyForByteStrings(fileSource, SizeLimit(4000000L))
val sink = StreamConverters.asInputStream()
val result = res.runWith(sink)

// 强制读取整个流，触发限制逻辑
val buffer = new ByteArrayOutputStream()
val readBuffer = new Array[Byte](8192)
var readCount: Int = 0
while ({ readCount = result.read(readBuffer); readCount != -1 }) {
  buffer.write(readBuffer, 0, readCount)
}
result.close()

4. 优化单个元素超限判断（可选） #

如果单个元素大小直接超过限制，可以提前触发异常，避免无效的计数：

override def onPush(): Unit = {
  numPushCalls += 1
  val elem = grab(in)
  val elemSize = sizeOf(elem)
  
  // 单个元素直接超过总限制，直接失败
  if (elemSize > maxBytes) {
    failStage(EntityStreamSizeException(maxBytes))
    return
  }
  
  bytesLeft -= elemSize
  if (bytesLeft >= 0) {
    push(out, elem)
  } else {
    failStage(EntityStreamSizeException(maxBytes))
  }
}

内容的提问来源于stack exchange，提问作者Sushant Somani