Repository: http4s/http4s
Pull request #6980 (2023-02-18 06:39:00)

Detect terminated `EmberConnection`s in ember client prior to attempting a request
Fixes the reproducer in https://github.com/http4s/http4s/issues/4935#issuecomment-1435333053, and hopefully will allow us to close that issue once and for all 🙏

`EmberConnection` now pro-actively starts the next read from a socket, upon completing a request/response.

In a normal situation, this read will hang until the next request has been sent, at which point it will become the head of the response.

However, in the situation that the server decides to end the connection (probably while the socket is sitting idle in the connection pool), this read will terminate early with EOF. By capturing this, we are able to find out if the socket is still alive _prior_ to starting a request. This allows us to avoid making a dud request (that we may not be able to retry) on a dead socket.


Patch:
diff --git a/ember-client/shared/src/main/scala/org/http4s/ember/client/EmberClientBuilder.scala b/ember-client/shared/src/main/scala/org/http4s/ember/client/EmberClientBuilder.scala
index e5dc6c6a48..46be07730d 100644
--- a/ember-client/shared/src/main/scala/org/http4s/ember/client/EmberClientBuilder.scala
+++ b/ember-client/shared/src/main/scala/org/http4s/ember/client/EmberClientBuilder.scala
@@ -247,24 +247,25 @@ final class EmberClientBuilder[F[_]: Async: Network] private (
           )
       builder =
         KeyPool.Builder
-          .apply[F, RequestKey, EmberConnection[F]](
-            (requestKey: RequestKey) =>
-              EmberConnection(
-                org.http4s.ember.client.internal.ClientHelpers
-                  .requestKeyToSocketWithKey[F](
-                    requestKey,
-                    tlsContextOptWithDefault,
-                    checkEndpointIdentification,
-                    serverNameIndication,
-                    sg,
-                    additionalSocketOptions,
-                  )
-              ) <* logger.trace(s"Created Connection - RequestKey: ${requestKey}"),
-            (connection: EmberConnection[F]) =>
-              logger.trace(
-                s"Shutting Down Connection - RequestKey: ${connection.keySocket.requestKey}"
-              ) >>
-                connection.cleanup,
+          .apply[F, RequestKey, EmberConnection[F]]((requestKey: RequestKey) =>
+            EmberConnection(
+              org.http4s.ember.client.internal.ClientHelpers
+                .requestKeyToSocketWithKey[F](
+                  requestKey,
+                  tlsContextOptWithDefault,
+                  checkEndpointIdentification,
+                  serverNameIndication,
+                  sg,
+                  additionalSocketOptions,
+                ),
+              chunkSize,
+            ) <* Resource
+              .eval(logger.trace(s"Created Connection - RequestKey: ${requestKey}"))
+              .onFinalize(
+                logger.trace(
+                  s"Shutting Down Connection - RequestKey: ${requestKey}"
+                )
+              )
           )
           .withDefaultReuseState(Reusable.DontReuse)
           .withIdleTimeAllowedInPool(idleTimeInPool)
@@ -317,6 +318,7 @@ final class EmberClientBuilder[F[_]: Async: Network] private (
                   drain,
                   managed.value.nextBytes,
                   managed.canBeReused,
+                  managed.value.startNextRead,
                 )
               case _ => Applicative[F].unit
             }
