Add a server connection manager

NOTICE.txt

@@ -59,3 +59,12 @@ This product contains a fuzz testing harness derived from Swift Protobuf.
     * https://github.com/apple/swift-protobuf/blob/main/LICENSE.txt
   * HOMEPAGE:
     * https://github.com/apple/swift-protobuf
+
+---
+
+This product contains a graceful shutdown connection manager derived from gRPC Swift NIO Transport.
+
+  * LICENSE (Apache License 2.0):
+    * https://github.com/grpc/grpc-swift-nio-transport/blob/main/LICENSE
+  * HOMEPAGE:
+    * https://github.com/gRPC/gRPC-swift-nio-transport

Sources/NIOHTTP2/HTTP2PipelineHelpers.swift

@@ -927,6 +927,60 @@ extension ChannelPipeline.SynchronousOperations {
         )
     }
 
+    /// Configures a `ChannelPipeline` to speak HTTP/2 and sets up mapping functions so that it may be interacted with
+    /// from concurrent code. This variant includes a ``NIOHTTP2ServerConnectionManagementHandler`` argument, and links
+    /// its stream and frame delegates with ``NIOHTTP2Handler``.
+    ///
+    /// This operation **must** be called on the event loop.
+    ///
+    /// In general this is not entirely useful by itself, as HTTP/2 is a negotiated protocol. This helper does not
+    /// handle negotiation. Instead, this simply adds the handler required to speak HTTP/2 after negotiation has
+    /// completed, or when agreed by prior knowledge. Use this function to setup a HTTP/2 pipeline if you wish to use
+    /// async sequence abstractions over inbound and outbound streams, as it allows that pipeline to evolve without
+    /// breaking your code.
+    ///
+    /// - Parameters:
+    ///   - mode: The mode this pipeline will operate in, server or client.
+    ///   - connectionManager: A ``NIOHTTP2ServerConnectionManagementHandler`` instance to use alongised the
+    ///   ``NIOHTTP2Handler``.
+    ///   - configuration: The settings that will be used when establishing the connection and new streams.
+    ///   - streamInitializer: A closure that will be called whenever the remote peer initiates a new stream.
+    ///     The output of this closure is the element type of the returned multiplexer
+    /// - Returns: An `EventLoopFuture` containing the `AsyncStreamMultiplexer` inserted into this pipeline, which can
+    ///   be used to initiate new streams and iterate over inbound HTTP/2 stream channels.
+    @inlinable
+    @available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
+    public func configureAsyncHTTP2Pipeline<Output: Sendable>(
+        mode: NIOHTTP2Handler.ParserMode,
+        connectionManager: NIOHTTP2ServerConnectionManagementHandler,
+        configuration: NIOHTTP2Handler.Configuration = NIOHTTP2Handler.Configuration(),
+        streamInitializer: @escaping NIOChannelInitializerWithOutput<Output>
+    ) throws -> NIOHTTP2Handler.AsyncStreamMultiplexer<Output> {
+        let handler = NIOHTTP2Handler(
+            mode: mode,
+            eventLoop: self.eventLoop,
+            connectionConfiguration: configuration.connection,
+            streamConfiguration: configuration.stream,
+            streamDelegate: connectionManager.http2StreamDelegate,
+            frameDelegate: nil,
+            inboundStreamInitializerWithAnyOutput: { channel in
+                streamInitializer(channel).map { $0 }
+            }
+        )
+
+        try self.addHandler(handler)
+        try self.addHandler(connectionManager)
+
+        let (inboundStreamChannels, continuation) = NIOHTTP2AsyncSequence.initialize(
+            inboundStreamInitializerOutput: Output.self
+        )
+
+        return try handler.syncAsyncStreamMultiplexer(
+            continuation: continuation,
+            inboundStreamChannels: inboundStreamChannels
+        )
+    }
+
     @inlinable
     func configureHTTP2AsyncSecureUpgrade<HTTP1Output: Sendable, HTTP2Output: Sendable>(
         on channel: any Channel,

Sources/NIOHTTP2/ServerConnectionManagementHandler+StateMachine.swift

@@ -0,0 +1,292 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftNIO open source project
+//
+// Copyright (c) 2026 Apple Inc. and the SwiftNIO project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of SwiftNIO project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+/*
+ * Copyright 2024, gRPC Authors All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import NIOCore
+
+extension NIOHTTP2ServerConnectionManagementHandler {
+    /// Tracks the state of TCP connections at the server.
+    ///
+    /// The state machine manages the state for the graceful shutdown procedure.
+    struct StateMachine {
+        /// Current state.
+        private var state: State
+
+        /// Opaque data sent to the client in a PING frame after emitting the first GOAWAY frame
+        /// as part of graceful shutdown.
+        private let goAwayPingData: HTTP2PingData
+
+        /// Whether the connection is currently closing.
+        var isClosing: Bool {
+            self.state.isClosing
+        }
+
+        /// Create a new state machine.
+        ///
+        /// - Parameters:
+        ///   - goAwayPingData: Opaque data sent to the client in a PING frame when the server
+        ///       initiates graceful shutdown.
+        init(goAwayPingData: HTTP2PingData = HTTP2PingData(withInteger: .random(in: .min ... .max))) {
+            self.state = .active(State.Active())
+            self.goAwayPingData = goAwayPingData
+        }
+
+        /// Record that the stream with the given ID has been opened.
+        mutating func streamOpened(_ id: HTTP2StreamID) {
+            switch self.state {
+            case .active(var state):
+                self.state = ._modifying
+                state.lastStreamID = id
+                let (inserted, _) = state.openStreams.insert(id)
+                assert(inserted, "Can't open stream \(Int(id)), it's already open")
+                self.state = .active(state)
+
+            case .closing(var state):
+                self.state = ._modifying
+                state.lastStreamID = id
+                let (inserted, _) = state.openStreams.insert(id)
+                assert(inserted, "Can't open stream \(Int(id)), it's already open")
+                self.state = .closing(state)
+
+            case .closed:
+                ()
+
+            case ._modifying:
+                preconditionFailure()
+            }
+        }
+
+        enum OnStreamClosed: Equatable {
+            /// Start the idle timer, after which the connection should be closed gracefully.
+            case startIdleTimer
+            /// Close the connection.
+            case close
+            /// Do nothing.
+            case none
+        }
+
+        /// Record that the stream with the given ID has been closed.
+        mutating func streamClosed(_ id: HTTP2StreamID) -> OnStreamClosed {
+            let onStreamClosed: OnStreamClosed
+
+            switch self.state {
+            case .active(var state):
+                self.state = ._modifying
+                let removedID = state.openStreams.remove(id)
+                assert(removedID != nil, "Can't close stream \(Int(id)), it wasn't open")
+                onStreamClosed = state.openStreams.isEmpty ? .startIdleTimer : .none
+                self.state = .active(state)
+
+            case .closing(var state):
+                self.state = ._modifying
+                let removedID = state.openStreams.remove(id)
+                assert(removedID != nil, "Can't close stream \(Int(id)), it wasn't open")
+                // If the second GOAWAY hasn't been sent it isn't safe to close if there are no open
+                // streams: the client may have opened a stream which the server doesn't know about yet.
+                let canClose = state.sentSecondGoAway && state.openStreams.isEmpty
+                onStreamClosed = canClose ? .close : .none
+                self.state = .closing(state)
+
+            case .closed:
+                onStreamClosed = .none
+
+            case ._modifying:
+                preconditionFailure()
+            }
+
+            return onStreamClosed
+        }
+
+        enum OnPing: Equatable {
+            /// Send a GOAWAY frame with the code "enhance your calm" and immediately close the connection.
+            case enhanceYourCalmThenClose(HTTP2StreamID)
+            /// Acknowledge the ping.
+            case sendAck
+            /// Ignore the ping.
+            case none
+        }
+
+        /// Received a ping with the given data.
+        ///
+        /// - Parameters:
+        ///   - time: The time at which the ping was received.
+        ///   - data: The data sent with the ping.
+        mutating func receivedPing(atTime time: NIODeadline, data: HTTP2PingData) -> OnPing {
+            let onPing: OnPing
+
+            switch self.state {
+            case .active(let state):
+                self.state = ._modifying
+
+                onPing = .sendAck
+                self.state = .active(state)
+
+            case .closing(let state):
+                self.state = ._modifying
+
+                onPing = .sendAck
+                self.state = .closing(state)
+
+            case .closed:
+                onPing = .none
+
+            case ._modifying:
+                preconditionFailure()
+            }
+
+            return onPing
+        }
+
+        enum OnPingAck: Equatable {
+            /// Send a GOAWAY frame with no error and the given last stream ID, optionally closing the
+            /// connection immediately afterwards.
+            case sendGoAway(lastStreamID: HTTP2StreamID, close: Bool)
+            /// Ignore the ack.
+            case none
+        }
+
+        /// Received a PING frame with the 'ack' flag set.
+        mutating func receivedPingAck(data: HTTP2PingData) -> OnPingAck {
+            let onPingAck: OnPingAck
+
+            switch self.state {
+            case .closing(var state):
+                self.state = ._modifying
+
+                // If only one GOAWAY has been sent and the data matches the data from the GOAWAY ping then
+                // the server should send another GOAWAY ratcheting down the last stream ID. If no streams
+                // are open then the server can close the connection immediately after, otherwise it must
+                // wait until all streams are closed.
+                if !state.sentSecondGoAway, data == self.goAwayPingData {
+                    state.sentSecondGoAway = true
+
+                    if state.openStreams.isEmpty {
+                        self.state = .closed
+                        onPingAck = .sendGoAway(lastStreamID: state.lastStreamID, close: true)
+                    } else {
+                        self.state = .closing(state)
+                        onPingAck = .sendGoAway(lastStreamID: state.lastStreamID, close: false)
+                    }
+                } else {
+                    onPingAck = .none
+                }
+
+                self.state = .closing(state)
+
+            case .active, .closed:
+                onPingAck = .none
+
+            case ._modifying:
+                preconditionFailure()
+            }
+
+            return onPingAck
+        }
+
+        enum OnStartGracefulShutdown: Equatable {
+            /// Initiate graceful shutdown by sending a GOAWAY frame with the last stream ID set as the max
+            /// stream ID and no error. Follow it immediately with a PING frame with the given data.
+            case sendGoAwayAndPing(HTTP2PingData)
+            /// Ignore the request to start graceful shutdown.
+            case none
+        }
+
+        /// Request that the connection begins graceful shutdown.
+        mutating func startGracefulShutdown() -> OnStartGracefulShutdown {
+            let onStartGracefulShutdown: OnStartGracefulShutdown
+
+            switch self.state {
+            case .active(let state):
+                self.state = .closing(State.Closing(from: state))
+                onStartGracefulShutdown = .sendGoAwayAndPing(self.goAwayPingData)
+
+            case .closing, .closed:
+                onStartGracefulShutdown = .none
+
+            case ._modifying:
+                preconditionFailure()
+            }
+
+            return onStartGracefulShutdown
+        }
+
+        /// Marks the state as closed.
+        mutating func markClosed() {
+            self.state = .closed
+        }
+    }
+}
+
+extension NIOHTTP2ServerConnectionManagementHandler.StateMachine {
+    fileprivate enum State {
+        /// The connection is active.
+        struct Active {
+            /// The number of open streams.
+            var openStreams: Set<HTTP2StreamID>
+            /// The ID of the most recently opened stream (zero indicates no streams have been opened yet).
+            var lastStreamID: HTTP2StreamID
+
+            init() {
+                self.openStreams = []
+                self.lastStreamID = .rootStream
+            }
+        }
+
+        /// The connection is closing gracefully, an initial GOAWAY frame has been sent (with the
+        /// last stream ID set to max).
+        struct Closing {
+            /// The number of open streams.
+            var openStreams: Set<HTTP2StreamID>
+            /// The ID of the most recently opened stream (zero indicates no streams have been opened yet).
+            var lastStreamID: HTTP2StreamID
+            /// Whether the second GOAWAY frame has been sent with a lower stream ID.
+            var sentSecondGoAway: Bool
+
+            init(from state: Active) {
+                self.openStreams = state.openStreams
+                self.lastStreamID = state.lastStreamID
+                self.sentSecondGoAway = false
+            }
+        }
+
+        case active(Active)
+        case closing(Closing)
+        case closed
+        case _modifying
+
+        var isClosing: Bool {
+            switch self {
+            case .closing:
+                return true
+            case .active, .closed, ._modifying:
+                return false
+            }
+        }
+    }
+}