Return iterator to state in deinit

Author: gjcairo

Sources/NIOHTTPServer/HTTPRequestConcludingAsyncReader.swift

@@ -38,7 +38,7 @@ public struct HTTPRequestConcludingAsyncReader: ConcludingAsyncReader, ~Copyable
         /// The type of errors that can occur during reading operations.
         public typealias ReadFailure = any Error
 
-        /// The shared reader state that holds the iterator and captures trailers.
+        /// The shared reader state for iterator recovery and trailer capture.
         fileprivate var state: ReaderState
 
         struct RequestBodyStateMachine {
@@ -69,10 +69,17 @@ public struct HTTPRequestConcludingAsyncReader: ConcludingAsyncReader, ~Copyable
 
             enum ReadResult {
                 case readBody(ByteBuffer)
+                /// The request end was received. The caller should return the iterator
+                /// to ReaderState and nil it out.
                 case requestFinished
+                /// The request was already finished in a previous read.
+                case alreadyFinished
             }
 
-            mutating func read(limit: Int?) async throws -> ReadResult {
+            mutating func read(
+                iterator: inout NIOAsyncChannelInboundStream<HTTPRequestPart>.AsyncIterator?,
+                limit: Int?
+            ) async throws -> ReadResult {
                 switch self.state {
                 case .readingBody(let readingBodyState):
                     var bodyElement: ByteBuffer
@@ -85,42 +92,20 @@ public struct HTTPRequestConcludingAsyncReader: ConcludingAsyncReader, ~Copyable
 
                     case .noExcess:
                         // There is no excess from previous reads. We obtain the next element from the stream.
-                        // Take the iterator from ReaderState, read one part, and put it back.
-                        // This ensures the iterator is always recoverable even if the reader
-                        // is dropped without consuming .end.
-                        guard var iterator = self.readerState.takeIterator() else {
-                            throw RequestBodyReadError.requestEndedBeforeReceivingEnd
-                        }
-
-                        let requestPart: HTTPRequestPart?
-                        do {
-                            requestPart = try await iterator.next(isolation: #isolation)
-                        } catch {
-                            // Put the iterator back before propagating the error.
-                            nonisolated(unsafe) let iter = iterator
-                            self.readerState.putIterator(iter)
-                            throw error
-                        }
+                        let requestPart = try await iterator?.next(isolation: #isolation)
 
                         switch requestPart {
                         case .head:
-                            nonisolated(unsafe) let iter = iterator
-                            self.readerState.putIterator(iter)
                             fatalError("Unexpectedly received a request head.")
 
                         case .none:
-                            // Stream ended without .end — don't put iterator back.
                             throw RequestBodyReadError.requestEndedBeforeReceivingEnd
 
                         case .body(let element):
-                            nonisolated(unsafe) let iter = iterator
-                            self.readerState.putIterator(iter)
                             bodyElement = element
 
                         case .end(let trailers):
                             self.state = .finished
-                            nonisolated(unsafe) let iter = iterator
-                            self.readerState.putIterator(iter)
                             self.readerState.wrapped.withLock { state in
                                 state.finishedReading = true
                                 state.trailers = trailers
@@ -142,17 +127,35 @@ public struct HTTPRequestConcludingAsyncReader: ConcludingAsyncReader, ~Copyable
                     return .readBody(bodyElement)
 
                 case .finished:
-                    return .requestFinished
+                    return .alreadyFinished
                 }
             }
         }
 
         var requestBodyStateMachine: RequestBodyStateMachine
 
+        /// The iterator that provides HTTP request parts from the underlying channel.
+        /// Owned by the reader so it can be returned to ReaderState in deinit
+        /// even if the handler didn't consume .end.
+        private var iterator: NIOAsyncChannelInboundStream<HTTPRequestPart>.AsyncIterator?
+
         /// Initializes a new request body reader.
+        ///
+        /// Takes the iterator from ReaderState so the reader owns it for the read cycle.
         fileprivate init(readerState: ReaderState) {
             self.requestBodyStateMachine = .init(readerState: readerState)
             self.state = readerState
+            self.iterator = readerState.takeIterator()
+        }
+
+        deinit {
+            // If the reader is dropped without having consumed .end, return the
+            // iterator to ReaderState so the server can drain unconsumed parts
+            // and reuse the connection.
+            if let iterator = self.iterator {
+                nonisolated(unsafe) let iter = iterator
+                self.state.putIterator(iter)
+            }
         }
 
         /// Reads a chunk of request body data.
@@ -167,7 +170,10 @@ public struct HTTPRequestConcludingAsyncReader: ConcludingAsyncReader, ~Copyable
         ) async throws(EitherError<ReadFailure, Failure>) -> Return {
             let readResult: RequestBodyStateMachine.ReadResult
             do {
-                readResult = try await self.requestBodyStateMachine.read(limit: maximumCount)
+                readResult = try await self.requestBodyStateMachine.read(
+                    iterator: &self.iterator,
+                    limit: maximumCount
+                )
             } catch {
                 throw .first(error)
             }
@@ -178,6 +184,15 @@ public struct HTTPRequestConcludingAsyncReader: ConcludingAsyncReader, ~Copyable
                     return try await body(Array(buffer: readElement).span)
 
                 case .requestFinished:
+                    // Return iterator to ReaderState for connection reuse.
+                    if let iter = self.iterator {
+                        nonisolated(unsafe) let iter = iter
+                        self.state.putIterator(iter)
+                    }
+                    self.iterator = nil
+                    return try await body(.init())
+
+                case .alreadyFinished:
                     return try await body(.init())
                 }
             } catch {