Sendable Response (#3082)

* Make Response Sendable * Fix a warning * Add preconcurrency annotations * Remove a lock * Bump ConsoleKit to Sendable version * Remove preconcurrency imports for ConsoleKit * Move all the response stuff into it's own box * Fix up some usages * Bypass multiple lock accesses for internal types * Bypass multiple lock accesses for internal types * Set the content length correctly * Update last use of response body * Migrate double lock access to single lock * Update Sources/Vapor/Response/Response.swift Co-authored-by: Mahdi Bahrami <github@mahdibm.com> --------- Co-authored-by: Mahdi Bahrami <github@mahdibm.com>
vapor · Nov 1, 2023 · 3bf4e73 · 3bf4e73
1 parent 600d666
commit 3bf4e73
Show file tree

Hide file tree

Showing 17 changed files with 196 additions and 99 deletions.
diff --git a/Package.swift b/Package.swift
@@ -22,7 +22,7 @@ let package = Package(
         .package(url: "https://github.com/vapor/async-kit.git", from: "1.15.0"),
 
         // 💻 APIs for creating interactive CLI tools.
-        .package(url: "https://github.com/vapor/console-kit.git", from: "4.0.0"),
+        .package(url: "https://github.com/vapor/console-kit.git", from: "4.9.0"),
 
         // 🔑 Hashing (SHA2, HMAC), encryption (AES), public-key (RSA), and random data generation.
         .package(url: "https://github.com/apple/swift-crypto.git", "1.0.0" ..< "4.0.0"),

diff --git a/Package@swift-5.9.swift b/Package@swift-5.9.swift
@@ -22,7 +22,7 @@ let package = Package(
         .package(url: "https://github.com/vapor/async-kit.git", from: "1.15.0"),
 
         // 💻 APIs for creating interactive CLI tools.
-        .package(url: "https://github.com/vapor/console-kit.git", from: "4.0.0"),
+        .package(url: "https://github.com/vapor/console-kit.git", from: "4.9.0"),
 
         // 🔑 Hashing (SHA2, HMAC), encryption (AES), public-key (RSA), and random data generation.
         .package(url: "https://github.com/apple/swift-crypto.git", "1.0.0" ..< "4.0.0"),

diff --git a/Sources/Vapor/Commands/ServeCommand.swift b/Sources/Vapor/Commands/ServeCommand.swift
@@ -1,5 +1,6 @@
 import Foundation
-@preconcurrency import ConsoleKit
+@preconcurrency import Dispatch
+import ConsoleKit
 import NIOConcurrencyHelpers
 
 /// Boots the application's server. Listens for `SIGINT` and `SIGTERM` for graceful shutdown.

diff --git a/Sources/Vapor/Concurrency/ResponseCodable+Concurrency.swift b/Sources/Vapor/Concurrency/ResponseCodable+Concurrency.swift
@@ -52,10 +52,12 @@ extension AsyncResponseEncodable {
     /// - returns: Newly encoded `Response`.
     public func encodeResponse(status: HTTPStatus, headers: HTTPHeaders = [:], for request: Request) async throws -> Response {
         let response = try await self.encodeResponse(for: request)
-        for (name, value) in headers {
-            response.headers.replaceOrAdd(name: name, value: value)
+        response.responseBox.withLockedValue { box in
+            for (name, value) in headers {
+                box.headers.replaceOrAdd(name: name, value: value)
+            }
+            box.status = status
         }
-        response.status = status
         return response
     }
 }

diff --git a/Sources/Vapor/Concurrency/ViewRenderer+Concurrency.swift b/Sources/Vapor/Concurrency/ViewRenderer+Concurrency.swift
@@ -13,8 +13,10 @@ public extension ViewRenderer {
 extension View: AsyncResponseEncodable {
     public func encodeResponse(for request: Request) async throws -> Response {
         let response = Response()
-        response.headers.contentType = .html
-        response.body = .init(buffer: self.data)
+        response.responseBox.withLockedValue { box in
+            box.headers.contentType = .html
+            box.body = .init(buffer: self.data)
+        }
         return response
     }
 }
diff --git a/Sources/Vapor/Core/Core.swift b/Sources/Vapor/Core/Core.swift
@@ -1,4 +1,4 @@
-@preconcurrency import ConsoleKit
+import ConsoleKit
 import NIOCore
 import NIOPosix
 import NIOConcurrencyHelpers

diff --git a/Sources/Vapor/HTTP/Server/HTTPServerHandler.swift b/Sources/Vapor/HTTP/Server/HTTPServerHandler.swift
@@ -31,7 +31,7 @@ final class HTTPServerHandler: ChannelInboundHandler, RemovableChannelHandler {
             self.errorCaught(context: context, error: error)
         case .success(let response):
             if request.method == .HEAD {
-                response.forHeadRequest = true
+                response.responseBox.withLockedValue { $0.forHeadRequest = true }
             }
             self.serialize(response, for: request, context: context)
         }

diff --git a/Sources/Vapor/HTTP/Server/HTTPServerResponseEncoder.swift b/Sources/Vapor/HTTP/Server/HTTPServerResponseEncoder.swift
@@ -20,23 +20,25 @@ final class HTTPServerResponseEncoder: ChannelOutboundHandler, RemovableChannelH
 
     func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
         let response = self.unwrapOutboundIn(data)
-        // add a RFC1123 timestamp to the Date header to make this
-        // a valid request
-        response.headers.add(name: "date", value: self.dateCache.currentTimestamp())
-
-        if let server = self.serverHeader {
-            response.headers.add(name: "server", value: server)
+        var headOrNoContentRequest = false
+        response.responseBox.withLockedValue { box in
+            // add a RFC1123 timestamp to the Date header to make this
+            // a valid request
+            box.headers.add(name: "date", value: self.dateCache.currentTimestamp())
+            if let server = self.serverHeader {
+                box.headers.add(name: "server", value: server)
+            }
+
+            // begin serializing
+            let responseHead = HTTPResponseHead(version: box.version, status: box.status, headers: box.headers)
+            context.write(wrapOutboundOut(.head(responseHead)), promise: nil)
+
+            if box.status == .noContent || box.forHeadRequest {
+                headOrNoContentRequest = true
+            }
         }
 
-        // begin serializing
-        context.write(wrapOutboundOut(.head(.init(
-            version: response.version,
-            status: response.status,
-            headers: response.headers
-        ))), promise: nil)
-
-
-        if response.status == .noContent || response.forHeadRequest {
+        if headOrNoContentRequest {
             // don't send bodies for 204 (no content) responses
             // or HEAD requests
             context.fireUserInboundEventTriggered(ResponseEndSentEvent())

diff --git a/Sources/Vapor/HTTP/Server/HTTPServerUpgradeHandler.swift b/Sources/Vapor/HTTP/Server/HTTPServerUpgradeHandler.swift
@@ -57,7 +57,10 @@ final class HTTPServerUpgradeHandler: ChannelDuplexHandler, RemovableChannelHand
         switch self.upgradeState {
         case .pending(let req, let buffer):
             self.upgradeState = .upgraded
-            if res.status == .switchingProtocols, let upgrader = res.upgrader {
+            let (status, upgrader) = res.responseBox.withLockedValue { box in
+                return (box.status, box.upgrader)
+            }
+            if status == .switchingProtocols, let upgrader = upgrader {
                 let protocolUpgrader = upgrader.applyUpgrade(req: req, res: res)
                 let sendableBox = SendableBox(
                     context: context,
@@ -128,7 +131,8 @@ private final class UpgradeBufferHandler: ChannelInboundHandler, RemovableChanne
 }
 
 /// Conformance for any struct that performs an HTTP Upgrade
-public protocol Upgrader {
+@preconcurrency
+public protocol Upgrader: Sendable {
     func applyUpgrade(req: Request, res: Response) -> HTTPServerProtocolUpgrader
 }
 

diff --git a/Sources/Vapor/Middleware/CORSMiddleware.swift b/Sources/Vapor/Middleware/CORSMiddleware.swift
@@ -138,26 +138,27 @@ public final class CORSMiddleware: Middleware {
         return response.map { response in
             // Modify response headers based on CORS settings
             let originBasedAccessControlAllowHeader = self.configuration.allowedOrigin.header(forRequest: request)
-            response.headers.replaceOrAdd(name: .accessControlAllowOrigin, value: originBasedAccessControlAllowHeader)
-            response.headers.replaceOrAdd(name: .accessControlAllowHeaders, value: self.configuration.allowedHeaders)
-            response.headers.replaceOrAdd(name: .accessControlAllowMethods, value: self.configuration.allowedMethods)
-
-            if let exposedHeaders = self.configuration.exposedHeaders {
-                response.headers.replaceOrAdd(name: .accessControlExpose, value: exposedHeaders)
-            }
-
-            if let cacheExpiration = self.configuration.cacheExpiration {
-                response.headers.replaceOrAdd(name: .accessControlMaxAge, value: String(cacheExpiration))
-            }
-
-            if self.configuration.allowCredentials {
-                response.headers.replaceOrAdd(name: .accessControlAllowCredentials, value: "true")
-            }
+            response.responseBox.withLockedValue { box in
+                box.headers.replaceOrAdd(name: .accessControlAllowOrigin, value: originBasedAccessControlAllowHeader)
+                box.headers.replaceOrAdd(name: .accessControlAllowHeaders, value: self.configuration.allowedHeaders)
+                box.headers.replaceOrAdd(name: .accessControlAllowMethods, value: self.configuration.allowedMethods)
+
+                if let exposedHeaders = self.configuration.exposedHeaders {
+                    box.headers.replaceOrAdd(name: .accessControlExpose, value: exposedHeaders)
+                }
+
+                if let cacheExpiration = self.configuration.cacheExpiration {
+                    box.headers.replaceOrAdd(name: .accessControlMaxAge, value: String(cacheExpiration))
+                }
+
+                if self.configuration.allowCredentials {
+                    box.headers.replaceOrAdd(name: .accessControlAllowCredentials, value: "true")
+                }
 
-            if case .originBased = self.configuration.allowedOrigin, !originBasedAccessControlAllowHeader.isEmpty {
-                response.headers.add(name: .vary, value: "origin")
+                if case .originBased = self.configuration.allowedOrigin, !originBasedAccessControlAllowHeader.isEmpty {
+                    box.headers.add(name: .vary, value: "origin")
+                }
             }
-
             return response
         }
     }

diff --git a/Sources/Vapor/Middleware/ErrorMiddleware.swift b/Sources/Vapor/Middleware/ErrorMiddleware.swift
@@ -51,11 +51,15 @@ public final class ErrorMiddleware: Middleware {
             // attempt to serialize the error to json
             do {
                 let errorResponse = ErrorResponse(error: true, reason: reason)
-                response.body = try .init(data: JSONEncoder().encode(errorResponse), byteBufferAllocator: req.byteBufferAllocator)
-                response.headers.replaceOrAdd(name: .contentType, value: "application/json; charset=utf-8")
+                try response.responseBox.withLockedValue { box in
+                    box.body = try .init(data: JSONEncoder().encode(errorResponse), byteBufferAllocator: req.byteBufferAllocator)
+                    box.headers.replaceOrAdd(name: .contentType, value: "application/json; charset=utf-8")
+                }
             } catch {
-                response.body = .init(string: "Oops: \(error)", byteBufferAllocator: req.byteBufferAllocator)
-                response.headers.replaceOrAdd(name: .contentType, value: "text/plain; charset=utf-8")
+                response.responseBox.withLockedValue { box in
+                    box.body = .init(string: "Oops: \(error)", byteBufferAllocator: req.byteBufferAllocator)
+                    box.headers.replaceOrAdd(name: .contentType, value: "text/plain; charset=utf-8")
+                }
             }
             return response
         }

diff --git a/Sources/Vapor/Request/Redirect.swift b/Sources/Vapor/Request/Redirect.swift
@@ -14,8 +14,10 @@ extension Request {
     @available(*, deprecated, renamed: "redirect(to:redirectType:)")
     public func redirect(to location: String, type: RedirectType) -> Response {
         let response = Response()
-        response.status = type.status
-        response.headers.replaceOrAdd(name: .location, value: location)
+        response.responseBox.withLockedValue { box in
+            box.status = type.status
+            box.headers.replaceOrAdd(name: .location, value: location)
+        }
         return response
     }
 
@@ -33,8 +35,10 @@ extension Request {
     /// - Returns: A response that redirects the client to the specified location
     public func redirect(to location: String, redirectType: Redirect = .normal) -> Response {
         let response = Response()
-        response.status = redirectType.status
-        response.headers.replaceOrAdd(name: .location, value: location)
+        response.responseBox.withLockedValue { box in
+            box.status = redirectType.status
+            box.headers.replaceOrAdd(name: .location, value: location)
+        }
         return response
     }
 }

diff --git a/Sources/Vapor/Response/Response+Body.swift b/Sources/Vapor/Response/Response+Body.swift
@@ -1,21 +1,22 @@
+@preconcurrency import Dispatch
 import Foundation
 import NIOCore
 import NIOConcurrencyHelpers
 
 extension Response {
-    struct BodyStream {
+    struct BodyStream: Sendable {
         let count: Int
-        let callback: (BodyStreamWriter) -> ()
+        let callback: @Sendable (BodyStreamWriter) -> ()
     }
 
     /// Represents a `Response`'s body.
     ///
     ///     let body = Response.Body(string: "Hello, world!")
     ///
     /// This can contain any data (streaming or static) and should match the message's `"Content-Type"` header.
-    public struct Body: CustomStringConvertible, ExpressibleByStringLiteral {
+    public struct Body: CustomStringConvertible, ExpressibleByStringLiteral, Sendable {
         /// The internal HTTP body storage enum. This is an implementation detail.
-        internal enum Storage {
+        internal enum Storage: Sendable {
             /// Cases
             case none
             case buffer(ByteBuffer)
@@ -151,12 +152,14 @@ extension Response {
             self.storage = .buffer(buffer)
         }
 
-        public init(stream: @escaping (BodyStreamWriter) -> (), count: Int, byteBufferAllocator: ByteBufferAllocator = ByteBufferAllocator()) {
+        @preconcurrency
+        public init(stream: @Sendable @escaping (BodyStreamWriter) -> (), count: Int, byteBufferAllocator: ByteBufferAllocator = ByteBufferAllocator()) {
             self.byteBufferAllocator = byteBufferAllocator
             self.storage = .stream(.init(count: count, callback: stream))
         }
 
-        public init(stream: @escaping (BodyStreamWriter) -> (), byteBufferAllocator: ByteBufferAllocator = ByteBufferAllocator()) {
+        @preconcurrency
+        public init(stream: @Sendable @escaping (BodyStreamWriter) -> (), byteBufferAllocator: ByteBufferAllocator = ByteBufferAllocator()) {
             self.init(stream: stream, count: -1, byteBufferAllocator: byteBufferAllocator)
         }