fix 784: dont crash on huge in-memory bodies

weissi · weissi · commit 7e55dbec486c · 2024-11-25T16:06:48.000Z
diff --git a/Sources/AsyncHTTPClient/AsyncAwait/Transaction+StateMachine.swift b/Sources/AsyncHTTPClient/AsyncAwait/Transaction+StateMachine.swift
@@ -34,14 +34,14 @@ extension Transaction {
             case finished(error: Error?)
         }
 
-        fileprivate enum RequestStreamState {
+        fileprivate enum RequestStreamState: Sendable {
             case requestHeadSent
             case producing
             case paused(continuation: CheckedContinuation<Void, Error>?)
             case finished
         }
 
-        fileprivate enum ResponseStreamState {
+        fileprivate enum ResponseStreamState: Sendable {
             // Waiting for response head. Valid transitions to: streamingBody.
             case waitingForResponseHead
             // streaming response body. Valid transitions to: finished.
diff --git a/Sources/AsyncHTTPClient/AsyncAwait/Transaction.swift b/Sources/AsyncHTTPClient/AsyncAwait/Transaction.swift
@@ -19,7 +19,9 @@ import NIOHTTP1
 import NIOSSL
 
 @available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
-@usableFromInline final class Transaction: @unchecked Sendable {
+@usableFromInline final class Transaction:
+    // until NIOLockedValueBox learns `sending` because StateMachine cannot be Sendable
+    @unchecked Sendable {
     let logger: Logger
 
     let request: HTTPClientRequest.Prepared
@@ -28,8 +30,7 @@ import NIOSSL
     let preferredEventLoop: EventLoop
     let requestOptions: RequestOptions
 
-    private let stateLock = NIOLock()
-    private var state: StateMachine
+    private let state: NIOLockedValueBox<StateMachine>
 
     init(
         request: HTTPClientRequest.Prepared,
@@ -44,7 +45,7 @@ import NIOSSL
         self.logger = logger
         self.connectionDeadline = connectionDeadline
         self.preferredEventLoop = preferredEventLoop
-        self.state = StateMachine(responseContinuation)
+        self.state = NIOLockedValueBox(StateMachine(responseContinuation))
     }
 
     func cancel() {
@@ -56,8 +57,8 @@ import NIOSSL
     private func writeOnceAndOneTimeOnly(byteBuffer: ByteBuffer) {
         // This method is synchronously invoked after sending the request head. For this reason we
         // can make a number of assumptions, how the state machine will react.
-        let writeAction = self.stateLock.withLock {
-            self.state.writeNextRequestPart()
+        let writeAction = self.state.withLockedValue { state in
+            state.writeNextRequestPart()
         }
 
         switch writeAction {
@@ -99,30 +100,33 @@ import NIOSSL
 
     struct BreakTheWriteLoopError: Swift.Error {}
 
+    // FIXME: Refactor this to not use `self.state.unsafe`.
     private func writeRequestBodyPart(_ part: ByteBuffer) async throws {
-        self.stateLock.lock()
-        switch self.state.writeNextRequestPart() {
+        self.state.unsafe.lock()
+        switch self.state.unsafe.withValueAssumingLockIsAcquired({ state in state.writeNextRequestPart() }) {
         case .writeAndContinue(let executor):
-            self.stateLock.unlock()
+            self.state.unsafe.unlock()
             executor.writeRequestBodyPart(.byteBuffer(part), request: self, promise: nil)
 
         case .writeAndWait(let executor):
             try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, Error>) in
-                self.state.waitForRequestBodyDemand(continuation: continuation)
-                self.stateLock.unlock()
+                self.state.unsafe.withValueAssumingLockIsAcquired({ state in
+                    state.waitForRequestBodyDemand(continuation: continuation)
+                })
+                self.state.unsafe.unlock()
 
                 executor.writeRequestBodyPart(.byteBuffer(part), request: self, promise: nil)
             }
 
         case .fail:
-            self.stateLock.unlock()
+            self.state.unsafe.unlock()
             throw BreakTheWriteLoopError()
         }
     }
 
     private func requestBodyStreamFinished() {
-        let finishAction = self.stateLock.withLock {
-            self.state.finishRequestBodyStream()
+        let finishAction = self.state.withLockedValue { state in
+            state.finishRequestBodyStream()
         }
 
         switch finishAction {
@@ -150,8 +154,8 @@ extension Transaction: HTTPSchedulableRequest {
     var requiredEventLoop: EventLoop? { nil }
 
     func requestWasQueued(_ scheduler: HTTPRequestScheduler) {
-        self.stateLock.withLock {
-            self.state.requestWasQueued(scheduler)
+        self.state.withLockedValue { state in
+            state.requestWasQueued(scheduler)
         }
     }
 }
@@ -165,8 +169,8 @@ extension Transaction: HTTPExecutableRequest {
     // MARK: Request
 
     func willExecuteRequest(_ executor: HTTPRequestExecutor) {
-        let action = self.stateLock.withLock {
-            self.state.willExecuteRequest(executor)
+        let action = self.state.withLockedValue { state in
+            state.willExecuteRequest(executor)
         }
 
         switch action {
@@ -183,8 +187,8 @@ extension Transaction: HTTPExecutableRequest {
     func requestHeadSent() {}
 
     func resumeRequestBodyStream() {
-        let action = self.stateLock.withLock {
-            self.state.resumeRequestBodyStream()
+        let action = self.state.withLockedValue { state in
+            state.resumeRequestBodyStream()
         }
 
         switch action {
@@ -214,16 +218,16 @@ extension Transaction: HTTPExecutableRequest {
     }
 
     func pauseRequestBodyStream() {
-        self.stateLock.withLock {
-            self.state.pauseRequestBodyStream()
+        self.state.withLockedValue { state in
+            state.pauseRequestBodyStream()
         }
     }
 
     // MARK: Response
 
     func receiveResponseHead(_ head: HTTPResponseHead) {
-        let action = self.stateLock.withLock {
-            self.state.receiveResponseHead(head, delegate: self)
+        let action = self.state.withLockedValue { state in
+            state.receiveResponseHead(head, delegate: self)
         }
 
         switch action {
@@ -243,8 +247,8 @@ extension Transaction: HTTPExecutableRequest {
     }
 
     func receiveResponseBodyParts(_ buffer: CircularBuffer<ByteBuffer>) {
-        let action = self.stateLock.withLock {
-            self.state.receiveResponseBodyParts(buffer)
+        let action = self.state.withLockedValue { state in
+            state.receiveResponseBodyParts(buffer)
         }
         switch action {
         case .none:
@@ -260,8 +264,8 @@ extension Transaction: HTTPExecutableRequest {
     }
 
     func succeedRequest(_ buffer: CircularBuffer<ByteBuffer>?) {
-        let succeedAction = self.stateLock.withLock {
-            self.state.succeedRequest(buffer)
+        let succeedAction = self.state.withLockedValue { state in
+            state.succeedRequest(buffer)
         }
         switch succeedAction {
         case .finishResponseStream(let source, let finalResponse):
@@ -276,8 +280,8 @@ extension Transaction: HTTPExecutableRequest {
     }
 
     func fail(_ error: Error) {
-        let action = self.stateLock.withLock {
-            self.state.fail(error)
+        let action = self.state.withLockedValue { state in
+            state.fail(error)
         }
         self.performFailAction(action)
     }
@@ -304,8 +308,8 @@ extension Transaction: HTTPExecutableRequest {
     }
 
     func deadlineExceeded() {
-        let action = self.stateLock.withLock {
-            self.state.deadlineExceeded()
+        let action = self.state.withLockedValue { state in
+            state.deadlineExceeded()
         }
         self.performDeadlineExceededAction(action)
     }
@@ -329,8 +333,8 @@ extension Transaction: HTTPExecutableRequest {
 extension Transaction: NIOAsyncSequenceProducerDelegate {
     @usableFromInline
     func produceMore() {
-        let action = self.stateLock.withLock {
-            self.state.produceMore()
+        let action = self.state.withLockedValue { state in
+            state.produceMore()
         }
         switch action {
         case .none:
diff --git a/Sources/AsyncHTTPClient/ConnectionPool/HTTP2/HTTP2ClientRequestHandler.swift b/Sources/AsyncHTTPClient/ConnectionPool/HTTP2/HTTP2ClientRequestHandler.swift
@@ -432,6 +432,9 @@ final class HTTP2ClientRequestHandler: ChannelDuplexHandler {
     }
 }
 
+@available(*, unavailable)
+extension HTTP2ClientRequestHandler: Sendable {}
+
 extension HTTP2ClientRequestHandler: HTTPRequestExecutor {
     func writeRequestBodyPart(_ data: IOData, request: HTTPExecutableRequest, promise: EventLoopPromise<Void>?) {
         if self.eventLoop.inEventLoop {
diff --git a/Sources/AsyncHTTPClient/ConnectionPool/HTTPExecutableRequest.swift b/Sources/AsyncHTTPClient/ConnectionPool/HTTPExecutableRequest.swift
@@ -132,7 +132,7 @@ import NIOSSL
 ///
 /// Use this handle to cancel the request, while it is waiting for a free connection, to execute the request.
 /// This protocol is only intended to be implemented by the `HTTPConnectionPool`.
-protocol HTTPRequestScheduler {
+protocol HTTPRequestScheduler: Sendable {
     /// Informs the task queuer that a request has been cancelled.
     func cancelRequest(_: HTTPSchedulableRequest)
 }
diff --git a/Sources/AsyncHTTPClient/HTTPClient.swift b/Sources/AsyncHTTPClient/HTTPClient.swift
@@ -222,21 +222,20 @@ public class HTTPClient {
                 """
             )
         }
-        let errorStorageLock = NIOLock()
-        let errorStorage: UnsafeMutableTransferBox<Error?> = .init(nil)
+        let errorStorage: NIOLockedValueBox<Error?> = NIOLockedValueBox(nil)
         let continuation = DispatchWorkItem {}
         self.shutdown(requiresCleanClose: requiresCleanClose, queue: DispatchQueue(label: "async-http-client.shutdown"))
         { error in
             if let error = error {
-                errorStorageLock.withLock {
-                    errorStorage.wrappedValue = error
+                errorStorage.withLockedValue { errorStorage in
+                    errorStorage = error
                 }
             }
             continuation.perform()
         }
         continuation.wait()
-        try errorStorageLock.withLock {
-            if let error = errorStorage.wrappedValue {
+        try errorStorage.withLockedValue { errorStorage in
+            if let error = errorStorage {
                 throw error
             }
         }
diff --git a/Sources/AsyncHTTPClient/HTTPHandler.swift b/Sources/AsyncHTTPClient/HTTPHandler.swift
@@ -47,25 +47,65 @@ extension HTTPClient {
             }
 
             @inlinable
-            func writeChunks<Bytes: Collection>(of bytes: Bytes, maxChunkSize: Int) -> EventLoopFuture<Void>
-            where Bytes.Element == UInt8 {
-                let iterator = UnsafeMutableTransferBox(bytes.chunks(ofCount: maxChunkSize).makeIterator())
-                guard let chunk = iterator.wrappedValue.next() else {
+            func writeChunks<Bytes: Collection>(
+                of bytes: Bytes,
+                maxChunkSize: Int
+            ) -> EventLoopFuture<Void> where Bytes.Element == UInt8 {
+                // `StreamWriter` is has design issues, for example
+                // - https://github.com/swift-server/async-http-client/issues/194
+                // - https://github.com/swift-server/async-http-client/issues/264
+                // - We're not told the EventLoop the task runs on and the user is free to return whatever EL they
+                //   want.
+                // One important consideration then is that we must lock around the iterator because we could be hopping
+                // between threads.
+                typealias Iterator = EnumeratedSequence<ChunksOfCountCollection<Bytes>>.Iterator
+                typealias Chunk = (offset: Int, element: ChunksOfCountCollection<Bytes>.Element)
+
+                func makeIteratorAndFirstChunk() -> (
+                    iterator: NIOLockedValueBox<Iterator>,
+                    chunk: Chunk
+                )? {
+                    var iterator = bytes.chunks(ofCount: maxChunkSize).enumerated().makeIterator()
+                    guard let chunk = iterator.next() else {
+                        return nil
+                    }
+
+                    return (NIOLockedValueBox(iterator), chunk)
+                }
+
+                guard let (iterator, chunk) = makeIteratorAndFirstChunk() else {
                     return self.write(IOData.byteBuffer(.init()))
                 }
 
                 @Sendable  // can't use closure here as we recursively call ourselves which closures can't do
-                func writeNextChunk(_ chunk: Bytes.SubSequence) -> EventLoopFuture<Void> {
-                    if let nextChunk = iterator.wrappedValue.next() {
-                        return self.write(.byteBuffer(ByteBuffer(bytes: chunk))).flatMap {
-                            writeNextChunk(nextChunk)
-                        }
+                func writeNextChunk(_ chunk: Chunk, allDone: EventLoopPromise<Void>) {
+                    if let nextElement = iterator.withLockedValue({ $0.next() }) {
+                        let (index, chunk) = nextElement
+                        self.write(.byteBuffer(ByteBuffer(bytes: chunk))).map {
+                            if (index + 1) % 4 == 0 {
+                                // Let's not stack-overflow if the futures insta-complete which they at least in HTTP/2
+                                // mode.
+                                // Also, we must frequently return to the EventLoop because we may get the pause signal
+                                // from another thread. If we fail to do that promptly, we may balloon our body chunks
+                                // into memory.
+                                allDone.futureResult.eventLoop.execute {
+                                    writeNextChunk(nextElement, allDone: allDone)
+                                }
+                            } else {
+                                writeNextChunk(nextElement, allDone: allDone)
+                            }
+                        }.cascadeFailure(to: allDone)
                     } else {
-                        return self.write(.byteBuffer(ByteBuffer(bytes: chunk)))
+                        self.write(.byteBuffer(ByteBuffer(bytes: chunk.element))).cascade(to: allDone)
                     }
                 }
 
-                return writeNextChunk(chunk)
+                // HACK (again, we're not told the right EventLoop): Let's write 0 bytes to make the user tell us...
+                return self.write(.byteBuffer(ByteBuffer())).flatMapWithEventLoop { (_, loop) in
+                    let allDone = loop.makePromise(of: Void.self)
+                    writeNextChunk(chunk, allDone: allDone)
+                    return allDone.futureResult
+                }
             }
         }
 
diff --git a/Sources/AsyncHTTPClient/UnsafeTransfer.swift b/Sources/AsyncHTTPClient/UnsafeTransfer.swift
diff --git a/Tests/AsyncHTTPClientTests/HTTP2ClientTests.swift b/Tests/AsyncHTTPClientTests/HTTP2ClientTests.swift
diff --git a/Tests/AsyncHTTPClientTests/HTTPClientTestUtils.swift b/Tests/AsyncHTTPClientTests/HTTPClientTestUtils.swift

Original file line number	Diff line number	Diff line change
`@@ -432,6 +432,9 @@ final class HTTP2ClientRequestHandler: ChannelDuplexHandler {`
`432`	`432`	`}`
`433`	`433`	`}`
`434`	`434`
	`435`	`+@available(*, unavailable)`
	`436`	`+extension HTTP2ClientRequestHandler: Sendable {}`
	`437`	`+`
`435`	`438`	`extension HTTP2ClientRequestHandler: HTTPRequestExecutor {`
`436`	`439`	`func writeRequestBodyPart(_ data: IOData, request: HTTPExecutableRequest, promise: EventLoopPromise<Void>?) {`
`437`	`440`	`if self.eventLoop.inEventLoop {`
Original file line number	Diff line number	Diff line change
`@@ -132,7 +132,7 @@ import NIOSSL`
`132`	`132`	`///`
`133`	`133`	`/// Use this handle to cancel the request, while it is waiting for a free connection, to execute the request.`
`134`	`134`	/// This protocol is only intended to be implemented by the `HTTPConnectionPool`.
`135`		`-protocol HTTPRequestScheduler {`
	`135`	`+protocol HTTPRequestScheduler: Sendable {`
`136`	`136`	`/// Informs the task queuer that a request has been cancelled.`
`137`	`137`	`func cancelRequest(_: HTTPSchedulableRequest)`
`138`	`138`	`}`
Original file line number	Diff line number	Diff line change
`@@ -222,21 +222,20 @@ public class HTTPClient {`
`222`	`222`	`"""`
`223`	`223`	`)`
`224`	`224`	`}`
`225`		`- let errorStorageLock = NIOLock()`
`226`		`- let errorStorage: UnsafeMutableTransferBox<Error?> = .init(nil)`
	`225`	`+ let errorStorage: NIOLockedValueBox<Error?> = NIOLockedValueBox(nil)`
`227`	`226`	`let continuation = DispatchWorkItem {}`
`228`	`227`	`self.shutdown(requiresCleanClose: requiresCleanClose, queue: DispatchQueue(label: "async-http-client.shutdown"))`
`229`	`228`	`{ error in`
`230`	`229`	`if let error = error {`
`231`		`- errorStorageLock.withLock {`
`232`		`- errorStorage.wrappedValue = error`
	`230`	`+ errorStorage.withLockedValue { errorStorage in`
	`231`	`+ errorStorage = error`
`233`	`232`	`}`
`234`	`233`	`}`
`235`	`234`	`continuation.perform()`
`236`	`235`	`}`
`237`	`236`	`continuation.wait()`
`238`		`- try errorStorageLock.withLock {`
`239`		`- if let error = errorStorage.wrappedValue {`
	`237`	`+ try errorStorage.withLockedValue { errorStorage in`
	`238`	`+ if let error = errorStorage {`
`240`	`239`	`throw error`
`241`	`240`	`}`
`242`	`241`	`}`