handle HttpEntity.Default in the Akka HTTP instrumentation

kamon-io · Nov 30, 2021 · d8a5791 · d8a5791
1 parent 558e837
commit d8a5791
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 2 deletions.
diff --git a/...on/kamon-akka-http/src/main/scala/kamon/instrumentation/akka/http/ServerFlowWrapper.scala b/...on/kamon-akka-http/src/main/scala/kamon/instrumentation/akka/http/ServerFlowWrapper.scala
@@ -137,9 +137,14 @@ object ServerFlowWrapper {
             requestSpan.mark("http.response.ready")
 
             responseWithContext.entity match {
-              case strict@HttpEntity.Strict(_, bs) =>
+              case strict @ HttpEntity.Strict(_, bs) =>
                 requestHandler.responseSent(bs.size)
                 strict
+
+              case default: HttpEntity.Default =>
+                requestHandler.responseSent(default.contentLength)
+                default
+
               case _ =>
                 val responseSizeCounter = new AtomicLong(0L)
                 responseWithContext.entity.transformDataBytes(

diff --git a/instrumentation/kamon-akka-http/src/test/scala/kamon/akka/http/ServerFlowWrapperSpec.scala b/instrumentation/kamon-akka-http/src/test/scala/kamon/akka/http/ServerFlowWrapperSpec.scala
@@ -1,9 +1,10 @@
 package kamon.akka.http
 
 import akka.actor.ActorSystem
-import akka.http.scaladsl.model.{HttpEntity, HttpRequest, HttpResponse, StatusCodes}
+import akka.http.scaladsl.model.{ContentTypes, HttpEntity, HttpRequest, HttpResponse, StatusCodes}
 import akka.stream.ActorMaterializer
 import akka.stream.scaladsl.{Flow, Sink, Source}
+import akka.util.ByteString
 import kamon.instrumentation.akka.http.ServerFlowWrapper
 import kamon.testkit.InitAndStopKamonAfterAll
 import org.scalatest.concurrent.ScalaFutures
@@ -20,6 +21,13 @@ class ServerFlowWrapperSpec extends AnyWordSpecLike with Matchers with ScalaFutu
     HttpResponse(status = StatusCodes.OK, entity = HttpEntity("OK"))
   }
 
+  private val defaultReturningFlow = Flow[HttpRequest].map { _ =>
+    HttpResponse(status = StatusCodes.OK, entity = HttpEntity.Default(
+      ContentTypes.`text/plain(UTF-8)`,
+      2,
+      Source.single(ByteString.apply("OK"))))
+  }
+
   "the server flow wrapper" should {
     "keep strict entities strict" in {
       val flow = ServerFlowWrapper(okReturningFlow, "localhost", 8080)
@@ -28,9 +36,23 @@ class ServerFlowWrapperSpec extends AnyWordSpecLike with Matchers with ScalaFutu
         .via(flow)
         .runWith(Sink.head)
         .futureValue
+
       response.entity should matchPattern {
         case HttpEntity.Strict(_, _) =>
       }
     }
+
+    "keep default entities default" in {
+      val flow = ServerFlowWrapper(defaultReturningFlow, "localhost", 8081)
+      val request = HttpRequest()
+      val response = Source.single(request)
+        .via(flow)
+        .runWith(Sink.head)
+        .futureValue
+
+      response.entity should matchPattern {
+        case _: HttpEntity.Default =>
+      }
+    }
   }
 }