Use a response writer at the Server level (#519)

* Use a response writer at the server level

* Comments, rename symbols

* Update Sources/HummingbirdCore/Response/ResponseWriter.swift

Co-authored-by: Joannis Orlandos <[email protected]>

* ResponseWriter second attempt

ResponseWriter returns a ResponseBodyWriter when headers are written.

* Comments

* Make ResponseWriter non Copyable

* Update Sources/HummingbirdCore/Response/ResponseWriter.swift

Co-authored-by: Joannis Orlandos <[email protected]>

* Fix trailing headers, by requiring user to finish ResponseBodyWriter

* comments

* formatting

---------

Co-authored-by: Joannis Orlandos <[email protected]>

Sources/Hummingbird/Application.swift

@@ -15,6 +15,7 @@
 import HummingbirdCore
 import Logging
 import NIOCore
+import NIOHTTPTypes
 import NIOPosix
 import NIOTransportServices
 import ServiceLifecycle
@@ -100,7 +101,7 @@ extension ApplicationProtocol {
             configuration: self.configuration.httpServer,
             eventLoopGroup: self.eventLoopGroup,
             logger: self.logger
-        ) { request, channel in
+        ) { (request, responseWriter: consuming ResponseWriter, channel) in
             let logger = self.logger.with(metadataKey: "hb_id", value: .stringConvertible(RequestID()))
             let context = Self.Responder.Context(
                 source: .init(
@@ -124,7 +125,9 @@ extension ApplicationProtocol {
             if let serverName = self.configuration.serverName {
                 response.headers[.server] = serverName
             }
-            return response
+            // Write response
+            let bodyWriter = try await responseWriter.writeHead(response.head)
+            try await response.body.write(bodyWriter)
         } onServerRunning: {
             await self.onServerRunning($0)
         }

Sources/HummingbirdCore/Response/ResponseBody.swift

@@ -17,28 +17,44 @@ import NIOCore
 
 /// Response body
 public struct ResponseBody: Sendable {
-    public let write: @Sendable (any ResponseBodyWriter) async throws -> HTTPFields?
+    @usableFromInline
+    let _write: @Sendable (any ResponseBodyWriter) async throws -> Void
     public let contentLength: Int?
 
-    /// Initialise ResponseBody with closure writing body contents
+    /// Initialise ResponseBody with closure writing body contents.
+    ///
+    /// When you have finished writing the response body you need to indicate you
+    /// have finished by calling ``ResponseBodyWriter.finish``. At this point you can also
+    /// send trailing headers by including them as a parameter in the finsh() call.
+    /// ```
+    /// let responseBody = ResponseBody(contentLength: contentLength) { writer in
+    ///     try await writer.write(buffer)
+    ///     try await writer.finish()
+    /// }
+    /// ```
     /// - Parameters:
     ///   - contentLength: Optional length of body
     ///   - write: closure provided with `writer` type that can be used to write to response body
     public init(contentLength: Int? = nil, _ write: @Sendable @escaping (any ResponseBodyWriter) async throws -> Void) {
-        self.write = { try await write($0); return nil }
+        self._write = { writer in
+            try await write(writer)
+        }
         self.contentLength = contentLength
     }
 
     /// Initialise empty ResponseBody
     public init() {
-        self.init(contentLength: 0) { _ in }
+        self.init(contentLength: 0) { writer in
+            try await writer.finish(nil)
+        }
     }
 
     /// Initialise ResponseBody that contains a single ByteBuffer
     /// - Parameter byteBuffer: ByteBuffer to write
     public init(byteBuffer: ByteBuffer) {
         self.init(contentLength: byteBuffer.readableBytes) { writer in
             try await writer.write(byteBuffer)
+            try await writer.finish(nil)
         }
     }
 
@@ -49,22 +65,13 @@ public struct ResponseBody: Sendable {
             for try await buffer in asyncSequence {
                 try await writer.write(buffer)
             }
-            return
+            try await writer.finish(nil)
         }
     }
 
-    /// Create ResponseBody that returns trailing headers from its closure once all the
-    /// body parts have been written
-    /// - Parameters:
-    ///   - contentLength: Optional length of body
-    ///   - write: closure provided with `writer` type that can be used to write to response body
-    ///         trailing headers are returned from the closure after all the body parts have been
-    ///         written
-    public static func withTrailingHeaders(
-        contentLength: Int? = nil,
-        _ write: @Sendable @escaping (any ResponseBodyWriter) async throws -> HTTPFields?
-    ) -> Self {
-        self.init(contentLength: contentLength, write: write)
+    @inlinable
+    public consuming func write(_ writer: any ResponseBodyWriter) async throws {
+        try await self._write(writer)
     }
 
     /// Returns a ResponseBody containing the results of mapping the given closure over the sequence of
@@ -73,13 +80,12 @@ public struct ResponseBody: Sendable {
     /// - Returns: The transformed ResponseBody
     public consuming func map(_ transform: @escaping @Sendable (ByteBuffer) async throws -> ByteBuffer) -> ResponseBody {
         let body = self
-        return Self.withTrailingHeaders { writer in
-            let tailHeaders = try await body.write(writer.map(transform))
-            return tailHeaders
+        return Self.init { writer in
+            try await body.write(writer.map(transform))
         }
     }
 
-    /// Create new response body that call a callback once original response body has been written
+    /// Create new response body that calls a closure 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
@@ -89,22 +95,13 @@ public struct ResponseBody: Sendable {
     package func withPostWriteClosure(_ postWrite: @escaping @Sendable () async -> Void) -> Self {
         return .init(contentLength: self.contentLength) { writer in
             do {
-                let result = try await self.write(writer)
+                try await self.write(writer)
                 await postWrite()
-                return result
             } catch {
                 await postWrite()
                 throw error
             }
+            return
         }
     }
-
-    /// Initialise ResponseBody with closure writing body contents
-    ///
-    /// This version of init is private and only available via ``withTrailingHeaders`` because
-    /// if it is public the compiler gets confused when a complex closure is provided.
-    private init(contentLength: Int? = nil, write: @Sendable @escaping (any ResponseBodyWriter) async throws -> HTTPFields?) {
-        self.write = { return try await write($0) }
-        self.contentLength = contentLength
-    }
 }

Sources/HummingbirdCore/Response/ResponseBodyWriter.swift

@@ -12,6 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+import HTTPTypes
 import NIOCore
 
 /// HTTP Response Body part writer
@@ -22,6 +23,9 @@ public protocol ResponseBodyWriter {
     /// Write a sequence of ByteBuffers
     /// - Parameter buffers: Sequence of buffers
     func write(contentsOf buffers: some Sequence<ByteBuffer>) async throws
+    /// Finish writing body
+    /// - Parameter trailingHeaders: Any trailing headers you want to include at end
+    consuming func finish(_ trailingHeaders: HTTPFields?) async throws
 }
 
 extension ResponseBodyWriter {
@@ -32,6 +36,15 @@ extension ResponseBodyWriter {
             try await self.write(part)
         }
     }
+
+    ///  Write AsyncSequence of ByteBuffers
+    /// - Parameter buffers: ByteBuffer AsyncSequence
+    @inlinable
+    public func write<BufferSequence: AsyncSequence>(_ buffers: BufferSequence) async throws where BufferSequence.Element == ByteBuffer {
+        for try await buffer in buffers {
+            try await self.write(buffer)
+        }
+    }
 }
 
 struct MappedResponseBodyWriter<ParentWriter: ResponseBodyWriter>: ResponseBodyWriter {
@@ -51,6 +64,12 @@ struct MappedResponseBodyWriter<ParentWriter: ResponseBodyWriter>: ResponseBodyW
             try await self.parentWriter.write(self.transform(part))
         }
     }
+
+    /// Finish writing body
+    /// - Parameter trailingHeaders: Any trailing headers you want to include at end
+    consuming func finish(_ trailingHeaders: HTTPFields?) async throws {
+        try await self.parentWriter.finish(trailingHeaders)
+    }
 }
 
 extension ResponseBodyWriter {

Sources/HummingbirdCore/Response/ResponseWriter.swift

@@ -0,0 +1,88 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Hummingbird server framework project
+//
+// Copyright (c) 2024 the Hummingbird authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See hummingbird/CONTRIBUTORS.txt for the list of Hummingbird authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import HTTPTypes
+import NIOCore
+import NIOHTTPTypes
+
+/// ResponseWriter that writes directly to AsyncChannel
+public struct ResponseWriter: ~Copyable {
+    @usableFromInline
+    let outbound: NIOAsyncChannelOutboundWriter<HTTPResponsePart>
+
+    /// Write HTTP head part and return ``ResponseBodyWriter`` to write response body
+    ///
+    /// - Parameter head: Response head
+    /// - Returns: Response body writer used to write HTTP response body
+    @inlinable
+    public consuming func writeHead(_ head: HTTPResponse) async throws -> some ResponseBodyWriter {
+        try await self.outbound.write(.head(head))
+        return RootResponseBodyWriter(outbound: self.outbound)
+    }
+
+    /// Write Informational HTTP head part
+    ///
+    /// Calling this with a non informational HTTP response head will cause a precondition error
+    /// - Parameter head: Informational response head
+    @inlinable
+    public func writeInformationalHead(_ head: HTTPResponse) async throws {
+        precondition((100..<200).contains(head.status.code), "Informational HTTP responses require a status code in the range of 100 through 199")
+        try await self.outbound.write(.head(head))
+    }
+
+    /// Write full HTTP response that doesn't include a body
+    ///
+    /// - Parameter head: Response head
+    @inlinable
+    public consuming func writeResponse(_ head: HTTPResponse) async throws {
+        try await self.outbound.write(contentsOf: [.head(head), .end(nil)])
+    }
+}
+
+/// ResponseBodyWriter that writes ByteBuffers to AsyncChannel outbound writer
+@usableFromInline
+struct RootResponseBodyWriter: Sendable, ResponseBodyWriter {
+    typealias Out = HTTPResponsePart
+    /// The components of a HTTP response from the view of a HTTP server.
+    public typealias OutboundWriter = NIOAsyncChannelOutboundWriter<Out>
+
+    @usableFromInline
+    let outbound: OutboundWriter
+
+    @usableFromInline
+    init(outbound: OutboundWriter) {
+        self.outbound = outbound
+    }
+
+    /// Write a single ByteBuffer
+    /// - Parameter buffer: single buffer to write
+    @inlinable
+    func write(_ buffer: ByteBuffer) async throws {
+        try await self.outbound.write(.body(buffer))
+    }
+
+    /// Write a sequence of ByteBuffers
+    /// - Parameter buffers: Sequence of buffers
+    @inlinable
+    func write(contentsOf buffers: some Sequence<ByteBuffer>) async throws {
+        try await self.outbound.write(contentsOf: buffers.map { .body($0) })
+    }
+
+    /// Finish writing body
+    /// - Parameter trailingHeaders: Any trailing headers you want to include at end
+    @inlinable
+    consuming func finish(_ trailingHeaders: HTTPFields?) async throws {
+        try await self.outbound.write(.end(trailingHeaders))
+    }
+}

Sources/HummingbirdCore/Server/HTTP/HTTPChannelHandler.swift

@@ -21,7 +21,7 @@ import ServiceLifecycle
 
 /// Protocol for HTTP channels
 public protocol HTTPChannelHandler: ServerChildChannel {
-    typealias Responder = @Sendable (Request, Channel) async -> Response
+    typealias Responder = @Sendable (Request, consuming ResponseWriter, Channel) async throws -> Void
     var responder: Responder { get }
 }
 
@@ -37,7 +37,6 @@ extension HTTPChannelHandler {
         do {
             try await withTaskCancellationHandler {
                 try await asyncChannel.executeThenClose { inbound, outbound in
-                    let responseWriter = HTTPServerBodyWriter(outbound: outbound)
                     var iterator = inbound.makeAsyncIterator()
 
                     // read first part, verify it is a head
@@ -49,11 +48,9 @@ extension HTTPChannelHandler {
                     while true {
                         let bodyStream = NIOAsyncChannelRequestBody(iterator: iterator)
                         let request = Request(head: head, body: .init(asyncSequence: bodyStream))
-                        let response = await self.responder(request, asyncChannel.channel)
+                        let responseWriter = ResponseWriter(outbound: outbound)
                         do {
-                            try await outbound.write(.head(response.head))
-                            let tailHeaders = try await response.body.write(responseWriter)
-                            try await outbound.write(.end(tailHeaders))
+                            try await self.responder(request, responseWriter, asyncChannel.channel)
                         } catch {
                             throw error
                         }
@@ -92,24 +89,3 @@ extension HTTPChannelHandler {
         }
     }
 }
-
-/// ResponseBodyWriter that writes ByteBuffers to AsyncChannel outbound writer
-struct HTTPServerBodyWriter: Sendable, ResponseBodyWriter {
-    typealias Out = HTTPResponsePart
-    /// The components of a HTTP response from the view of a HTTP server.
-    public typealias OutboundWriter = NIOAsyncChannelOutboundWriter<Out>
-
-    let outbound: OutboundWriter
-
-    /// Write a single ByteBuffer
-    /// - Parameter buffer: single buffer to write
-    func write(_ buffer: ByteBuffer) async throws {
-        try await self.outbound.write(.body(buffer))
-    }
-
-    /// Write a sequence of ByteBuffers
-    /// - Parameter buffers: Sequence of buffers
-    func write(contentsOf buffers: some Sequence<ByteBuffer>) async throws {
-        try await self.outbound.write(contentsOf: buffers.map { .body($0) })
-    }
-}

Sources/HummingbirdTesting/RouterTestFramework.swift

@@ -119,9 +119,9 @@ struct RouterTestFramework<Responder: HTTPResponder>: ApplicationTestFramework w
                         response = Response(status: .internalServerError)
                     }
                     let responseWriter = RouterResponseWriter()
-                    let trailerHeaders = try await response.body.write(responseWriter)
-                    return responseWriter.collated.withLockedValue { collated in
-                        TestResponse(head: response.head, body: collated, trailerHeaders: trailerHeaders)
+                    try await response.body.write(responseWriter)
+                    return responseWriter.values.withLockedValue { values in
+                        TestResponse(head: response.head, body: values.body, trailerHeaders: values.trailingHeaders)
                     }
                 }
 
@@ -141,15 +141,21 @@ struct RouterTestFramework<Responder: HTTPResponder>: ApplicationTestFramework w
     }
 
     final class RouterResponseWriter: ResponseBodyWriter {
-        let collated: NIOLockedValueBox<ByteBuffer>
+        let values: NIOLockedValueBox<(body: ByteBuffer, trailingHeaders: HTTPFields?)>
 
         init() {
-            self.collated = .init(.init())
+            self.values = .init((body: .init(), trailingHeaders: nil))
         }
 
         func write(_ buffer: ByteBuffer) async throws {
-            _ = self.collated.withLockedValue { collated in
-                collated.writeImmutableBuffer(buffer)
+            _ = self.values.withLockedValue { values in
+                values.body.writeImmutableBuffer(buffer)
+            }
+        }
+
+        func finish(_ headers: HTTPTypes.HTTPFields?) async throws {
+            self.values.withLockedValue { values in
+                values.trailingHeaders = headers
             }
         }
     }