Add support for graceful shutdown to the RedisCommandHandler

Author: fabianfett

Sources/RediStack/ChannelHandlers/RedisCommandHandler.swift

@@ -44,7 +44,9 @@ public final class RedisCommandHandler {
     }
 
     private enum State {
-        case `default`, error(Error)
+        case `default`
+        case draining(EventLoopPromise<Void>?)
+        case error(Error)
     }
 }
 
@@ -70,6 +72,7 @@ extension RedisCommandHandler: ChannelInboundHandler {
     /// See `NIO.ChannelInboundHandler.channelInactive(context:)`
     /// - Note: `RedisMetrics.commandFailureCount` is **not** incremented from this method.
     public func channelInactive(context: ChannelHandlerContext) {
+        self.state = .error(RedisClientError.connectionClosed)
         self._failCommandQueue(because: RedisClientError.connectionClosed)
     }
 
@@ -100,6 +103,16 @@ extension RedisCommandHandler: ChannelInboundHandler {
             leadPromise.succeed(value)
             RedisMetrics.commandSuccessCount.increment()
         }
+
+        switch self.state {
+        case .draining(let promise):
+            if self.commandResponseQueue.isEmpty {
+                context.close(mode: .all, promise: promise)
+            }
+
+        case .error, .`default`:
+            break
+        }
     }
 }
 
@@ -119,7 +132,11 @@ extension RedisCommandHandler: ChannelOutboundHandler {
         let commandPayload = self.unwrapOutboundIn(data)
 
         switch self.state {
-        case let .error(e): commandPayload.responsePromise.fail(e)
+        case let .error(e):
+            commandPayload.responsePromise.fail(e)
+
+        case .draining:
+            commandPayload.responsePromise.fail(RedisClientError.connectionClosed)
 
         case .default:
             self.commandResponseQueue.append(commandPayload.responsePromise)
@@ -129,4 +146,39 @@ extension RedisCommandHandler: ChannelOutboundHandler {
             )
         }
     }
+
+    /// Listens for ``RedisGracefulConnectionCloseEvent``. If such an event is received the handler will wait
+    /// until all currently running commands have returned. Once all requests are fulfilled the handler will close the channel.
+    ///
+    /// If a command is sent on the channel, after the ``RedisGracefulConnectionCloseEvent`` was scheduled,
+    /// the command will be failed with a ``RedisClientError/connectionClosed``.
+    ///
+    /// See `NIO.ChannelOutboundHandler.triggerUserOutboundEvent(context:event:promise:)`
+    public func triggerUserOutboundEvent(context: ChannelHandlerContext, event: Any, promise: EventLoopPromise<Void>?) {
+        switch event {
+        case is RedisGracefulConnectionCloseEvent:
+            switch self.state {
+            case .default:
+                if self.commandResponseQueue.isEmpty {
+                    self.state = .error(RedisClientError.connectionClosed)
+                    context.close(mode: .all, promise: promise)
+                } else {
+                    self.state = .draining(promise)
+                }
+
+            case .error, .draining:
+                promise?.succeed(())
+                break
+            }
+
+        default:
+            context.triggerUserOutboundEvent(event, promise: promise)
+        }
+    }
+}
+
+/// A channel event that informs the ``RedisCommandHandler`` that it should close the channel gracefully
+public struct RedisGracefulConnectionCloseEvent {
+    /// Creates a ``RedisGracefulConnectionCloseEvent``
+    public init() {}
 }

Tests/RediStackTests/ChannelHandlers/RedisCommandHandlerTests.swift

@@ -14,6 +14,8 @@
 
 import NIOCore
 import NIOPosix
+import NIOEmbedded
+import Atomics
 @testable import RediStack
 import XCTest
 
@@ -44,6 +46,83 @@ final class RedisCommandHandlerTests: XCTestCase {
             XCTAssertEqual(error, .connectionClosed)
         }
     }
