Generic helpers for HTTP/2 async pipelines

Sources/NIOHTTP2/HTTP2ChannelHandler+InlineStreamMultiplexer.swift

@@ -147,6 +147,10 @@ extension InlineStreamMultiplexer {
     internal func createStreamChannel(_ streamStateInitializer: @escaping (Channel) -> EventLoopFuture<Void>) -> EventLoopFuture<Channel> {
         self.commonStreamMultiplexer.createStreamChannel(multiplexer: .inline(self), streamStateInitializer)
     }
+
+    internal func createStreamChannel() -> EventLoopFuture<Channel> {
+        self.commonStreamMultiplexer.createStreamChannel(multiplexer: .inline(self))
+    }
 }
 
 extension NIOHTTP2Handler {
@@ -201,3 +205,46 @@ extension NIOHTTP2Handler {
         }
     }
 }
+
+extension InlineStreamMultiplexer {
+    func setChannelContinuation(_ streamChannels: any ChannelContinuation) {
+        self.commonStreamMultiplexer.setChannelContinuation(streamChannels)
+    }
+}
+
+extension NIOHTTP2Handler {
+    /// A variant of `NIOHTTP2Handler.StreamMultiplexer` which creates a child channel for each HTTP/2 stream and
+    /// provides access to inbound HTTP/2 streams.
+    ///
+    /// In general in NIO applications it is helpful to consider each HTTP/2 stream as an
+    /// independent stream of HTTP/2 frames. This multiplexer achieves this by creating a
+    /// number of in-memory `HTTP2StreamChannel` objects, one for each stream. These operate
+    /// on ``HTTP2Frame/FramePayload`` objects as their base communication
+    /// atom, as opposed to the regular NIO `SelectableChannel` objects which use `ByteBuffer`
+    /// and `IOData`.
+    ///
+    /// The stream channel objects are initialized upon creation using the supplied `streamStateInitializer` which returns a type
+    /// `Output`. This type may be `HTTP2Frame` or changed to any other type.
+    @available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
+    @_spi(AsyncChannel)
+    public struct AsyncStreamMultiplexer<Output> {
+        private let inlineStreamMultiplexer: InlineStreamMultiplexer
+        public let inbound: NIOHTTP2InboundStreamChannels<Output>
+
+        // Cannot be created by users.
+        internal init(_ inlineStreamMultiplexer: InlineStreamMultiplexer, continuation: any ChannelContinuation, inboundStreamChannels: NIOHTTP2InboundStreamChannels<Output>) {
+            self.inlineStreamMultiplexer = inlineStreamMultiplexer
+            self.inlineStreamMultiplexer.setChannelContinuation(continuation)
+            self.inbound = inboundStreamChannels
+        }
+
+        public func createStreamChannel<Output>(_ streamStateInitializer: @escaping @Sendable (Channel) -> EventLoopFuture<Output>) async throws -> Output {
+            return try await self.inlineStreamMultiplexer.createStreamChannel().flatMap { streamChannel in
+                let http2StreamChannel = streamChannel as! HTTP2StreamChannel
+                let promise = streamChannel.eventLoop.makePromise(of: Output.self)
+                http2StreamChannel.configure(initializer: streamStateInitializer, userPromise: promise)
+                return promise.futureResult
+            }.get()
+        }
+    }
+}

Sources/NIOHTTP2/HTTP2ChannelHandler.swift

@@ -1076,4 +1076,16 @@ extension NIOHTTP2Handler {
             throw NIOHTTP2Errors.missingMultiplexer()
         }
     }
+
+    @available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
+    internal func syncAsyncStreamMultiplexer<Output>(continuation: any ChannelContinuation, inboundStreamChannels: NIOHTTP2InboundStreamChannels<Output>) throws -> AsyncStreamMultiplexer<Output> {
+        self.eventLoop!.preconditionInEventLoop()
+
+        switch self.inboundStreamMultiplexer {
+        case let .some(.inline(multiplexer)):
+            return AsyncStreamMultiplexer(multiplexer, continuation: continuation, inboundStreamChannels: inboundStreamChannels)
+        case .some(.legacy), .none:
+            throw NIOHTTP2Errors.missingMultiplexer()
+        }
+    }
 }

Sources/NIOHTTP2/HTTP2CommonInboundStreamMultiplexer.swift