@@ -340,19 +342,17 @@ final class EmberClientBuilder[F[_]: Async: Network] private (
               )
           )
           .flatMap(unixSockets =>
-            Resource
-              .make(
-                EmberConnection(
-                  ClientHelpers.unixSocket(
-                    request,
-                    unixSockets,
-                    address,
-                    tlsContextOpt,
-                    enableEndpointValidation,
-                    enableServerNameIndication,
-                  )
-                )
-              )(ec => ec.shutdown)
+            EmberConnection(
+              ClientHelpers.unixSocket(
+                request,
+                unixSockets,
+                address,
+                tlsContextOpt,
+                enableEndpointValidation,
+                enableServerNameIndication,
+              ),
+              chunkSize,
+            )
           )
           .flatMap(connection =>
             Resource.eval(
diff --git a/ember-client/shared/src/main/scala/org/http4s/ember/client/EmberConnection.scala b/ember-client/shared/src/main/scala/org/http4s/ember/client/EmberConnection.scala
index 5d96f5c320..7ce88422ce 100644
--- a/ember-client/shared/src/main/scala/org/http4s/ember/client/EmberConnection.scala
+++ b/ember-client/shared/src/main/scala/org/http4s/ember/client/EmberConnection.scala
@@ -16,17 +16,55 @@
 
 package org.http4s.ember.client
 
-import cats._
 import cats.effect.Concurrent
 import cats.effect.Resource
+import cats.effect.kernel.Deferred
 import cats.effect.kernel.Ref
+import cats.effect.std.Hotswap
 import cats.syntax.all._
+import fs2.Chunk
 
 private[ember] final case class EmberConnection[F[_]](
     keySocket: RequestKeySocket[F],
+    chunkSize: Int,
     shutdown: F[Unit],
     nextBytes: Ref[F, Array[Byte]],
-)(implicit F: MonadThrow[F]) {
+    /* These variables store the currently running read operation and its result.
+     * The next read is started pre-emptively so that we may receieve EOF from server ASAP.
+     * On the happy path when we reuse a connection, this will block until the next request is sent.
+     * On the less-happy path, this read will complete with EOF, which we check before sending a request.
+     */
+    hotRead: Hotswap[F, Deferred[F, Either[Throwable, Option[Chunk[Byte]]]]],
+    nextRead: Ref[F, Deferred[F, Either[Throwable, Option[Chunk[Byte]]]]],
+)(implicit F: Concurrent[F]) {
+
+  /** For the connection to be valid, the socket must be open,
+    * and its pre-emptive read must not have terminated in an error or EOF.
+    */
+  def isValid: F[Boolean] = {
+    val isOpen = keySocket.socket.isOpen
+    val isEof = nextRead.get.flatMap(_.tryGet).map {
+      case Some(result) => result.fold(_ => true, _.isEmpty) // if Left or None this socket is dead
+      case None => false // no read yet, which is good!
+    }
+    (isOpen, isEof).mapN((open, eof) => open && !eof)
+  }
+
+  /** We must start the next read after completing a request/response pair,
+    * and before returning this connection to the pool.
+    *
+    * This way [[isValid]] can check for EOF when the connection is retrieved from the pool.
+    */
+  def startNextRead: F[Unit] =
+    hotRead
+      .swap {
+        Resource.eval(F.deferred[Either[Throwable, Option[Chunk[Byte]]]]).flatTap { result =>
+          val read = keySocket.socket.read(chunkSize)
+          F.background(read.attempt.flatMap(result.complete(_)).void.voidError)
+        }
+      }
+      .flatMap(nextRead.set(_))
+
   def cleanup: F[Unit] =
     nextBytes.set(Array.emptyByteArray) >>
       keySocket.socket.endOfInput.attempt.void >>
@@ -35,14 +73,24 @@ private[ember] final case class EmberConnection[F[_]](
 }
 
 private[ember] object EmberConnection {
-  def apply[F[_]: Concurrent](
-      keySocketResource: Resource[F, RequestKeySocket[F]]
-  ): F[EmberConnection[F]] =
-    Ref[F].of(Array.emptyByteArray).flatMap { nextBytes =>
-      val keySocketResourceAllocated: F[(RequestKeySocket[F], F[Unit])] =
-        keySocketResource.allocated
-      keySocketResourceAllocated.map { case (keySocket, release) =>
-        EmberConnection(keySocket, release, nextBytes)
-      }
+  def apply[F[_]](
+      keySocketResource: Resource[F, RequestKeySocket[F]],
+      chunkSize: Int,
+  )(implicit F: Concurrent[F]): Resource[F, EmberConnection[F]] =
+    (
+      Resource.eval(keySocketResource.allocated),
+      Resource.eval(F.ref(Array.emptyByteArray)),
+      Hotswap.create[F, Deferred[F, Either[Throwable, Option[Chunk[Byte]]]]],
+      Resource.eval(
+        F.deferred[Either[Throwable, Option[Chunk[Byte]]]]
+          .flatTap(_.complete(Right(Some(Chunk.empty))))
+          .flatMap(F.ref(_))
+      ),
+    ).flatMapN { case ((keySocket, release), nextBytes, hotRead, nextRead) =>
+      Resource.make(
+        F.pure(EmberConnection(keySocket, chunkSize, release, nextBytes, hotRead, nextRead))
+      )(
+        _.cleanup
+      )
     }
 }
diff --git a/ember-client/shared/src/main/scala/org/http4s/ember/client/internal/ClientHelpers.scala b/ember-client/shared/src/main/scala/org/http4s/ember/client/internal/ClientHelpers.scala
index b39ec91a11..521445a2df 100644
--- a/ember-client/shared/src/main/scala/org/http4s/ember/client/internal/ClientHelpers.scala
+++ b/ember-client/shared/src/main/scala/org/http4s/ember/client/internal/ClientHelpers.scala
@@ -29,6 +29,7 @@ import cats.effect.kernel.Concurrent
 import cats.effect.kernel.Ref
 import cats.effect.kernel.Resource
 import cats.effect.kernel.Sync
+import cats.effect.syntax.all._
 import cats.syntax.all._
 import com.comcast.ip4s.Host
 import com.comcast.ip4s.Port
@@ -162,25 +163,32 @@ private[client] object ClientHelpers {
         .drain
 
     def writeRead(req: Request[F]): F[(Response[F], F[Option[Array[Byte]]])] =
-      writeRequestToSocket(req, connection.keySocket.socket) >>
-        connection.nextBytes.getAndSet(Array.emptyByteArray).flatMap { head =>
-          val parse = Parser.Response.parser(maxResponseHeaderSize)(
-            head,
+      writeRequestToSocket(req, connection.keySocket.socket).productR {
+        val parse = (
+          // leftover bytes from last request
+          connection.nextBytes.getAndSet(Array.emptyByteArray),
+          // bytes from the pre-emptive read in this connection, which probably has not completed yet, hence the timout
+          connection.nextRead.get.flatMap(_.get).rethrow.timeout(idleTimeout),
+        ).flatMapN { (head, firstRead) =>
+          Parser.Response.parser(maxResponseHeaderSize)(
+            firstRead.foldLeft(head)(Util.concatBytes(_, _)),
             timeoutMaybe(connection.keySocket.socket.read(chunkSize), idleTimeout),
           )
-          timeoutToMaybe(
-            parse,
-            timeout,
-            Defer[F].defer(
-              ApplicativeThrow[F].raiseError(
-                new java.util.concurrent.TimeoutException(
-                  s"Timed Out on EmberClient Header Receive Timeout: $timeout"
-                )
-              )
-            ),
-          )
         }
 
+        timeoutToMaybe(
+          parse,
+          timeout,
+          Defer[F].defer(
+            ApplicativeThrow[F].raiseError(
+              new java.util.concurrent.TimeoutException(
+                s"Timed Out on EmberClient Header Receive Timeout: $timeout"
+              )
+            )
+          ),
+        )
+      }
+
     for {
       processedReq <- preprocessRequest(request, userAgent)
       res <- writeRead(processedReq)
@@ -213,6 +221,7 @@ private[client] object ClientHelpers {
       drain: F[Option[Array[Byte]]],
       nextBytes: Ref[F, Array[Byte]],
       canBeReused: Ref[F, Reusable],
+      startNextRead: F[Unit],
   )(implicit F: Concurrent[F]): F[Unit] =
     drain.flatMap {
       case Some(bytes) =>
@@ -220,7 +229,10 @@ private[client] object ClientHelpers {
         val responseClose = connectionFor(resp.httpVersion, resp.headers).hasClose
 
         if (requestClose || responseClose) F.unit
-        else nextBytes.set(bytes) >> canBeReused.set(Reusable.Reuse)
+        else
+          nextBytes.set(bytes) *>
+            startNextRead *> // start the next read before returning to pool
+            canBeReused.set(Reusable.Reuse) // now it is safe to mark as re-usable
       case None => F.unit
     }
 
@@ -242,7 +254,7 @@ private[client] object ClientHelpers {
   ): Resource[F, Managed[F, EmberConnection[F]]] =
     pool.take(RequestKey.fromRequest(request)).flatMap { managed =>
       Resource
-        .eval(managed.value.keySocket.socket.isOpen)
+        .eval(managed.value.isValid)
         .ifM(
           managed.pure[Resource[F, *]],
           // Already Closed,
diff --git a/ember-core/shared/src/main/scala/org/http4s/ember/core/Parser.scala b/ember-core/shared/src/main/scala/org/http4s/ember/core/Parser.scala
index 7cb77e4f0d..a66d4b681a 100644
--- a/ember-core/shared/src/main/scala/org/http4s/ember/core/Parser.scala
+++ b/ember-core/shared/src/main/scala/org/http4s/ember/core/Parser.scala
@@ -25,7 +25,6 @@ import org.http4s._
 import org.typelevel.ci.CIString
 import scodec.bits.ByteVector
 
-import java.util.Arrays
 import scala.annotation.switch
 import scala.util.control.NonFatal
 
@@ -48,7 +47,7 @@ private[ember] object Parser {
           else {
             read.flatMap {
               case Some(chunk) =>
-                val nextBuffer = concatBytes(currentBuffer, chunk)
+                val nextBuffer = Util.concatBytes(currentBuffer, chunk)
                 recurseFind(nextBuffer, read, maxHeaderSize, s)(f)(idx)
               case None if currentBuffer.length > 0 =>
                 F.raiseError(EmberException.ReachedEndOfStream())
@@ -631,22 +630,6 @@ private[ember] object Parser {
     }
   }
 
-  private[this] def concatBytes(a1: Array[Byte], a2: Chunk[Byte]): Array[Byte] =
-    if (a1.length == 0) {
-      a2 match {
-        case slice: Chunk.ArraySlice[Byte]
-            if slice.values.isInstanceOf[Array[Byte]] &&
-              slice.offset == 0 &&
-              slice.values.length == slice.length =>
-          slice.values
-        case _ => a2.toArray
-      }
-    } else {
-      val res = Arrays.copyOf(a1, a1.length + a2.size)
-      a2.copyToArray(res, a1.length)
-      res
-    }
-
   private[this] final val space = 32 // ' '
   private[this] final val cr = 13 // '\r'
   private[this] final val lf = 10 // '\n'
diff --git a/ember-core/shared/src/main/scala/org/http4s/ember/core/Util.scala b/ember-core/shared/src/main/scala/org/http4s/ember/core/Util.scala
index 6b8891d119..32b4a1a99b 100644
--- a/ember-core/shared/src/main/scala/org/http4s/ember/core/Util.scala
+++ b/ember-core/shared/src/main/scala/org/http4s/ember/core/Util.scala
@@ -28,6 +28,7 @@ import org.http4s.headers.Connection
 import org.typelevel.ci._
 
 import java.time.Instant
+import java.util.Arrays
 import java.util.Locale
 import scala.concurrent.duration._
 
@@ -129,4 +130,20 @@ private[ember] object Util extends UtilPlatform {
     }
   }
 
+  def concatBytes(a1: Array[Byte], a2: Chunk[Byte]): Array[Byte] =
+    if (a1.length == 0) {
+      a2 match {
+        case slice: Chunk.ArraySlice[Byte]
+            if slice.values.isInstanceOf[Array[Byte]] &&
+              slice.offset == 0 &&
+              slice.values.length == slice.length =>
+          slice.values
+        case _ => a2.toArray
+      }
+    } else {
+      val res = Arrays.copyOf(a1, a1.length + a2.size)
+      a2.copyToArray(res, a1.length)
+      res
+    }
+
 }


Test patch:
diff --git a/ember-client/shared/src/test/scala/org/http4s/ember/client/internal/ClientHelpersSuite.scala b/ember-client/shared/src/test/scala/org/http4s/ember/client/internal/ClientHelpersSuite.scala
index 02140a1778..2ee44a2283 100644
--- a/ember-client/shared/src/test/scala/org/http4s/ember/client/internal/ClientHelpersSuite.scala
+++ b/ember-client/shared/src/test/scala/org/http4s/ember/client/internal/ClientHelpersSuite.scala
@@ -115,6 +115,7 @@ class ClientHelpersSuite extends Http4sSuite {
           IO.pure(Some(Array.emptyByteArray)),
           nextBytes,
           reuse,
+          IO.unit,
         )
       testResult <- reuse.get.map { case r =>
         assertEquals(r, Reusable.Reuse)
@@ -133,6 +134,7 @@ class ClientHelpersSuite extends Http4sSuite {
         IO.pure(Some(Array[Byte](1, 2, 3))),
         nextBytes,
         reuse,
+        IO.unit,
       )
       drained <- nextBytes.get
     } yield assertEquals(drained.toList, List[Byte](1, 2, 3))
@@ -149,6 +151,7 @@ class ClientHelpersSuite extends Http4sSuite {
           IO.pure(Some(Array.emptyByteArray)),
           nextBytes,
           reuse,
+          IO.unit,
         )
       testResult <- reuse.get.map { case r =>
         assertEquals(r, Reusable.DontReuse)
@@ -167,6 +170,7 @@ class ClientHelpersSuite extends Http4sSuite {
           IO.pure(Some(Array.emptyByteArray)),
           nextBytes,
           reuse,
+          IO.unit,
         )
       testResult <- reuse.get.map { case r =>
         assertEquals(r, Reusable.DontReuse)
@@ -185,6 +189,7 @@ class ClientHelpersSuite extends Http4sSuite {
           IO.pure(None),
           nextBytes,
           reuse,
+          IO.unit,
         )
       testResult <- reuse.get.map { case r =>
         assertEquals(r, Reusable.DontReuse)
diff --git a/ember-server/shared/src/test/scala/org/http4s/ember/server/EmberServerSuite.scala b/ember-server/shared/src/test/scala/org/http4s/ember/server/EmberServerSuite.scala
index 19981eb31c..565f9d9624 100644
--- a/ember-server/shared/src/test/scala/org/http4s/ember/server/EmberServerSuite.scala
+++ b/ember-server/shared/src/test/scala/org/http4s/ember/server/EmberServerSuite.scala
@@ -23,7 +23,6 @@ import fs2.Stream
 import fs2.io.net.BindException
 import fs2.io.net.ConnectException
 import org.http4s._
-import org.http4s.dsl.Http4sDsl
 import org.http4s.ember.client.EmberClientBuilder
 import org.http4s.ember.core.EmberException
 import org.http4s.implicits._
@@ -33,24 +32,25 @@ import scala.concurrent.duration._
 
 class EmberServerSuite extends Http4sSuite {
 
-  def service[F[_]](implicit F: Async[F]): HttpApp[F] = {
-    val dsl = new Http4sDsl[F] {}
-    import dsl._
+  def service: HttpApp[IO] = {
+    import org.http4s.dsl.io._
 
     HttpRoutes
-      .of[F] {
+      .of[IO] {
         case GET -> Root =>
           Ok("Hello!")
         case req @ POST -> Root / "echo" =>
           Ok(req.body)
         case GET -> Root / "failed-stream" =>
-          Ok(Stream.raiseError[F](new RuntimeException("BOOM")).covaryOutput[String])
+          Ok(Stream.raiseError[IO](new RuntimeException("BOOM")).covaryOutput[String])
       }
       .orNotFound
   }
 
-  def url(address: SocketAddress[Host], path: String = ""): String =
-    s"http://${Uri.Host.fromIp4sHost(address.host).renderString}:${address.port.value}$path"
+  def url(address: SocketAddress[Host], path: String = ""): Uri =
+    Uri.unsafeFromString(
+      s"http://${Uri.Host.fromIp4sHost(address.host).renderString}:${address.port.value}$path"
+    )
 
   def serverResource(
       f: EmberServerBuilder[IO] => EmberServerBuilder[IO] = identity
@@ -59,7 +59,7 @@ class EmberServerSuite extends Http4sSuite {
       EmberServerBuilder
         .default[IO]
         .withPort(port"0")
-        .withHttpApp(service[IO])
+        .withHttpApp(service)
     ).build
 
   private val client = ResourceFunFixture(EmberClientBuilder.default[IO].build)
@@ -67,7 +67,7 @@ class EmberServerSuite extends Http4sSuite {
   private def server(receiveBufferSize: Int = 256 * 1024) = ResourceFunFixture(
     EmberServerBuilder
       .default[IO]
-      .withHttpApp(service[IO])
+      .withHttpApp(service)
       .withPort(port"0")
       .withReceiveBufferSize(receiveBufferSize)
       .build
@@ -83,9 +83,7 @@ class EmberServerSuite extends Http4sSuite {
   }
 
   fixture().test("connection closed when response body stream fails") { case (server, client) =>
-    import org.http4s.dsl.io._
-    import org.http4s.client.dsl.io._
-    val req = GET(Uri.unsafeFromString(url(server.addressIp4s, "/failed-stream")))
+    val req = Request[IO](uri = url(server.addressIp4s, "/failed-stream"))
     client
       .stream(req)
       .flatMap(_.body)
@@ -133,14 +131,27 @@ class EmberServerSuite extends Http4sSuite {
         Stream.emit("hello").repeatN(256).through(fs2.text.utf8.encode)
       val expected = "hello" * 256
 
-      val uri = Uri
-        .fromString(url(server.addressIp4s, "/echo"))
-        .toOption
-        .get
+      val uri = url(server.addressIp4s, "/echo")
       val request = POST(body, uri)
       for {
         r1 <- client.fetchAs[String](request)
         r2 <- client.fetchAs[String](request)
       } yield assertEquals(expected, r1) && assertEquals(expected, r2)
   }
+
+  client.test("#4935 - client can detect a terminated connection") { client =>
+    def runReq(server: Server) = {
+      val req =
+        Request[IO](Method.POST, uri = url(server.addressIp4s, "/echo")).withEntity("Hello!")
+      client.expect[String](req).assertEquals("Hello!")
+    }
+
+    serverResource(_.withShutdownTimeout(0.nanos))
+      .use(server => runReq(server).as(server.addressIp4s.port))
+      .flatMap { port =>
+        IO.sleep(1.second) *> // so server shutdown propagates
+          serverResource(_.withPort(port).withShutdownTimeout(0.nanos)).use(runReq(_))
+      }
+  }
+
 }