Fix metrics timers (#412)

* Fix metrics timers to record at correct time

* Add ResponseBody.withPostWriteClosure

* Give withPostWriteClosure package scope

Sources/Hummingbird/Middleware/MetricsMiddleware.swift

@@ -26,19 +26,21 @@ public struct MetricsMiddleware<Context: BaseRequestContext>: RouterMiddleware {
         let startTime = DispatchTime.now().uptimeNanoseconds
 
         do {
-            let response = try await next(request, context)
-            // need to create dimensions once request has been responded to ensure
-            // we have the correct endpoint path
-            let dimensions: [(String, String)] = [
-                ("hb_uri", context.endpointPath ?? request.uri.path),
-                ("hb_method", request.method.rawValue),
-            ]
-            Counter(label: "hb_requests", dimensions: dimensions).increment()
-            Metrics.Timer(
-                label: "hb_request_duration",
-                dimensions: dimensions,
-                preferredDisplayUnit: .seconds
-            ).recordNanoseconds(DispatchTime.now().uptimeNanoseconds - startTime)
+            var response = try await next(request, context)
+            response.body = response.body.withPostWriteClosure {
+                // need to create dimensions once request has been responded to ensure
+                // we have the correct endpoint path
+                let dimensions: [(String, String)] = [
+                    ("hb_uri", context.endpointPath ?? request.uri.path),
+                    ("hb_method", request.method.rawValue),
+                ]
+                Counter(label: "hb_requests", dimensions: dimensions).increment()
+                Metrics.Timer(
+                    label: "hb_request_duration",
+                    dimensions: dimensions,
+                    preferredDisplayUnit: .seconds
+                ).recordNanoseconds(DispatchTime.now().uptimeNanoseconds - startTime)
+            }
             return response
         } catch {
             // need to create dimensions once request has been responded to ensure

Sources/HummingbirdCore/Response/ResponseBody.swift

@@ -71,6 +71,26 @@ public struct ResponseBody: Sendable {
         self.init(contentLength: contentLength, write: write)
     }
 
+    /// Create new response body that call a callback once original response body has been written
+    /// to the channel
+    ///
+    /// When you return a response from a handler, this cannot be considered to be the point the
+    /// response was written. This functions provides you a method for catching the point when the
+    /// response has been fully written. If you drop the response in a middleware run after this
+    /// point the post write closure will not get run.
+    package func withPostWriteClosure(_ postWrite: @escaping @Sendable () async -> Void) -> Self {
+        return .init(contentLength: self.contentLength) { writer in
+            do {
+                let result = try await self.write(writer)
+                await postWrite()
+                return result
+            } catch {
+                await postWrite()
+                throw error
+            }
+        }
+    }
+
     /// Initialise ResponseBody with closure writing body contents
     ///
     /// This version of init is private and only available via ``withTrailingHeaders`` because

Tests/HummingbirdTests/MetricsTests.swift

@@ -261,4 +261,21 @@ final class MetricsTests: XCTestCase {
         XCTAssertEqual(counter.dimensions[1].0, "hb_method")
         XCTAssertEqual(counter.dimensions[1].1, "GET")
     }
+
+    func testRecordingBodyWriteTime() async throws {
+        let router = Router()
+        router.middlewares.add(MetricsMiddleware())
+        router.get("/hello") { _, _ -> Response in
+            return Response(status: .ok, body: .init { _ in
+                try await Task.sleep(for: .milliseconds(5))
+            })
+        }
+        let app = Application(responder: router.buildResponder())
+        try await app.test(.router) { client in
+            try await client.execute(uri: "/hello", method: .get) { _ in }
+        }
+
+        let timer = try XCTUnwrap(Self.testMetrics.timers["hb_request_duration"] as? TestTimer)
+        XCTAssertGreaterThan(timer.values[0].1, 5_000_000)
+    }
 }