@@ -38,7 +38,16 @@ internal class HTTP2CommonInboundStreamMultiplexer {
     private var isReading = false
     private var flushPending = false
 
-    init(mode: NIOHTTP2Handler.ParserMode, channel: Channel, inboundStreamStateInitializer: MultiplexerAbstractChannel.InboundStreamStateInitializer, targetWindowSize: Int, streamChannelOutboundBytesHighWatermark: Int, streamChannelOutboundBytesLowWatermark: Int) {
+    var streamChannels: (any ChannelContinuation)?
+
+    init(
+        mode: NIOHTTP2Handler.ParserMode,
+        channel: Channel,
+        inboundStreamStateInitializer: MultiplexerAbstractChannel.InboundStreamStateInitializer,
+        targetWindowSize: Int,
+        streamChannelOutboundBytesHighWatermark: Int,
+        streamChannelOutboundBytesLowWatermark: Int
+    ) {
         self.channel = channel
         self.inboundStreamStateInitializer = inboundStreamStateInitializer
         self.targetWindowSize = targetWindowSize
@@ -93,6 +102,11 @@ extension HTTP2CommonInboundStreamMultiplexer {
             )
 
             self.streams[streamID] = channel
+
+            // If we have an async sequence of inbound stream channels yield the channel to it
+            // This also implicitly performs the stream initialization step
+            self.streamChannels?.yield(channel: channel.baseChannel)
+
             channel.configureInboundStream(initializer: self.inboundStreamStateInitializer)
             channel.receiveInboundFrame(frame)
 
@@ -178,6 +192,8 @@ extension HTTP2CommonInboundStreamMultiplexer {
         for channel in self.pendingStreams.values {
             channel.receiveStreamClosed(nil)
         }
+        // there cannot be any more inbound streams now that the connection channel is inactive
+        self.streamChannels?.finish()
     }
 
     internal func propagateChannelWritabilityChanged(context: ChannelHandlerContext) {
@@ -265,20 +281,59 @@ extension HTTP2CommonInboundStreamMultiplexer {
 }
 
 extension HTTP2CommonInboundStreamMultiplexer {
+    internal func _createStreamChannel(_ multiplexer: HTTP2StreamChannel.OutboundStreamMultiplexer, _ promise: EventLoopPromise<Channel>?) {
+        let channel = MultiplexerAbstractChannel(
+            allocator: self.channel.allocator,
+            parent: self.channel,
+            multiplexer: multiplexer,
+            streamID: nil,
+            targetWindowSize: Int32(self.targetWindowSize),
+            outboundBytesHighWatermark: self.streamChannelOutboundBytesHighWatermark,
+            outboundBytesLowWatermark: self.streamChannelOutboundBytesLowWatermark,
+            inboundStreamStateInitializer: .excludesStreamID(nil)
+        )
+        self.pendingStreams[channel.channelID] = channel
+        promise?.succeed(channel.baseChannel)
+    }
+
+    internal func createStreamChannel(multiplexer: HTTP2StreamChannel.OutboundStreamMultiplexer, promise: EventLoopPromise<Channel>?) {
+        if self.channel.eventLoop.inEventLoop {
+            self._createStreamChannel(multiplexer, promise)
+        } else {
+            self.channel.eventLoop.execute {
+                self._createStreamChannel(multiplexer, promise)
+            }
+        }
+    }
+
+    internal func createStreamChannel(multiplexer: HTTP2StreamChannel.OutboundStreamMultiplexer) -> EventLoopFuture<Channel> {
+        let promise = self.channel.eventLoop.makePromise(of: Channel.self)
+        self.createStreamChannel(multiplexer: multiplexer, promise: promise)
+        return promise.futureResult
+    }
+
+    internal func _createStreamChannel(_ multiplexer: HTTP2StreamChannel.OutboundStreamMultiplexer, _ promise: EventLoopPromise<Channel>?, _ streamStateInitializer: @escaping (Channel) -> EventLoopFuture<Void>) {
+        let channel = MultiplexerAbstractChannel(
+            allocator: self.channel.allocator,
+            parent: self.channel,
+            multiplexer: multiplexer,
+            streamID: nil,
+            targetWindowSize: Int32(self.targetWindowSize),
+            outboundBytesHighWatermark: self.streamChannelOutboundBytesHighWatermark,
+            outboundBytesLowWatermark: self.streamChannelOutboundBytesLowWatermark,
+            inboundStreamStateInitializer: .excludesStreamID(nil)
+        )
+        self.pendingStreams[channel.channelID] = channel
+        channel.configure(initializer: streamStateInitializer, userPromise: promise)
+    }
+
     internal func createStreamChannel(multiplexer: HTTP2StreamChannel.OutboundStreamMultiplexer, promise: EventLoopPromise<Channel>?, _ streamStateInitializer: @escaping (Channel) -> EventLoopFuture<Void>) {
-        self.channel.eventLoop.execute {
-            let channel = MultiplexerAbstractChannel(
-                allocator: self.channel.allocator,
-                parent: self.channel,
-                multiplexer: multiplexer,
-                streamID: nil,
-                targetWindowSize: Int32(self.targetWindowSize),
-                outboundBytesHighWatermark: self.streamChannelOutboundBytesHighWatermark,
-                outboundBytesLowWatermark: self.streamChannelOutboundBytesLowWatermark,
-                inboundStreamStateInitializer: .excludesStreamID(nil)
-            )
-            self.pendingStreams[channel.channelID] = channel
-            channel.configure(initializer: streamStateInitializer, userPromise: promise)
+        if self.channel.eventLoop.inEventLoop {
+            self._createStreamChannel(multiplexer, promise, streamStateInitializer)
+        } else {
+            self.channel.eventLoop.execute {
+                self._createStreamChannel(multiplexer, promise, streamStateInitializer)
+            }
         }
     }
 
@@ -321,3 +376,112 @@ extension HTTP2CommonInboundStreamMultiplexer {
         }
     }
 }
+
+extension HTTP2CommonInboundStreamMultiplexer {
+    func setChannelContinuation(_ streamChannels: any ChannelContinuation) {
+        self.streamChannels = streamChannels
+    }
+}
+
+/// `ChannelContinuation` is used to generic async-sequence-like objects to deal with `Channel`s. This is so that they may be held
+/// by the `HTTP2ChannelHandler` without causing it to become generic itself.
+internal protocol ChannelContinuation {
+    func yield(channel: Channel)
+    func finish()
+    func finish(throwing error: Error)
+}
+
+
+/// `StreamChannelContinuation` is a wrapper for a generic `AsyncThrowingStream` which holds the inbound HTTP2 stream channels.
+///
+@available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
+struct StreamChannelContinuation<Output>: ChannelContinuation {
+    private var continuation: AsyncThrowingStream<Output, Error>.Continuation
+    private let streamStateInitializer: (Channel) -> EventLoopFuture<Output>
+
+    private init(
+        continuation: AsyncThrowingStream<Output, Error>.Continuation,
+        streamStateInitializer: @escaping (Channel) -> EventLoopFuture<Output>
+    ) {
+        self.continuation = continuation
+        self.streamStateInitializer = streamStateInitializer
+    }
+
+    /// `initialize` creates a new `StreamChannels` object and returns it along with its backing `AsyncThrowingStream`.
+    ///
+    /// - Parameters:
+    ///   - streamInitializer: A closure which initializes the newly-created inbound stream channel and returns a generic.
+    ///   The returned type corresponds to the output of the channel once the operations in the initializer have been performed.
+    ///   For example a `streamStateInitializer` which inserts handlers before wrapping the channel in a `NIOAsyncChannel` would
+    ///   have a `Output` corresponding to that `NIOAsyncChannel` type. Another example is in cases where there is
+    ///   per-stream protocol negotiation where `Output` would be some form of `NIOProtocolNegotiationResult`.
+    static func initialize(
+        with streamStateInitializer: @escaping (Channel) -> EventLoopFuture<Output>
+    ) -> (StreamChannelContinuation<Output>, NIOHTTP2InboundStreamChannels<Output>) {
+        var continuation: AsyncThrowingStream<Output, Error>.Continuation? = nil
+        let stream = AsyncThrowingStream { continuation = $0 }
+        return (StreamChannelContinuation(continuation: continuation!, streamStateInitializer: streamStateInitializer), NIOHTTP2InboundStreamChannels(stream))
+    }
+
+    /// `yield` takes a channel, executes the stored `streamInitializer` upon it and then yields the *derived* type to
+    /// the wrapped `AsyncThrowingStream`.
+    func yield(channel: Channel) {
+        channel.eventLoop.assertInEventLoop()
+        self.streamStateInitializer(channel).whenSuccess{ streamChannel in
+            let yieldResult = self.continuation.yield(streamChannel)
+            switch yieldResult {
+            case .enqueued:
+                break // success, nothing to do
+            case .dropped:
+                preconditionFailure("Attempted to yield channel when AsyncThrowingStream is over capacity. This shouldn't be possible for an unbounded stream.")
+            case .terminated:
+                channel.close(mode: .all, promise: nil)
+                assertionFailure("Attempted to yield channel to AsyncThrowingStream in terminated state.")
+            default:
+                channel.close(mode: .all, promise: nil)
+                assertionFailure("Attempt to yield channel to AsyncThrowingStream failed for unhandled reason.")
+            }
+        }
+    }
+
+    /// `finish` marks the continuation as finished.
+    func finish() {
+        self.continuation.finish()
+    }
+
+    /// `finish` marks the continuation as finished with the supplied error.
+    func finish(throwing error: Error) {
+        self.continuation.finish(throwing: error)
+    }
+}
+
+/// `NIOHTTP2InboundStreamChannels` provides access to inbound stream channels as an `AsyncSequence`.
+@available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
+@_spi(AsyncChannel)
+public struct NIOHTTP2InboundStreamChannels<Output>: AsyncSequence {
+    public struct AsyncIterator: AsyncIteratorProtocol {
+        public typealias Element = Output
+
+        internal var iterator: AsyncThrowingStream<Output, Error>.AsyncIterator
+
+        public mutating func next() async throws -> Output? {
+            try await self.iterator.next()
+        }
+
+        init(_ iterator: AsyncThrowingStream<Output, Error>.AsyncIterator) {
+            self.iterator = iterator
+        }
+    }
+
+    public typealias Element = Output
+
+    public func makeAsyncIterator() -> AsyncIterator {
+        AsyncIterator(self.asyncThrowingStream.makeAsyncIterator())
+    }
+
+    let asyncThrowingStream: AsyncThrowingStream<Output, Error>
+
+    init(_ asyncThrowingStream: AsyncThrowingStream<Output, Error>) {
+        self.asyncThrowingStream = asyncThrowingStream
+    }
+}