Coalesce outbound writes into a single buffer before a flush (#9)

Motivation:

The HTTP3 stream handler emits a ByteBuffer for each write it does. This
creates more ByteBuffer's and more pipeline traversals than is
necessary.

Modifications:

- Store the pending buffer in the stream handler and append to it for
each write. Coalesce write promises as well.
- Only emit the buffer when a flush happens.

Result:

Fewer allocations, fewer pipeline traversals.

Author: glbrntt

Sources/HTTP3/HTTP3StreamStateMachine.swift

@@ -370,14 +370,14 @@ package struct HTTP3StreamStateMachine: ~Copyable {
         }
 
         enum WriteAction {
-            /// Bytes are ready to be written out.
-            case writeBytes(ByteBuffer)
+            /// The frame's bytes were appended to the provided buffer.
+            case wroteBytes
             /// We need this header to be encoder.
             case encodeHeaders([HTTPField])
         }
 
-        /// Write a frame out.
-        mutating func write(frame: HTTP3Frame) -> WriteAction {
+        /// Write a frame out by appending its encoded bytes to `buffer`.
+        mutating func write(frame: HTTP3Frame, into buffer: inout ByteBuffer) -> WriteAction {
             switch consume self.state {
             case .idle(let idleState):
                 let maybePartial = MaybePartialFrame(frame)
@@ -394,24 +394,24 @@ package struct HTTP3StreamStateMachine: ~Copyable {
                     // So it won't get this far.
                     fatalError("Tried to write a push promise, which is not supported")
                 case .partial(let partial):
-                    var buffer = ByteBuffer()
                     buffer.writeHTTP3PartialFrame(partial, preferHuffmanEncoding: idleState.preferHuffmanEncoding)
                     self = .init(state: .idle(idleState))
-                    return .writeBytes(buffer)
+                    return .wroteBytes
                 }
             case .waitingForEncode:
                 fatalError("Cannot call write whilst waiting for a QPACK encode result")
             }
         }
 
         enum HeaderEncodeResultAction {
-            /// Bytes are ready to be written out.
-            case writeBytes(ByteBuffer)
+            /// The header's bytes were appended to the provided buffer.
+            case wroteBytes
         }
 
         mutating func gotHeaderEncodeResult(
             _ result: HTTP3PartialFrame.Headers,
-            from: [HTTPField]
+            from: [HTTPField],
+            into buffer: inout ByteBuffer
         ) -> HeaderEncodeResultAction {
             switch consume self.state {
             case .idle:
@@ -420,13 +420,12 @@ package struct HTTP3StreamStateMachine: ~Copyable {
                 guard from == waitingState.fields else {
                     fatalError("Unexpected encode result")
                 }
-                var buffer = ByteBuffer()
                 buffer.writeHTTP3PartialFrame(
                     .headers(result),
                     preferHuffmanEncoding: waitingState.preferHuffmanEncoding
                 )
                 self = .init(state: .idle(.init(preferHuffmanEncoding: waitingState.preferHuffmanEncoding)))
-                return .writeBytes(buffer)
+                return .wroteBytes
             }
         }
     }
@@ -466,8 +465,8 @@ package struct HTTP3StreamStateMachine: ~Copyable {
     }
 
     package enum WriteFrameAction {
-        /// You should write out the following bytes to the wire.
-        case returnBytes(ByteBuffer)
+        /// The frame's bytes were appended to the buffer you provided.
+        case wroteBytes
         /// You should encode the given headers and call back with the result.
         case encodeHeaders([HTTPField])
         /// The frame can't be written, because doing so would be a stream error.
@@ -480,8 +479,8 @@ package struct HTTP3StreamStateMachine: ~Copyable {
         case previousError
     }
 
-    /// Write out a frame.
-    package mutating func writeFrame(frame: HTTP3Frame) -> WriteFrameAction {
+    /// Write out a frame by appending its encoded bytes to `buffer`.
+    package mutating func writeFrame(frame: HTTP3Frame, into buffer: inout ByteBuffer) -> WriteFrameAction {
         switch self.state {
         case .idle(var idleState):
             guard idleState.readState.checkCanWrite() else {
@@ -491,11 +490,11 @@ package struct HTTP3StreamStateMachine: ~Copyable {
             let validationResult = idleState.validator.processOutboundFrame(frame)
             switch validationResult {
             case .forwardFrame(let validatedFrame):
-                let writeAction = idleState.writeState.write(frame: validatedFrame)
+                let writeAction = idleState.writeState.write(frame: validatedFrame, into: &buffer)
                 switch writeAction {
-                case .writeBytes(let bytes):
+                case .wroteBytes:
                     self = .init(state: .idle(idleState))
-                    return .returnBytes(bytes)
+                    return .wroteBytes
                 case .encodeHeaders(let fields):
                     self = .init(state: .idle(idleState))
                     return .encodeHeaders(fields)
@@ -520,8 +519,8 @@ package struct HTTP3StreamStateMachine: ~Copyable {
     }
 
     package enum HeaderEncodeResultAction {
-        /// You should write out the following bytes to the wire.
-        case returnBytes(ByteBuffer)
+        /// The header's bytes were appended to the buffer you provided.
+        case wroteBytes
         /// This header can't be encoded because the stream is already in an error state.
         case previousError(HTTP3Error)
         /// You should fail the current write because the stream is already closed
@@ -530,15 +529,16 @@ package struct HTTP3StreamStateMachine: ~Copyable {
 
     package mutating func gotHeaderEncodeResult(
         _ result: HTTP3PartialFrame.Headers,
-        from: [HTTPField]
+        from: [HTTPField],
+        into buffer: inout ByteBuffer
     ) -> HeaderEncodeResultAction {
         switch self.state {
         case .idle(var idleState):
-            let writeAction = idleState.writeState.gotHeaderEncodeResult(result, from: from)
+            let writeAction = idleState.writeState.gotHeaderEncodeResult(result, from: from, into: &buffer)
             self = .init(state: .idle(idleState))
             switch writeAction {
-            case .writeBytes(let bytes):
-                return .returnBytes(bytes)
+            case .wroteBytes:
+                return .wroteBytes
             }
         case .finished:
             // We shouldn't get a header decode result on a finished stream.

Sources/NIOHTTP3/HTTP3StreamHandler.swift

@@ -45,6 +45,12 @@ package final class HTTP3StreamHandler: ChannelDuplexHandler {
     /// The channel context. This handler can only be in one channel at a time.
     private var context: ChannelHandlerContext?
 
+    /// Bytes for frames that have been written but not yet flushed.
+    private var pendingBytes: ByteBuffer?
+
+    /// The promise which will be fulfilled when `pendingBytes` has been written.
+    private var pendingPromise: EventLoopPromise<Void>?
+
     /// The state machine which handles processing incoming bytes into frames, including validating them and decoding QPACK.
     private var stateMachine: HTTP3StreamStateMachine
 
@@ -79,6 +85,11 @@ package final class HTTP3StreamHandler: ChannelDuplexHandler {
 
     package func channelInactive(context: ChannelHandlerContext) {
         self.logger.trace("HTTP3StreamHandler.channelInactive")
+
+        // Don't leak the pending promise.
+        self.pendingBytes = nil
+        self.pendingPromise.take()?.fail(ChannelError.ioOnClosedChannel)
+
         // We want to flush out anything that's buffered which can be flushed.
         // There's unlikely to be anything...only if we got a channelInactive between a read and a readComplete.
         // We need to buffer any such actions into an array and save it for after we close the state machine
@@ -134,6 +145,10 @@ package final class HTTP3StreamHandler: ChannelDuplexHandler {
     }
 
     package func handlerRemoved(context: ChannelHandlerContext) {
+        // Don't leak the pending promise.
+        self.pendingBytes = nil
+        self.pendingPromise.take()?.fail(ChannelError.ioOnClosedChannel)
+
         // Cleanup reference to avoid leaks.
         self.context = nil
     }
@@ -187,7 +202,13 @@ package final class HTTP3StreamHandler: ChannelDuplexHandler {
     package func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
         let frame = self.unwrapOutboundIn(data)
         self.logger.trace("HTTP3StreamHandler.write", metadata: [LoggingKeys.h3FrameType: "\(frame.type)"])
-        let action = self.stateMachine.writeFrame(frame: frame)
+
+        if self.pendingBytes == nil {
+            self.pendingBytes = context.channel.allocator.buffer(capacity: 256)
+        }
+
+        let action = self.stateMachine.writeFrame(frame: frame, into: &self.pendingBytes!)
+
         switch action {
         case .previousError:
             // Just drop the byte
@@ -200,8 +221,8 @@ package final class HTTP3StreamHandler: ChannelDuplexHandler {
                     location: .here()
                 )
             )
-        case .returnBytes(let bytes):
-            context.write(wrapOutboundOut(bytes), promise: promise)
+        case .wroteBytes:
+            self.pendingPromise.setOrCascade(to: promise)
         case .wouldBeStreamError(let error):
             context.fireErrorCaught(error)
             promise?.fail(error)
@@ -213,7 +234,8 @@ package final class HTTP3StreamHandler: ChannelDuplexHandler {
             promise?.fail(error)
         case .encodeHeaders(let fields):
             let encoded = self.qpackEncoder(fields, self.streamID)
-            let action = self.stateMachine.gotHeaderEncodeResult(encoded, from: fields)
+            let action = self.stateMachine.gotHeaderEncodeResult(encoded, from: fields, into: &self.pendingBytes!)
+
             switch action {
             case .previousError(let previousError):
                 promise?.fail(
@@ -225,14 +247,41 @@ package final class HTTP3StreamHandler: ChannelDuplexHandler {
                         location: .here()
                     )
                 )
-            case .returnBytes(let bytes):
-                context.write(wrapOutboundOut(bytes), promise: promise)
+            case .wroteBytes:
+                self.pendingPromise.setOrCascade(to: promise)
             case .alreadyClosed:
                 promise?.fail(ChannelError.ioOnClosedChannel)
             }
         }
     }
 
+    package func flush(context: ChannelHandlerContext) {
+        self.emitPendingBytes(context: context)
+        context.flush()
+    }
+
+    package func close(
+        context: ChannelHandlerContext,
+        mode: CloseMode,
+        promise: EventLoopPromise<Void>?
+    ) {
+        switch mode {
+        case .output, .all:
+            self.emitPendingBytes(context: context)
+        case .input:
+            ()
+        }
+        context.close(mode: mode, promise: promise)
+    }
+
+    /// Write any pending bytes.
+    private func emitPendingBytes(context: ChannelHandlerContext) {
+        if let bytes = self.pendingBytes.take() {
+            let promise = self.pendingPromise.take()
+            context.write(HTTP3StreamHandler.wrapOutboundOut(bytes), promise: promise)
+        }
+    }
+
     package func errorCaught(context: ChannelHandlerContext, error: any Error) {
         switch error {
         case let error as QUICStreamResetError:

Tests/HTTP3Tests/HTTP3StreamStateMachineTests.swift

@@ -836,6 +836,21 @@ extension HTTP3StreamStateMachine.WriteFrameAction {
 }
 
 extension HTTP3StreamStateMachine {
+    /// Test convenience: write a frame using a throwaway buffer, for assertions that don't inspect bytes.
+    fileprivate mutating func writeFrame(frame: HTTP3Frame) -> WriteFrameAction {
+        var buffer = ByteBuffer()
+        return self.writeFrame(frame: frame, into: &buffer)
+    }
+
+    /// Test convenience: deliver an encode result using a throwaway buffer.
+    fileprivate mutating func gotHeaderEncodeResult(
+        _ result: HTTP3PartialFrame.Headers,
+        from: [HTTPField]
+    ) -> HeaderEncodeResultAction {
+        var buffer = ByteBuffer()
+        return self.gotHeaderEncodeResult(result, from: from, into: &buffer)
+    }
+
     fileprivate enum ResolvedAction {
         case returnBytes(ByteBuffer)
         case wouldBeStreamError(HTTP3Error)
@@ -846,12 +861,13 @@ extension HTTP3StreamStateMachine {
     /// Do a write, and do the qpack too, and return just one action.
     fileprivate mutating func writeFrameAndQPACK(frame: HTTP3Frame) -> ResolvedAction? {
         let encoder = StaticQPACKEncoder()
-        let action = self.writeFrame(frame: frame)
+        var buffer = ByteBuffer()
+        let action = self.writeFrame(frame: frame, into: &buffer)
         switch action {
         case .previousError:
             return nil
-        case .returnBytes(let bytes):
-            return .returnBytes(bytes)
+        case .wroteBytes:
+            return .returnBytes(buffer)
         case .wouldBeStreamError(let error):
             return .wouldBeStreamError(error)
         case .wouldBeConnectionError(let error):
@@ -860,10 +876,10 @@ extension HTTP3StreamStateMachine {
             return .alreadyClosed
         case .encodeHeaders(let fields):
             let qpackResult = encoder.encode(headers: fields)
-            let action2 = self.gotHeaderEncodeResult(.init(fieldSection: qpackResult), from: fields)
+            let action2 = self.gotHeaderEncodeResult(.init(fieldSection: qpackResult), from: fields, into: &buffer)
             switch action2 {
-            case .returnBytes(let bytes):
-                return .returnBytes(bytes)
+            case .wroteBytes:
+                return .returnBytes(buffer)
             case .previousError:
                 return nil
             case .alreadyClosed: