Improve the Client's performance for non-stream bodies

zio-http-benchmarks/src/main/scala/zhttp.benchmarks/ClientBenchmark.scala

@@ -0,0 +1,94 @@
+package zhttp.benchmarks
+
+import org.openjdk.jmh.annotations._
+import zio._
+import zio.http._
+
+import java.util.concurrent.TimeUnit
+import scala.annotation.nowarn
+
+@nowarn
+@State(org.openjdk.jmh.annotations.Scope.Benchmark)
+@BenchmarkMode(Array(Mode.Throughput))
+@OutputTimeUnit(TimeUnit.SECONDS)
+@Warmup(iterations = 3, time = 3)
+@Measurement(iterations = 3, time = 3)
+@Fork(1)
+class ClientBenchmark {
+  private val random = scala.util.Random
+  random.setSeed(42)
+
+  private implicit val unsafe: Unsafe = Unsafe.unsafe(identity)
+
+  @Param(Array("small", "large"))
+  var path: String = _
+
+  private val smallString = "Hello!".getBytes
+  private val largeString = random.alphanumeric.take(10000).mkString.getBytes
+
+  private val smallRequest = Request(url = url"http://0.0.0.0:8080/small")
+  private val largeRequest = Request(url = url"http://0.0.0.0:8080/large")
+
+  private val smallResponse = Response(status = Status.Ok, body = Body.fromArray(smallString))
+  private val largeResponse = Response(status = Status.Ok, body = Body.fromArray(largeString))
+
+  private val smallRoute = Route.route(Method.GET / "small")(handler(smallResponse))
+  private val largeRoute = Route.route(Method.GET / "large")(handler(largeResponse))
+
+  private val shutdownResponse = Response.text("shutting down")
+
+  private def shutdownRoute(shutdownSignal: Promise[Nothing, Unit]) =
+    Route.route(Method.GET / "shutdown")(handler(shutdownSignal.succeed(()).as(shutdownResponse)))
+
+  private def http(shutdownSignal: Promise[Nothing, Unit]) =
+    Routes(smallRoute, largeRoute, shutdownRoute(shutdownSignal))
+
+  private val rtm     = Runtime.unsafe.fromLayer(ZClient.default)
+  private val runtime = rtm.unsafe
+
+  private def run(f: RIO[Client, Any]): Any = runtime.run(f).getOrThrow()
+
+  @Setup(Level.Trial)
+  def setup(): Unit = {
+    val startServer: Task[Unit] = (for {
+      shutdownSignal <- Promise.make[Nothing, Unit]
+      fiber          <- Server.serve(http(shutdownSignal)).fork
+      _              <- shutdownSignal.await *> fiber.interrupt
+    } yield ()).provideLayer(Server.default)
+
+    val waitForServerStarted: Task[Unit] = (for {
+      client <- ZIO.service[Client]
+      _      <- client.request(smallRequest)
+    } yield ()).provide(ZClient.default, zio.Scope.default)
+
+    run(startServer.forkDaemon *> waitForServerStarted.retry(Schedule.fixed(1.second)))
+  }
+
+  @TearDown(Level.Trial)
+  def tearDown(): Unit = {
+    val stopServer = (for {
+      client <- ZIO.service[Client]
+      _      <- client.request(Request(url = url"http://localhost:8080/shutdown"))
+    } yield ()).provide(ZClient.default, zio.Scope.default)
+    run(stopServer)
+    rtm.shutdown0()
+  }
+
+  @Benchmark
+  @OperationsPerInvocation(100)
+  def zhttpChunkBenchmark(): Any = run {
+    val req = if (path == "small") smallRequest else largeRequest
+    ZIO.serviceWithZIO[Client] { client =>
+      ZIO.scoped(client.request(req).flatMap(_.body.asChunk)).repeatN(100)
+    }
+  }
+
+  @Benchmark
+  @OperationsPerInvocation(100)
+  def zhttpStreamToChunkBenchmark(): Any = run {
+    val req = if (path == "small") smallRequest else largeRequest
+    ZIO.serviceWithZIO[Client] { client =>
+      ZIO.scoped(client.request(req).flatMap(_.body.asStream.runCollect)).repeatN(100)
+    }
+  }
+}

