Cherry pick GitLab 187: Graceful connection close without sending QUIT command (#85)

Cherry pick: https://gitlab.com/swift-server-community/RediStack/-/merge_requests/187

> Note: Clients should not use this command. Instead, clients should simply close the connection when they're not used anymore. Terminating a connection on the client side is preferable, as it eliminates TIME_WAIT lingering sockets on the server side.

https://redis.io/commands/quit/

Author: fabianfett

Sources/RediStack/RedisConnection.swift

@@ -170,12 +170,17 @@ public final class RedisConnection: RedisClient, RedisClientWithUserContext {
         self.channel.closeFuture.whenSuccess {
             // if our state is still open, that means we didn't cause the closeFuture to resolve.
             // update state, metrics, and logging
-            guard self.state.isConnected else { return }
-
+            let oldState = self.state
             self.state = .closed
-            self.logger.error("connection was closed unexpectedly")
             RedisMetrics.activeConnectionCount.decrement()
-            self.onUnexpectedClosure?()
+
+            switch oldState {
+            case .shuttingDown, .closed:
+                break
+            case .open, .pubsub:
+                logger.warning("connection was closed unexpectedly")
+                self.onUnexpectedClosure?()
+            }
         }
 
         self.logger.trace("connection created")
@@ -333,62 +338,21 @@ extension RedisConnection {
         // we're now in a shutdown state, starting with the command queue.
         self.state = .shuttingDown
 
-        let notification = self.sendQuitCommand(logger: logger) // send "QUIT" so that all the responses are written out
-            .flatMap { self.closeChannel() } // close the channel from our end
+        // Inform ChannelHandler about close intent using "RedisGracefulConnectionCloseEvent"
+        let closePromise = self.eventLoop.makePromise(of: Void.self)
+        let closeFuture = closePromise.futureResult
+        self.channel.triggerUserOutboundEvent(RedisGracefulConnectionCloseEvent(), promise: closePromise)
 
-        notification.whenFailure {
+        closeFuture.whenFailure {
             logger.error("error while closing connection", metadata: [
                 RedisLogging.MetadataKeys.error: "\($0)"
             ])
         }
-        notification.whenSuccess {
-            self.state = .closed
+        closeFuture.whenSuccess {
             logger.trace("connection is now closed")
-            RedisMetrics.activeConnectionCount.decrement()
         }
 
-        return notification
-    }
-
-    /// Bypasses everything for a normal command and explicitly just sends a "QUIT" command to Redis.
-    /// - Note: If the command fails, the `NIO.EventLoopFuture` will still succeed - as it's not critical for the command to succeed.
-    private func sendQuitCommand(logger: Logger) -> EventLoopFuture<Void> {
-        let promise = channel.eventLoop.makePromise(of: RESPValue.self)
-        let command = RedisCommand(
-            message: .array([RESPValue(bulk: "QUIT")]),
-            responsePromise: promise
-        )
-
-        logger.trace("sending QUIT command")
-
-        return channel.writeAndFlush(command) // write the command
-            .flatMap { promise.futureResult } // chain the callback to the response's
-            .map { _ in logger.trace("sent QUIT command") } // ignore the result's value
-            .recover { _ in logger.debug("recovered from error sending QUIT") } // if there's an error, just return to void
-    }
-
-    /// Attempts to close the `NIO.Channel`.
-    /// SwiftNIO throws a `NIO.EventLoopError.shutdown` if the channel is already closed,
-    /// so that case is captured to let this method's `NIO.EventLoopFuture` still succeed.
-    private func closeChannel() -> EventLoopFuture<Void> {
-        let promise = self.channel.eventLoop.makePromise(of: Void.self)
-
-        self.channel.close(promise: promise)
-
-        // if we succeed, great, if not - check the error that happened
-        return promise.futureResult
-            .flatMapError { error in
-                guard let e = error as? EventLoopError else {
-                    return self.eventLoop.makeFailedFuture(error)
-                }
-
-                // if the error is that the channel is already closed, great - just succeed.
-                // otherwise, fail the chain
-                switch e {
-                case .shutdown: return self.eventLoop.makeSucceededFuture(())
-                default: return self.eventLoop.makeFailedFuture(e)
-                }
-            }
+        return closeFuture
     }
 }
 

Sources/RediStackTestUtils/EmbeddedMockRedisServer.swift

@@ -68,7 +68,7 @@ internal final class EmbeddedMockRedisServer {
     }
 
     func createConnectedChannel() -> Channel {
-        let channel = EmbeddedChannel(loop: self.loop)
+        let channel = EmbeddedChannel(handler: GracefulShutdownToCloseHandler(), loop: self.loop)
         channel.closeFuture.whenComplete { _ in
             self.channels.removeAll(where: { $0 === channel })
         }
@@ -84,3 +84,18 @@ internal final class EmbeddedMockRedisServer {
         try self.loop.close()
     }
 }
+
+/// A `ChannelHandler` that triggers a channel close once `RedisGracefulConnectionCloseEvent` is received
+private final class GracefulShutdownToCloseHandler: ChannelHandler, ChannelOutboundHandler {
+    typealias OutboundIn = NIOAny
+
+    func triggerUserOutboundEvent(context: ChannelHandlerContext, event: Any, promise: EventLoopPromise<Void>?) {
+        switch event {
+        case is RedisGracefulConnectionCloseEvent:
+            context.close(mode: .all, promise: promise)
+
+        default:
+            context.triggerUserOutboundEvent(event, promise: promise)
+        }
+    }
+}