From f4fa8823811823faaf680b8565099203cdf13500 Mon Sep 17 00:00:00 2001
From: Amos Ng <amosng1@gmail.com>
Date: Sun, 14 Jun 2020 19:06:10 +0800
Subject: [PATCH] Added response timings to trace logs and response headers

---
 CHANGELOG.md                                  |  1 +
 .../kotlin/mdnet/base/server/Application.kt   |  3 +-
 .../kotlin/mdnet/base/server/ImageServer.kt   |  2 ++
 src/main/kotlin/mdnet/base/server/common.kt   | 30 +++++++++++++++----
 4 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0ce114e..90b1e8e 100755
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 ### Added
+- [2020-06-14] Added response timings to trace logs and response headers by [@lflare].
 
 ### Changed
 
diff --git a/src/main/kotlin/mdnet/base/server/Application.kt b/src/main/kotlin/mdnet/base/server/Application.kt
index 8e7fe7c..196944e 100644
--- a/src/main/kotlin/mdnet/base/server/Application.kt
+++ b/src/main/kotlin/mdnet/base/server/Application.kt
@@ -18,7 +18,8 @@ import java.util.concurrent.atomic.AtomicReference
 fun getServer(cache: DiskLruCache, serverSettings: ServerSettings, clientSettings: ClientSettings, statistics: AtomicReference<Statistics>): Http4kServer {
     val imageServer = ImageServer(cache, statistics, serverSettings.imageServer)
 
-    return catchAllHideDetails()
+    return Timer
+            .then(catchAllHideDetails())
             .then(ServerFilters.CatchLensFailure)
             .then(addCommonHeaders())
             .then(
diff --git a/src/main/kotlin/mdnet/base/server/ImageServer.kt b/src/main/kotlin/mdnet/base/server/ImageServer.kt
index ccbf01d..41e2f8d 100644
--- a/src/main/kotlin/mdnet/base/server/ImageServer.kt
+++ b/src/main/kotlin/mdnet/base/server/ImageServer.kt
@@ -68,8 +68,10 @@ class ImageServer(private val cache: DiskLruCache, private val statistics: Atomi
         val snapshot = cache.get(cacheId)
         if (snapshot != null) {
             request.handleCacheHit(sanitizedUri, getRc4(rc4Bytes), snapshot)
+                .header("X-Uri", sanitizedUri)
         } else {
             request.handleCacheMiss(sanitizedUri, getRc4(rc4Bytes), cacheId)
+                .header("X-Uri", sanitizedUri)
         }
     }
 
diff --git a/src/main/kotlin/mdnet/base/server/common.kt b/src/main/kotlin/mdnet/base/server/common.kt
index 38f1f1c..0a6a0b4 100644
--- a/src/main/kotlin/mdnet/base/server/common.kt
+++ b/src/main/kotlin/mdnet/base/server/common.kt
@@ -1,6 +1,10 @@
 /* ktlint-disable no-wildcard-imports */
 package mdnet.base.server
 
+import java.time.ZoneOffset
+import java.time.ZonedDateTime
+import java.time.format.DateTimeFormatter
+import java.util.*
 import mdnet.base.Constants
 import org.http4k.core.Filter
 import org.http4k.core.HttpHandler
@@ -8,13 +12,8 @@ import org.http4k.core.Request
 import org.http4k.core.Response
 import org.http4k.core.Status
 import org.slf4j.LoggerFactory
-import java.time.ZoneOffset
-import java.time.ZonedDateTime
-import java.time.format.DateTimeFormatter
-import java.util.*
 
 private val HTTP_TIME_FORMATTER = DateTimeFormatter.ofPattern("EEE, dd MMM yyyy HH:mm:ss O", Locale.ENGLISH)
-
 private val LOGGER = LoggerFactory.getLogger("Application")
 
 fun addCommonHeaders(): Filter {
@@ -41,3 +40,24 @@ fun catchAllHideDetails(): Filter {
         }
     }
 }
+
+val Timer = Filter {
+    next: HttpHandler -> {
+        request: Request ->
+            val start = System.currentTimeMillis()
+            val response = next(request)
+            val latency = System.currentTimeMillis() - start
+            if (LOGGER.isTraceEnabled && response.header("X-Uri") != null) {
+                // Dirty hack to get sanitizedUri from ImageServer
+                val sanitizedUri = response.header("X-Uri")
+
+                // Log in TRACE
+                LOGGER.trace("Request for $sanitizedUri completed in ${latency}ms")
+
+                // Delete response header entirely
+                response.header("X-Uri", null)
+            }
+            // Set response header with processing times
+            response.header("X-Time-Taken", latency.toString())
+    }
+}