diff --git a/project/Dependencies.scala b/project/Dependencies.scala index 18b2953b88..830380984e 100644 --- a/project/Dependencies.scala +++ b/project/Dependencies.scala @@ -2,7 +2,7 @@ import sbt.* object Dependencies { val JwtCoreVersion = "11.0.4" - val NettyVersion = "4.2.12.Final" + val NettyVersion = "4.2.13.Final" val ScalaCompatCollectionVersion = "2.14.0" val ZioVersion = "2.1.25" val ZioCliVersion = "0.8.0" diff --git a/zio-http/jvm/src/main/scala/zio/http/netty/AsyncBodyReader.scala b/zio-http/jvm/src/main/scala/zio/http/netty/AsyncBodyReader.scala index 9d06f20749..d9b20b041d 100644 --- a/zio-http/jvm/src/main/scala/zio/http/netty/AsyncBodyReader.scala +++ b/zio-http/jvm/src/main/scala/zio/http/netty/AsyncBodyReader.scala @@ -32,12 +32,20 @@ import io.netty.channel.{ChannelHandlerContext, SimpleChannelInboundHandler} import io.netty.handler.codec.http.{HttpContent, LastHttpContent} import io.netty.util.concurrent.ScheduledFuture -private[netty] abstract class AsyncBodyReader(timeoutMillis: Option[Long]) - extends SimpleChannelInboundHandler[HttpContent](true) { self => +private[netty] abstract class AsyncBodyReader( + timeoutMillis: Option[Long], + maxPreConnectBufferSize: Int = AsyncBodyReader.DefaultMaxPreConnectBufferSize, +) extends SimpleChannelInboundHandler[HttpContent](true) { self => import zio.http.netty.AsyncBodyReader._ + require(maxPreConnectBufferSize >= 0, "maxPreConnectBufferSize must be >= 0") + private var state: State = State.Buffering private val buffer = new mutable.ArrayBuilder.ofByte() + // Tracks `buffer` size while in `State.Buffering` so we can apply back-pressure + // before the consumer connects. `ArrayBuilder#knownSize` is unreliable on + // Scala 2.12, so we count bytes ourselves. + private var bufferedBytes: Int = 0 private var previousAutoRead: Boolean = false private var readingDone: Boolean = false private var ctx: ChannelHandlerContext = _ @@ -114,9 +122,12 @@ private[netty] abstract class AsyncBodyReader(timeoutMillis: Option[Long]) val readMore = state match { case State.Buffering => - // `connect` method hasn't been called yet, add all incoming content to the buffer + // `connect` method hasn't been called yet, add all incoming content to the buffer. + // Cap the pre-connect buffer to avoid unbounded heap growth when a fast producer + // outpaces a slow-to-start consumer (see issue #3173). buffer0.addAll(content) - true + bufferedBytes += content.length + bufferedBytes < maxPreConnectBufferSize case State.Direct(callback) if isLast && buffer0.knownSize == 0 => // Buffer is empty, we can just use the array directly callback(Chunk.fromArray(content), isLast = true) @@ -219,6 +230,13 @@ private[netty] abstract class AsyncBodyReader(timeoutMillis: Option[Long]) private[netty] object AsyncBodyReader { + /** + * Default cap on bytes that may be buffered before the consumer calls + * `connect`. Once the buffer reaches this size we stop calling `ctx.read()`, + * relying on Netty's TCP back-pressure until the consumer drains the buffer. + */ + final val DefaultMaxPreConnectBufferSize: Int = 64 * 1024 + sealed trait State object State { diff --git a/zio-http/jvm/src/main/scala/zio/http/netty/server/ServerAsyncBodyHandler.scala b/zio-http/jvm/src/main/scala/zio/http/netty/server/ServerAsyncBodyHandler.scala index ade061bbea..700b539845 100644 --- a/zio-http/jvm/src/main/scala/zio/http/netty/server/ServerAsyncBodyHandler.scala +++ b/zio-http/jvm/src/main/scala/zio/http/netty/server/ServerAsyncBodyHandler.scala @@ -18,4 +18,5 @@ package zio.http.netty.server import zio.http.netty.AsyncBodyReader -private[zio] final class ServerAsyncBodyHandler extends AsyncBodyReader(timeoutMillis = None) {} +private[zio] final class ServerAsyncBodyHandler(maxPreConnectBufferSize: Int) + extends AsyncBodyReader(timeoutMillis = None, maxPreConnectBufferSize = maxPreConnectBufferSize) {} diff --git a/zio-http/jvm/src/main/scala/zio/http/netty/server/ServerInboundHandler.scala b/zio-http/jvm/src/main/scala/zio/http/netty/server/ServerInboundHandler.scala index 0637b418bc..84f3b0ed6c 100644 --- a/zio-http/jvm/src/main/scala/zio/http/netty/server/ServerInboundHandler.scala +++ b/zio-http/jvm/src/main/scala/zio/http/netty/server/ServerInboundHandler.scala @@ -131,7 +131,7 @@ private[zio] final case class ServerInboundHandler( } private def addAsyncBodyHandler(ctx: ChannelHandlerContext): AsyncBodyReader = { - val handler = new ServerAsyncBodyHandler + val handler = new ServerAsyncBodyHandler(config.requestBodyPreConnectBufferSize) ctx .channel() .pipeline() diff --git a/zio-http/jvm/src/test/scala/zio/http/netty/AsyncBodyReaderSpec.scala b/zio-http/jvm/src/test/scala/zio/http/netty/AsyncBodyReaderSpec.scala new file mode 100644 index 0000000000..c642ebee74 --- /dev/null +++ b/zio-http/jvm/src/test/scala/zio/http/netty/AsyncBodyReaderSpec.scala @@ -0,0 +1,109 @@ +/* + * Copyright 2021 - 2023 Sporta Technologies PVT LTD & the ZIO HTTP contributors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package zio.http.netty + +import java.util.concurrent.atomic.AtomicInteger + +import zio._ +import zio.test._ + +import zio.http.ZIOHttpSpec + +import io.netty.buffer.Unpooled +import io.netty.channel.embedded.EmbeddedChannel +import io.netty.channel.{ChannelHandlerContext, ChannelOutboundHandlerAdapter} +import io.netty.handler.codec.http.DefaultHttpContent + +/** + * Regression test for the pre-connect Buffering window in `AsyncBodyReader`. + * + * Before any consumer calls `connect`, the reader buffers incoming body chunks + * on the heap and asks Netty for more (`ctx.read()`). Without a cap, a fast + * producer (e.g. a localhost upload) can fill the heap during the few + * milliseconds it takes the route handler to call `request.body.asStream` and + * start pulling. See issue #3173. + * + * `EmbeddedChannel.writeInbound` ignores `autoRead`, so we cannot verify the + * fix by checking that buffered bytes stay below the cap. Instead we count the + * outbound `read()` requests AsyncBodyReader makes — the cap should make those + * calls stop, which is what enforces TCP back-pressure on a real socket. + */ +object AsyncBodyReaderSpec extends ZIOHttpSpec { + + private final class TestReader(maxPreConnectBufferSize: Int) + extends AsyncBodyReader(timeoutMillis = None, maxPreConnectBufferSize = maxPreConnectBufferSize) + + /** Outbound handler placed before AsyncBodyReader to count `ctx.read()`. */ + private final class ReadCounter extends ChannelOutboundHandlerAdapter { + val count = new AtomicInteger(0) + override def read(ctx: ChannelHandlerContext): Unit = { + count.incrementAndGet() + ctx.read(): Unit + } + } + + private def setup(cap: Int): (ReadCounter, EmbeddedChannel) = { + val counter = new ReadCounter + val ch = new EmbeddedChannel(new TestReader(maxPreConnectBufferSize = cap)) + // addFirst places counter at the head of the pipeline, so outbound `read` + // events from AsyncBodyReader pass through it on their way to the channel. + ch.pipeline().addFirst(counter): Unit + (counter, ch) + } + + private def feed(ch: EmbeddedChannel, chunkBytes: Int, count: Int): Unit = { + val payload = new Array[Byte](chunkBytes) + var i = 0 + while (i < count) { + ch.writeInbound(new DefaultHttpContent(Unpooled.wrappedBuffer(payload))): Unit + i += 1 + } + } + + override def spec: Spec[TestEnvironment with Scope, Any] = + suite("AsyncBodyReader")( + test("Buffering state stops requesting reads once pre-connect cap is reached") { + ZIO.attempt { + val cap = 64 * 1024 + val chunkSz = 8 * 1024 + val (counter, ch) = setup(cap) + // Feed 100 chunks of 8 KB = 800 KB, far above the 64 KB cap. + feed(ch, chunkBytes = chunkSz, count = 100) + + // Without the fix, AsyncBodyReader would request a read after every + // chunk (~100 reads). With the cap, a read is requested only while + // bufferedBytes < cap; the chunk that pushes bufferedBytes to `cap` + // is buffered but does not trigger another read. Expected reads: + // (cap - 1) / chunkSz = 7. + val expectedReads = (cap - 1) / chunkSz + val reads = counter.count.get() + assertTrue(reads == expectedReads) + } + }, + test("Buffering state keeps requesting reads while under cap") { + ZIO.attempt { + val cap = 64 * 1024 + val chunkSz = 1024 + val (counter, ch) = setup(cap) + feed(ch, chunkBytes = chunkSz, count = 4) + // 4 KB buffered, well below 64 KB cap → reader still asking for more. + val reads = counter.count.get() + assertTrue(reads == 4) + } + }, + ) +} diff --git a/zio-http/shared/src/main/scala/zio/http/Server.scala b/zio-http/shared/src/main/scala/zio/http/Server.scala index 6d22498271..c73664697d 100644 --- a/zio-http/shared/src/main/scala/zio/http/Server.scala +++ b/zio-http/shared/src/main/scala/zio/http/Server.scala @@ -76,6 +76,8 @@ object Server extends ServerPlatformSpecific { tcpNoDelay: Boolean, @unroll generateHeadRoutes: Boolean = false, + @unroll + requestBodyPreConnectBufferSize: Int = 64 * 1024, ) { self => /** @@ -197,6 +199,24 @@ object Server extends ServerPlatformSpecific { def requestStreaming(requestStreaming: RequestStreaming): Config = self.copy(requestStreaming = requestStreaming) + /** + * Soft cap on the bytes the server will buffer from a streamed request body + * before the route handler starts consuming it. The cap may be overshot by + * up to one `HttpContent` chunk: the chunk that crosses the threshold is + * appended to the buffer before reads stop. Once the threshold is reached, + * Netty's TCP back-pressure kicks in and the client is paused until the + * handler calls `request.body.asStream` (or similar) and begins reading. + * Prevents heap exhaustion when an async handler defers consuming the body + * and the client uploads faster than the handler can connect. Applies when + * [[RequestStreaming.Enabled]] is configured, and also with + * [[RequestStreaming.Hybrid]] for request bodies that exceed the aggregated + * threshold and are therefore handled as streamed bodies. + */ + def requestBodyPreConnectBufferSize(size: Int): Config = { + require(size >= 0, "requestBodyPreConnectBufferSize must be >= 0") + self.copy(requestBodyPreConnectBufferSize = size) + } + /** * Sets the maximum number of connection requests that will be queued before * being rejected @@ -232,7 +252,10 @@ object Server extends ServerPlatformSpecific { zio.Config.duration("idle-timeout").optional.withDefault(Config.default.idleTimeout) ++ zio.Config.boolean("avoid-context-switching").withDefault(Config.default.avoidContextSwitching) ++ zio.Config.int("so-backlog").withDefault(Config.default.soBacklog) ++ - zio.Config.boolean("tcp-nodelay").withDefault(Config.default.tcpNoDelay) + zio.Config.boolean("tcp-nodelay").withDefault(Config.default.tcpNoDelay) ++ + zio.Config + .int("request-body-pre-connect-buffer-size") + .withDefault(Config.default.requestBodyPreConnectBufferSize) }.map { case ( @@ -252,6 +275,7 @@ object Server extends ServerPlatformSpecific { avoidCtxSwitch, soBacklog, tcpNoDelay, + requestBodyPreConnectBufferSize, ) => default.copy( sslConfig = sslConfig, @@ -269,6 +293,7 @@ object Server extends ServerPlatformSpecific { avoidContextSwitching = avoidCtxSwitch, soBacklog = soBacklog, tcpNoDelay = tcpNoDelay, + requestBodyPreConnectBufferSize = requestBodyPreConnectBufferSize, ) } @@ -289,6 +314,7 @@ object Server extends ServerPlatformSpecific { avoidContextSwitching = false, soBacklog = 100, tcpNoDelay = true, + requestBodyPreConnectBufferSize = 64 * 1024, ) final case class ResponseCompressionConfig(