zio-http/jvm/src/main/scala/zio/http/netty/AsyncBodyReader.scala

@@ -18,45 +18,50 @@ package zio.http.netty
 
 import java.io.IOException
 
+import scala.collection.mutable
+
+import zio.Chunk
 import zio.stacktracer.TracingImplicits.disableAutoTrace
-import zio.{Chunk, ChunkBuilder}
 
-import zio.http.netty.AsyncBodyReader.State
 import zio.http.netty.NettyBody.UnsafeAsync
 
 import io.netty.buffer.ByteBufUtil
 import io.netty.channel.{ChannelHandlerContext, SimpleChannelInboundHandler}
 import io.netty.handler.codec.http.{HttpContent, LastHttpContent}
 
 abstract class AsyncBodyReader extends SimpleChannelInboundHandler[HttpContent](true) {
+  import zio.http.netty.AsyncBodyReader._
+
+  private var state: State               = State.Buffering
+  private val buffer                     = new mutable.ArrayBuilder.ofByte()
+  private var previousAutoRead: Boolean  = false
+  private var readingDone: Boolean       = false
+  private var ctx: ChannelHandlerContext = _
 
-  private var state: State                                 = State.Buffering
-  private val buffer: ChunkBuilder[(Chunk[Byte], Boolean)] = ChunkBuilder.make[(Chunk[Byte], Boolean)]()
-  private var previousAutoRead: Boolean                    = false
-  private var ctx: ChannelHandlerContext                   = _
+  private def result(buffer: mutable.ArrayBuilder.ofByte): Chunk[Byte] = {
+    val arr = buffer.result()
+    Chunk.ByteArray(arr, 0, arr.length)
+  }
 
   private[zio] def connect(callback: UnsafeAsync): Unit = {
+    val buffer0 = buffer // Avoid reading it from the heap in the synchronized block
     this.synchronized {
       state match {
         case State.Buffering =>
-          val result: Chunk[(Chunk[Byte], Boolean)] = buffer.result()
-          val readingDone: Boolean                  = result.lastOption match {
-            case None              => false
-            case Some((_, isLast)) => isLast
-          }
-          buffer.clear() // GC
-
-          if (ctx.channel.isOpen || readingDone) {
-            state = State.Direct(callback)
-            result.foreach { case (chunk, isLast) =>
-              callback(chunk, isLast)
+          state = State.Direct(callback)
+
+          if (readingDone) {
+            callback(result(buffer0), isLast = true)
+          } else if (ctx.channel().isOpen) {
+            callback match {
+              case UnsafeAsync.Aggregating(bufSize) => buffer.sizeHint(bufSize)
+              case cb                               => cb(result(buffer0), isLast = false)
             }
             ctx.read(): Unit
           } else {
             throw new IllegalStateException("Attempting to read from a closed channel, which will never finish")
           }
-
-        case State.Direct(_) =>
+        case _               =>
           throw new IllegalStateException("Cannot connect twice")
       }
     }
@@ -76,22 +81,36 @@ abstract class AsyncBodyReader extends SimpleChannelInboundHandler[HttpContent](
     ctx: ChannelHandlerContext,
     msg: HttpContent,
   ): Unit = {
-    val isLast = msg.isInstanceOf[LastHttpContent]
-    val chunk  = Chunk.fromArray(ByteBufUtil.getBytes(msg.content()))
+    val buffer0 = buffer // Avoid reading it from the heap in the synchronized block
 
     this.synchronized {
+      val isLast  = msg.isInstanceOf[LastHttpContent]
+      val content = ByteBufUtil.getBytes(msg.content())
+
       state match {
-        case State.Buffering        =>
-          buffer += ((chunk, isLast))
-        case State.Direct(callback) =>
-          callback(chunk, isLast)
-          ctx.read()
+        case State.Buffering                                            =>
+          // `connect` method hasn't been called yet, add all incoming content to the buffer
+          buffer0.addAll(content)
+        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)
+        case State.Direct(callback: UnsafeAsync.Aggregating)            =>
+          // We're aggregating the full response, only call the callback on the last message
+          buffer0.addAll(content)
+          if (isLast) callback(result(buffer0), isLast = true)
+        case State.Direct(callback)                                     =>
+          // We're streaming, emit chunks as they come
+          callback(Chunk.fromArray(content), isLast)
       }
-    }
 
-    if (isLast) {
-      ctx.channel().pipeline().remove(this)
-    }: Unit
+      if (isLast) {
+        readingDone = true
+        ctx.channel().pipeline().remove(this)
+      } else {
+        ctx.read()
+      }
+      ()
+    }
   }
 
   override def exceptionCaught(ctx: ChannelHandlerContext, cause: Throwable): Unit = {
@@ -125,4 +144,10 @@ object AsyncBodyReader {
 
     final case class Direct(callback: UnsafeAsync) extends State
   }
+
+  // For Scala 2.12. In Scala 2.13+, the methods directly implemented on ArrayBuilder[Byte] are selected over syntax.
+  private implicit class ByteArrayBuilderOps[A](private val self: mutable.ArrayBuilder[Byte]) extends AnyVal {
+    def addAll(as: Array[Byte]): Unit = self ++= as
+    def knownSize: Int                = -1
+  }
 }

zio-http/jvm/src/main/scala/zio/http/netty/NettyBody.scala

@@ -18,18 +18,18 @@ package zio.http.netty
 
 import java.nio.charset.Charset
 
+import zio._
 import zio.stacktracer.TracingImplicits.disableAutoTrace
-import zio.{Chunk, Task, Trace, Unsafe, ZIO}
 
 import zio.stream.ZStream
 
-import zio.http.Body.{UnsafeBytes, UnsafeWriteable}
+import zio.http.Body.UnsafeBytes
 import zio.http.internal.BodyEncoding
-import zio.http.{Body, Boundary, Header, Headers, MediaType}
+import zio.http.{Body, Boundary, MediaType}
 
 import io.netty.buffer.{ByteBuf, ByteBufUtil}
-import io.netty.channel.{Channel => JChannel}
 import io.netty.util.AsciiString
+
 object NettyBody extends BodyEncoding {
 
   /**
@@ -73,7 +73,6 @@ object NettyBody extends BodyEncoding {
     override val mediaType: Option[MediaType] = None,
     override val boundary: Option[Boundary] = None,
   ) extends Body
-      with UnsafeWriteable
       with UnsafeBytes {
 
     override def asArray(implicit trace: Trace): Task[Array[Byte]] = ZIO.succeed(asciiString.array())
@@ -101,11 +100,10 @@ object NettyBody extends BodyEncoding {
   }
 
   private[zio] final case class ByteBufBody(
-    val byteBuf: ByteBuf,
+    byteBuf: ByteBuf,
     override val mediaType: Option[MediaType] = None,
     override val boundary: Option[Boundary] = None,
   ) extends Body
-      with UnsafeWriteable
       with UnsafeBytes {
 
     override def asArray(implicit trace: Trace): Task[Array[Byte]] = ZIO.succeed(ByteBufUtil.getBytes(byteBuf))
@@ -137,38 +135,37 @@ object NettyBody extends BodyEncoding {
     knownContentLength: Option[Long],
     override val mediaType: Option[MediaType] = None,
     override val boundary: Option[Boundary] = None,
-  ) extends Body
-      with UnsafeWriteable {
+  ) extends Body {
+
     override def asArray(implicit trace: Trace): Task[Array[Byte]] = asChunk.map(_.toArray)
 
-    override def asChunk(implicit trace: Trace): Task[Chunk[Byte]] = asStream.runCollect
+    override def asChunk(implicit trace: Trace): Task[Chunk[Byte]] =
+      ZIO.async { cb =>
+        try {
+          // Cap at 100kB as a precaution in case the server sends an invalid content length
+          unsafeAsync(UnsafeAsync.Aggregating(bufferSize(1024 * 100))(cb))
+        } catch {
+          case e: Throwable => cb(ZIO.fail(e))
+        }
+      }
 
     override def asStream(implicit trace: Trace): ZStream[Any, Throwable, Byte] =
       ZStream
         .async[Any, Throwable, Byte](
           emit =>
             try {
-              unsafeAsync(new UnsafeAsync {
-                override def apply(message: Chunk[Byte], isLast: Boolean): Unit = {
-                  emit(ZIO.succeed(message))
-                  if (isLast) {
-                    emit(ZIO.fail(None))
-                  }
-                }
-                override def fail(cause: Throwable): Unit                       =
-                  emit(ZIO.fail(Some(cause)))
-              })
+              unsafeAsync(new UnsafeAsync.Streaming(emit))
             } catch {
               case e: Throwable => emit(ZIO.fail(Option(e)))
             },
-          streamBufferSize,
+          bufferSize(4096),
         )
 
     // No need to create a large buffer when we know the response is small
-    private[this] def streamBufferSize: Int = {
+    private[this] def bufferSize(maxSize: Int): Int = {
       val cl = knownContentLength.getOrElse(4096L)
       if (cl <= 16L) 16
-      else if (cl >= 4096) 4096
+      else if (cl >= maxSize) maxSize
       else Integer.highestOneBit(cl.toInt - 1) << 1 // Round to next power of 2
     }
 
@@ -188,4 +185,30 @@ object NettyBody extends BodyEncoding {
     def apply(message: Chunk[Byte], isLast: Boolean): Unit
     def fail(cause: Throwable): Unit
   }
+
+  private[zio] object UnsafeAsync {
+    private val FailNone = Exit.fail(None)
+
+    final case class Aggregating(bufferInitialSize: Int)(callback: Task[Chunk[Byte]] => Unit)(implicit trace: Trace)
+        extends UnsafeAsync {
+
+      def apply(message: Chunk[Byte], isLast: Boolean): Unit = {
+        assert(isLast)
+        callback(ZIO.succeed(message))
+      }
+
+      def fail(cause: Throwable): Unit =
+        callback(ZIO.fail(cause))
+    }
+
+    final class Streaming(emit: ZStream.Emit[Any, Throwable, Byte, Unit])(implicit trace: Trace) extends UnsafeAsync {
+      def apply(message: Chunk[Byte], isLast: Boolean): Unit = {
+        if (message.nonEmpty) emit(ZIO.succeed(message))
+        if (isLast) emit(FailNone)
+      }
+
+      def fail(cause: Throwable): Unit =
+        emit(ZIO.fail(Some(cause)))
+    }
+  }
 }

zio-http/shared/src/main/scala-2/zio/http/UrlInterpolator.scala

@@ -34,13 +34,8 @@ private[http] object UrlInterpolatorMacro {
         val result = URL.decode(p) match {
           case Left(error) => c.abort(c.enclosingPosition, s"Invalid URL: ${error.getMessage}")
           case Right(url)  =>
-            if (url.isAbsolute) {
-              val uri = url.encode
-              q"_root_.zio.http.URL.fromAbsoluteURI(new _root_.java.net.URI($uri)).get"
-            } else {
-              val uri = url.encode
-              q"_root_.zio.http.URL.fromRelativeURI(new _root_.java.net.URI($uri)).get"
-            }
+            val uri = url.encode
+            q"_root_.zio.http.URL.fromURI(new _root_.java.net.URI($uri)).get"
         }
         c.Expr[URL](result)
       case Apply(_, List(Apply(_, staticPartLiterals)))                  =>
@@ -86,13 +81,7 @@ private[http] object UrlInterpolatorMacro {
                 q"$acc + $part"
               }
 
-            val result = if (url.isAbsolute) {
-              q"_root_.zio.http.URL.fromAbsoluteURI(new _root_.java.net.URI($concatenated)).get"
-            } else {
-              q"_root_.zio.http.URL.fromRelativeURI(new _root_.java.net.URI($concatenated)).get"
-            }
-
-            c.Expr[URL](result)
+            c.Expr[URL](q"_root_.zio.http.URL.fromURI(new _root_.java.net.URI($concatenated)).get")
         }
     }
   }