Working on @kristomasette comments

caoilte · Jan 11, 2016 · 34e147a · 34e147a
1 parent 6d9f36a
commit 34e147a
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 12 deletions.
diff --git a/src/main/scala/org/caoilte/spray/routing/LogAccessRouting.scala b/src/main/scala/org/caoilte/spray/routing/LogAccessRouting.scala
@@ -40,6 +40,12 @@ class RequestAccessLogger(ctx: RequestContext, accessLogger: AccessLogger,
   var cancellable:Cancellable = context.system.scheduler.scheduleOnce(requestTimeout, self, RequestLoggingTimeout)
   var chunkedResponse:HttpResponse = null
 
+  def handleChunkedResponseStart(wrapper: Any, response: HttpResponse) = {
+    // In case a chunk transfer is started, store the response for later logging
+    chunkedResponse = response
+    responder forward wrapper
+  }
+
   def receive = {
     case RequestLoggingTimeout => {
       val errorResponse = HttpResponse(
@@ -56,27 +62,27 @@ class RequestAccessLogger(ctx: RequestContext, accessLogger: AccessLogger,
       accessLogger.logAccess(request, response, timeStampCalculator(Unit))
       forwardMsgAndStop(response)
     }
-    case Confirmed(ChunkedResponseStart(response), sentAck) => {
-      // In case a chunk transfer is started, store the response for later logging
-      chunkedResponse = response
-      responder forward Confirmed(ChunkedResponseStart(response), sentAck)
-    }
-    case Confirmed(MessageChunk(data, extension), sentAck) => {
-      // Each junk is considered a heart-beat: We are still producing output and the client is still consuming
+
+    case confirmed@Confirmed(ChunkedResponseStart(response), _) => handleChunkedResponseStart(confirmed, response)
+    case unconfirmed@ChunkedResponseStart(response) => handleChunkedResponseStart(unconfirmed, response)
+
+    case confirmed@(Confirmed(MessageChunk(_,_), _) | MessageChunk) => {
+      // Each chunk is considered a heart-beat: We are still producing output and the client is still consuming
       // - therefore reset the requestTimeout schedule
       cancellable.cancel()
       cancellable = context.system.scheduler.scheduleOnce(requestTimeout, self, RequestLoggingTimeout)
 
-      responder forward Confirmed(MessageChunk(data, extension), sentAck)
+      responder forward confirmed
     }
-    case response:ChunkedMessageEnd => {
+
+    case confirmed@(Confirmed(ChunkedMessageEnd, _) | ChunkedMessageEnd) => {
       // Handled like HttpResponse: provide the (previously saved) ChunkedResponseStart for logging and quit
       cancellable.cancel()
       accessLogger.logAccess(request, chunkedResponse, timeStampCalculator(Unit))
-      forwardMsgAndStop(response)
+      forwardMsgAndStop(confirmed)
     }
+
     case other => {
-      println(s"other $other")
       responder forward other
     }
   }

diff --git a/src/test/scala/org/caoilte/spray/routing/TestLogAccessRoutingActor.scala b/src/test/scala/org/caoilte/spray/routing/TestLogAccessRoutingActor.scala
@@ -58,7 +58,7 @@ class ChunkedResponseServiceActor(val accessLogger: AccessLogger, chunkedRespons
         ctx.responder ! ChunkedResponseStart(HttpResponse()).withAck(Ok(ctx, amountOfChunks))
       }
       case Ok(ctx, 0) => {
-        ctx.responder ! ChunkedMessageEnd
+        ctx.responder ! ChunkedMessageEnd.withAck(new Object)
       }
       case Ok(ctx, remainingChunks) => {
         blocking {