Add `Connection: close` header to all responses in order to mitigate connection pool issues with envoy (#175)v1.11.2

2 files changed, 24 insertions, 2 deletions

diff --git a/src/main/scala/xyz/driver/core/rest/DriverRoute.scala b/src/main/scala/xyz/driver/core/rest/DriverRoute.scala
index fe3790f..55f39ba 100644
--- a/src/main/scala/xyz/driver/core/rest/DriverRoute.scala
+++ b/src/main/scala/xyz/driver/core/rest/DriverRoute.scala
@@ -32,7 +32,11 @@ trait DriverRoute {
       val tracingHeader = RawHeader(ContextHeaders.TrackingIdHeader, trackingId)
       MDC.put("trackingId", trackingId)
 
-      respondWithHeader(tracingHeader)
+      // This header will eliminate the risk of LB trying to reuse a connection
+      // that already timed out on the server side by completely rejecting keep-alive
+      val rejectKeepAlive = Connection("close")
+
+      respondWithHeaders(tracingHeader, rejectKeepAlive)
     }
   }
 
@@ -103,4 +107,5 @@ trait DriverRoute {
   protected def errorResponse[T <: Exception](statusCode: StatusCode, entity: ResponseEntity, exception: T): Route = {
     complete(HttpResponse(statusCode, entity = entity))
   }
+
 }
diff --git a/src/test/scala/xyz/driver/core/rest/DriverRouteTest.scala b/src/test/scala/xyz/driver/core/rest/DriverRouteTest.scala
index aca8fdc..d32fefd 100644
--- a/src/test/scala/xyz/driver/core/rest/DriverRouteTest.scala
+++ b/src/test/scala/xyz/driver/core/rest/DriverRouteTest.scala
@@ -2,8 +2,9 @@ package xyz.driver.core.rest
 
 import akka.http.scaladsl.marshallers.sprayjson.SprayJsonSupport
 import akka.http.scaladsl.model.StatusCodes
+import akka.http.scaladsl.model.headers.Connection
 import akka.http.scaladsl.server.Directives.{complete => akkaComplete}
-import akka.http.scaladsl.server.{Directives, Route}
+import akka.http.scaladsl.server.{Directives, Rejection, RejectionHandler, Route}
 import akka.http.scaladsl.testkit.ScalatestRouteTest
 import com.typesafe.scalalogging.Logger
 import org.scalatest.{AsyncFlatSpec, Matchers}
@@ -12,6 +13,7 @@ import xyz.driver.core.json.serviceExceptionFormat
 import xyz.driver.core.FutureExtensions
 import xyz.driver.core.rest.errors._
 
+import scala.collection.immutable
 import scala.concurrent.Future
 
 class DriverRouteTest
@@ -103,4 +105,19 @@ class DriverRouteTest
       responseAs[ServiceException] shouldBe DatabaseException()
     }
   }
+
+  it should "add a `Connection: close` header to avoid clashing with envoy's timeouts" in {
+    val rejectionHandler = RejectionHandler.newBuilder().handleNotFound(complete(StatusCodes.NotFound)).result()
+    val route            = new TestRoute(handleRejections(rejectionHandler)((get & path("foo"))(complete("OK"))))
+
+    Get("/foo") ~> route.routeWithDefaults ~> check {
+      status shouldBe StatusCodes.OK
+      headers should contain(Connection("close"))
+    }
+
+    Get("/bar") ~> route.routeWithDefaults ~> check {
+      status shouldBe StatusCodes.NotFound
+      headers should contain(Connection("close"))
+    }
+  }
 }