+
+    func testCloseIsTriggeredOnceCommandQueueIsEmpty() {
+        let loop = EmbeddedEventLoop()
+        let channel = EmbeddedChannel(handler: RedisCommandHandler(), loop: loop)
+
+        XCTAssertNoThrow(try channel.connect(to: .init(unixDomainSocketPath: "/foo")).wait())
+        XCTAssertTrue(channel.isActive)
+
+        let getFoo = RESPValue.array([.bulkString(.init(string: "GET")), .bulkString(.init(string: "foo"))])
+        let promiseFoo = loop.makePromise(of: RESPValue.self)
+        let commandFoo = (message: getFoo, responsePromise: promiseFoo)
+        XCTAssertNoThrow(try channel.writeOutbound(commandFoo))
+        XCTAssertEqual(try channel.readOutbound(as: RESPValue.self), getFoo)
+
+        let getBar = RESPValue.array([.bulkString(.init(string: "GET")), .bulkString(.init(string: "bar"))])
+        let promiseBar = loop.makePromise(of: RESPValue.self)
+        let commandBar = (message: getBar, responsePromise: promiseBar)
+        XCTAssertNoThrow(try channel.writeOutbound(commandBar))
+        XCTAssertEqual(try channel.readOutbound(as: RESPValue.self), getBar)
+
+        let getBaz = RESPValue.array([.bulkString(.init(string: "GET")), .bulkString(.init(string: "baz"))])
+        let promiseBaz = loop.makePromise(of: RESPValue.self)
+        let commandBaz = (message: getBaz, responsePromise: promiseBaz)
+        XCTAssertNoThrow(try channel.writeOutbound(commandBaz))
+        XCTAssertEqual(try channel.readOutbound(as: RESPValue.self), getBaz)
+
+        let gracefulClosePromise = loop.makePromise(of: Void.self)
+        let channelCloseHitCounter = ManagedAtomic<Int>(0)
+        gracefulClosePromise.futureResult.whenComplete { _ in
+            channelCloseHitCounter.wrappingIncrement(ordering: .relaxed)
+        }
+        channel.triggerUserOutboundEvent(RedisGracefulConnectionCloseEvent(), promise: gracefulClosePromise)
+        XCTAssertEqual(channelCloseHitCounter.load(ordering: .relaxed), 0)
+
+        let fooResponse = RESPValue.simpleString(.init(string: "fooresult"))
+        XCTAssertNoThrow(try channel.writeInbound(fooResponse))
+        XCTAssertTrue(channel.isActive)
+        XCTAssertEqual(channelCloseHitCounter.load(ordering: .relaxed), 0)
+        XCTAssertEqual(try promiseFoo.futureResult.wait(), fooResponse)
+
+        let barResponse = RESPValue.simpleString(.init(string: "barresult"))
+        XCTAssertNoThrow(try channel.writeInbound(barResponse))
+        XCTAssertTrue(channel.isActive)
+        XCTAssertEqual(channelCloseHitCounter.load(ordering: .relaxed), 0)
+        XCTAssertEqual(try promiseBar.futureResult.wait(), barResponse)
+
+        let bazResponse = RESPValue.simpleString(.init(string: "bazresult"))
+        XCTAssertNoThrow(try channel.writeInbound(bazResponse))
+        XCTAssertEqual(try promiseBaz.futureResult.wait(), bazResponse)
+        XCTAssertFalse(channel.isActive)
+        XCTAssertEqual(channelCloseHitCounter.load(ordering: .relaxed), 1)
+        XCTAssertNoThrow(try gracefulClosePromise.futureResult.wait())
+    }
+
+    func testCloseIsTriggeredRightAwayIfCommandQueueIsEmpty() {
+        let loop = EmbeddedEventLoop()
+        let channel = EmbeddedChannel(handler: RedisCommandHandler(), loop: loop)
+        XCTAssertNoThrow(try channel.connect(to: .init(unixDomainSocketPath: "/foo")).wait())
+        XCTAssertTrue(channel.isActive)
+
+        let gracefulClosePromise = loop.makePromise(of: Void.self)
+        let gracefulCloseHitCounter = ManagedAtomic<Int>(0)
+        gracefulClosePromise.futureResult.whenComplete { _ in
+            gracefulCloseHitCounter.wrappingIncrement(ordering: .relaxed)
+        }
+        channel.triggerUserOutboundEvent(RedisGracefulConnectionCloseEvent(), promise: gracefulClosePromise)
+        XCTAssertFalse(channel.isActive)
+        XCTAssertEqual(gracefulCloseHitCounter.load(ordering: .relaxed), 1)
+
+        let getBar = RESPValue.array([.bulkString(.init(string: "GET")), .bulkString(.init(string: "bar"))])
+        let promiseBar = loop.makePromise(of: RESPValue.self)
+        let commandBar = (message: getBar, responsePromise: promiseBar)
+        channel.write(commandBar, promise: nil)
+        XCTAssertThrowsError(try promiseBar.futureResult.wait()) {
+            XCTAssertEqual($0 as? RedisClientError, .connectionClosed)
+        }
+    }
 }
 
 private final class RemoteCloseHandler: ChannelInboundHandler {