Close active connections while stopping an HTTPServer

Package.swift

@@ -30,6 +30,7 @@ let package = Package(
         // Dependencies declare other packages that this package depends on.
         .package(url: "https://github.com/apple/swift-nio.git", from: "2.3.0"),
         .package(url: "https://github.com/apple/swift-nio-ssl.git", from: "2.0.0"),
+        .package(url: "https://github.com/apple/swift-nio-extras.git", from: "1.0.0"),
         .package(url: "https://github.com/IBM-Swift/BlueSSLService.git", from: "1.0.0"),
         .package(url: "https://github.com/IBM-Swift/LoggerAPI.git", from: "1.7.3")
     ],
@@ -41,7 +42,7 @@ let package = Package(
             dependencies: []),
         .target(
             name: "KituraNet",
-            dependencies: ["NIO", "NIOFoundationCompat", "NIOHTTP1", "NIOSSL", "SSLService", "LoggerAPI", "NIOWebSocket", "CLinuxHelpers"]),
+            dependencies: ["NIO", "NIOFoundationCompat", "NIOHTTP1", "NIOSSL", "SSLService", "LoggerAPI", "NIOWebSocket", "CLinuxHelpers", "NIOExtras"]),
         .testTarget(
             name: "KituraNetTests",
             dependencies: ["KituraNet"])

Sources/KituraNet/HTTP/HTTPServer.swift

@@ -22,6 +22,7 @@ import SSLService
 import LoggerAPI
 import NIOWebSocket
 import CLinuxHelpers
+import NIOExtras
 
 #if os(Linux)
 import Glibc
@@ -124,6 +125,8 @@ public class HTTPServer: Server {
     /// The event loop group on which the HTTP handler runs
     private let eventLoopGroup: MultiThreadedEventLoopGroup
 
+    var quiescingHelper: ServerQuiescingHelper?
+
     /**
      Creates an HTTP server object.
 
@@ -299,6 +302,11 @@ public class HTTPServer: Server {
             .serverChannelOption(ChannelOptions.backlog, value: BacklogOption.Value(self.maxPendingConnections))
             .serverChannelOption(ChannelOptions.socket(SocketOptionLevel(SOL_SOCKET), SO_REUSEADDR), value: 1)
             .serverChannelOption(ChannelOptions.socket(SocketOptionLevel(SOL_SOCKET), SO_REUSEPORT), value: allowPortReuse ? 1 : 0)
+            .serverChannelInitializer { channel in
+                // Adding the quiescing helper will help us do a graceful stop()
+                self.quiescingHelper = ServerQuiescingHelper(group: self.eventLoopGroup)
+                return channel.pipeline.addHandler(self.quiescingHelper!.makeServerChannelHandler(channel: channel))
+            }
             .childChannelInitializer { channel in
                 let httpHandler = HTTPRequestHandler(for: self)
                 let config: NIOHTTPServerUpgradeConfiguration = (upgraders: upgraders, completionHandler: { _ in
@@ -463,13 +471,21 @@ public class HTTPServer: Server {
      ````
     */
     public func stop() {
+        // Close the listening channel
         guard let serverChannel = serverChannel else { return }
         do {
             try serverChannel.close().wait()
         } catch let error {
             Log.error("Failed to close the server channel. Error: \(error)")
         }
-        self.state = .stopped
+
+        // Now close all the open channels
+        guard let quiescingHelper = self.quiescingHelper else { return }
+        let fullShutdownPromise: EventLoopPromise<Void> = eventLoopGroup.next().makePromise()
+        quiescingHelper.initiateShutdown(promise: fullShutdownPromise)
+        fullShutdownPromise.futureResult.whenComplete { _ in
+            self.state = .stopped
+        }
     }
 
     /**

Tests/KituraNetTests/ChannelQuiescingTests.swift

@@ -0,0 +1,72 @@
+import NIO
+import NIOHTTP1
+import XCTest
+import KituraNet
+
+
+class ChannelQuiescingTests: KituraNetTest {
+
+    static var allTests: [(String, (ChannelQuiescingTests) -> () throws -> Void)] {
+        return [
+            ("testChannelQuiescing", testChannelQuiescing),
+        ]
+    }
+
+    func testChannelQuiescing() {
+        let server = HTTP.createServer()
+        try! server.listen(on: 0)
+        let port = server.port ?? -1
+        server.delegate = SleepingDelegate()
+
+        let connectionClosedExpectation = expectation(description: "Server closes connections")
+        let bootstrap = ClientBootstrap(group: MultiThreadedEventLoopGroup(numberOfThreads: 1))
+            .channelOption(ChannelOptions.socket(SocketOptionLevel(SOL_SOCKET), SO_REUSEADDR), value: 1)
+            .channelInitializer { channel in
+                channel.pipeline.addHTTPClientHandlers().flatMap {
+                    channel.pipeline.addHandler(HTTPHandler(connectionClosedExpectation))
+                }
+        }
+        let request = HTTPRequestHead(version: HTTPVersion.init(major: 1, minor: 1), method: .GET, uri: "/")
+
+        // Make the first connection
+        let channel1 = try! bootstrap.connect(host: "localhost", port: port).wait()
+        _ = channel1.write(NIOAny(HTTPClientRequestPart.head(request)), promise: nil)
+        try! channel1.writeAndFlush(NIOAny(HTTPClientRequestPart.end(nil))).wait()
+
+        // Make the second connection
+        let channel2 = try! bootstrap.connect(host: "localhost", port: port).wait()
+        _ = channel2.write(NIOAny(HTTPClientRequestPart.head(request)), promise: nil)
+        try! channel2.writeAndFlush(NIOAny(HTTPClientRequestPart.end(nil))).wait()
+
+        // The server must close both the connections
+        connectionClosedExpectation.expectedFulfillmentCount = 2
+
+        // Give time for the route handlers to kick in
+        sleep(1)
+
+        // Stop the server
+        server.stop()
+        waitForExpectations(timeout: 10)
+    }
+}
+
+class SleepingDelegate: ServerDelegate {
+    public func handle(request: ServerRequest, response: ServerResponse) {
+        sleep(2)
+        try! response.end()
+    }
+}
+
+class HTTPHandler: ChannelInboundHandler {
+    typealias InboundIn = HTTPClientResponsePart
+
+    private let expectation: XCTestExpectation
+
+    public init(_ expectation: XCTestExpectation) {
+        self.expectation = expectation
+    }
+
+    func channelInactive(context: ChannelHandlerContext) {
+        expectation.fulfill()
+    }
+}

Tests/LinuxMain.swift

@@ -55,4 +55,5 @@ XCTMain([
     testCase(RegressionTests.allTests.shuffled()),
     testCase(MonitoringTests.allTests.shuffled()),
     testCase(KituraNetWebSocketUpgradeTest.allTests.shuffled()),
+    testCase(ChannelQuiescingTests.allTests.shuffled()),
 ].shuffled())