#103 -- Provide error in unsubscribe pub/sub handler if available

Sources/RediStack/ChannelHandlers/RedisPubSubHandler.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the RediStack open source project
 //
-// Copyright (c) 2020 RediStack project authors
+// Copyright (c) 2020-2022 RediStack project authors
 // Licensed under Apache License v2.0
 //
 // See LICENSE.txt for license information
@@ -24,18 +24,35 @@ import NIO
 ///     - message: The message data that was received from the `publisher`.
 public typealias RedisSubscriptionMessageReceiver = (_ publisher: RedisChannelName, _ message: RESPValue) -> Void
 
-/// A closure handler invoked for Pub/Sub subscription changes.
+/// The details of the subscription change.
+/// - Parameters:
+///     - subscriptionKey: The subscribed channel or pattern that had its subscription status changed.
+///     - currentSubscriptionCount: The current total number of subscriptions the connection has.
+public typealias RedisSubscriptionChangeDetails = (subscriptionKey: String, currentSubscriptionCount: Int)
+
+/// A closure handler invoked for Pub/Sub subscribe commands.
 ///
 /// This closure will be invoked only *once* for each individual channel or pattern that is having its subscription changed,
-/// even if it was done as a single PSUBSCRIBE, SUBSCRIBE, PUNSUBSCRIBE, or UNSUBSCRIBE command.
+/// even if it was done as a single PSUBSCRIBE or SUBSCRIBE command.
+/// - Warning: The receiver is called on the same `NIO.EventLoop` that processed the message.
+///
+///     If you are doing non-trivial work in response to PubSub messages, it is **highly recommended** that the work be dispatched to another thread
+///     so as to not block further messages from being processed.
+/// - Parameter details: The details of the subscription.
+public typealias RedisSubscribeHandler = (_ details: RedisSubscriptionChangeDetails) -> Void
+
+/// A closure handler invoked for Pub/Sub unsubscribe commands.
+///
+/// This closure will be invoked only *once* for each individual channel or pattern that is having its subscription changed,
+/// even if it was done as a single PUNSUBSCRIBE or UNSUBSCRIBE command.
 /// - Warning: The receiver is called on the same `NIO.EventLoop` that processed the message.
 ///
 ///     If you are doing non-trivial work in response to PubSub messages, it is **highly recommended** that the work be dispatched to another thread
 ///     so as to not block further messages from being processed.
 /// - Parameters:
-///     - subscriptionKey: The subscribed channel or pattern that had its subscription status changed.
-///     - currentSubscriptionCount: The current total number of subscriptions the connection has.
-public typealias RedisSubscriptionChangeHandler = (_ subscriptionKey: String, _ currentSubscriptionCount: Int) -> Void
+///     - details: The details of the subscription.
+///     - error: The error triggering the unsubscribe, if any. If this is `nil`, then the subscription was expected as a response to a user-initiated unsubscribe.
+public typealias RedisUnsubscribeHandler = (_ details: RedisSubscriptionChangeDetails, _ error: Error?) -> Void
 
 /// A list of patterns or channels that a Pub/Sub subscription change is targetting.
 ///
@@ -146,7 +163,7 @@ extension RedisPubSubHandler {
         
         guard let subscription = self.subscriptions[prefixedKey] else { return }
 
-        subscription.onSubscribe?(subscriptionKey, subscriptionCount)
+        subscription.onSubscribe?((subscriptionKey, subscriptionCount))
         subscription.onSubscribe = nil // nil to free memory
         self.subscriptions[prefixedKey] = subscription
         
@@ -162,7 +179,7 @@ extension RedisPubSubHandler {
         let prefixedKey = self.prefixKey(subscriptionKey, with: keyPrefix)
         guard let subscription = self.subscriptions.removeValue(forKey: prefixedKey) else { return }
     
-        subscription.onUnsubscribe?(subscriptionKey, subscriptionCount)
+        subscription.onUnsubscribe?((subscriptionKey, subscriptionCount), nil)
         subscription.type.gauge.decrement()
 
         switch self.pendingUnsubscribes.removeValue(forKey: prefixedKey) {
@@ -208,8 +225,8 @@ extension RedisPubSubHandler {
     public func addSubscription(
         for target: RedisSubscriptionTarget,
         messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver,
-        onSubscribe subscribeHandler: RedisSubscriptionChangeHandler?,
-        onUnsubscribe unsubscribeHandler: RedisSubscriptionChangeHandler?
+        onSubscribe subscribeHandler: RedisSubscribeHandler?,
+        onUnsubscribe unsubscribeHandler: RedisUnsubscribeHandler?
     ) -> EventLoopFuture<Int> {
         guard self.eventLoop.inEventLoop else {
             return self.eventLoop.flatSubmit {
@@ -481,7 +498,7 @@ extension RedisPubSubHandler: ChannelInboundHandler {
         let receivers = self.subscriptions
         self.subscriptions.removeAll()
         receivers.forEach {
-            $0.value.onUnsubscribe?($0.key, 0)
+            $0.value.onUnsubscribe?(($0.key, 0), error)
             $0.value.type.gauge.decrement()
         }
     }
@@ -521,14 +538,14 @@ extension RedisPubSubHandler {
     fileprivate final class Subscription {
         let type: SubscriptionType
         let onMessage: RedisSubscriptionMessageReceiver
-        var onSubscribe: RedisSubscriptionChangeHandler? // will be set to nil after first call
-        let onUnsubscribe: RedisSubscriptionChangeHandler?
+        var onSubscribe: RedisSubscribeHandler? // will be set to nil after first call
+        let onUnsubscribe: RedisUnsubscribeHandler?
         
         init(
             type: SubscriptionType,
             messageReceiver: @escaping RedisSubscriptionMessageReceiver,
-            subscribeHandler: RedisSubscriptionChangeHandler?,
-            unsubscribeHandler: RedisSubscriptionChangeHandler?
+            subscribeHandler: RedisSubscribeHandler?,
+            unsubscribeHandler: RedisUnsubscribeHandler?
         ) {
             self.type = type
             self.onMessage = messageReceiver

Sources/RediStack/RedisClient.swift

@@ -71,8 +71,8 @@ public protocol RedisClient {
         eventLoop: EventLoop?,
         logger: Logger?,
         messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver,
-        onSubscribe subscribeHandler: RedisSubscriptionChangeHandler?,
-        onUnsubscribe unsubscribeHandler: RedisSubscriptionChangeHandler?
+        onSubscribe subscribeHandler: RedisSubscribeHandler?,
+        onUnsubscribe unsubscribeHandler: RedisUnsubscribeHandler?
     ) -> EventLoopFuture<Void>
 
     /// Subscribes the client to the specified Redis channel name patterns, invoking the provided message receiver each time a message is published to
@@ -100,8 +100,8 @@ public protocol RedisClient {
         eventLoop: EventLoop?,
         logger: Logger?,
         messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver,
-        onSubscribe subscribeHandler: RedisSubscriptionChangeHandler?,
-        onUnsubscribe unsubscribeHandler: RedisSubscriptionChangeHandler?
+        onSubscribe subscribeHandler: RedisSubscribeHandler?,
+        onUnsubscribe unsubscribeHandler: RedisUnsubscribeHandler?
     ) -> EventLoopFuture<Void>
     
     /// Unsubscribes the client from a specific Redis channel from receiving any future published messages.
@@ -194,8 +194,8 @@ extension RedisClient {
         eventLoop: EventLoop? = nil,
         logger: Logger? = nil,
         messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver,
-        onSubscribe subscribeHandler: RedisSubscriptionChangeHandler? = nil,
-        onUnsubscribe unsubscribeHandler: RedisSubscriptionChangeHandler? = nil
+        onSubscribe subscribeHandler: RedisSubscribeHandler? = nil,
+        onUnsubscribe unsubscribeHandler: RedisUnsubscribeHandler? = nil
     ) -> EventLoopFuture<Void> {
         return self.subscribe(to: channels, eventLoop: eventLoop, logger: logger, messageReceiver: receiver, onSubscribe: subscribeHandler, onUnsubscribe: unsubscribeHandler)
     }
@@ -205,8 +205,8 @@ extension RedisClient {
         eventLoop: EventLoop? = nil,
         logger: Logger? = nil,
         messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver,
-        onSubscribe subscribeHandler: RedisSubscriptionChangeHandler? = nil,
-        onUnsubscribe unsubscribeHandler: RedisSubscriptionChangeHandler? = nil
+        onSubscribe subscribeHandler: RedisSubscribeHandler? = nil,
+        onUnsubscribe unsubscribeHandler: RedisUnsubscribeHandler? = nil
     ) -> EventLoopFuture<Void> {
         return self.subscribe(to: channels, eventLoop: eventLoop, logger: logger, messageReceiver: receiver, onSubscribe: subscribeHandler, onUnsubscribe: unsubscribeHandler)
     }
@@ -216,8 +216,8 @@ extension RedisClient {
         eventLoop: EventLoop? = nil,
         logger: Logger? = nil,
         messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver,
-        onSubscribe subscribeHandler: RedisSubscriptionChangeHandler? = nil,
-        onUnsubscribe unsubscribeHandler: RedisSubscriptionChangeHandler? = nil
+        onSubscribe subscribeHandler: RedisSubscribeHandler? = nil,
+        onUnsubscribe unsubscribeHandler: RedisUnsubscribeHandler? = nil
     ) -> EventLoopFuture<Void> {
         return self.psubscribe(to: patterns, eventLoop: eventLoop, logger: logger, messageReceiver: receiver, onSubscribe: subscribeHandler, onUnsubscribe: unsubscribeHandler)
     }
@@ -227,8 +227,8 @@ extension RedisClient {
         eventLoop: EventLoop? = nil,
         logger: Logger? = nil,
         messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver,
-        onSubscribe subscribeHandler: RedisSubscriptionChangeHandler? = nil,
-        onUnsubscribe unsubscribeHandler: RedisSubscriptionChangeHandler? = nil
+        onSubscribe subscribeHandler: RedisSubscribeHandler? = nil,
+        onUnsubscribe unsubscribeHandler: RedisUnsubscribeHandler? = nil
     ) -> EventLoopFuture<Void> {
         return self.psubscribe(to: patterns, eventLoop: eventLoop, logger: logger, messageReceiver: receiver, onSubscribe: subscribeHandler, onUnsubscribe: unsubscribeHandler)
     }

Sources/RediStack/RedisConnection.swift

@@ -381,8 +381,8 @@ extension RedisConnection {
         eventLoop: EventLoop? = nil,
         logger: Logger? = nil,
         messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver,
-        onSubscribe subscribeHandler: RedisSubscriptionChangeHandler?,
-        onUnsubscribe unsubscribeHandler: RedisSubscriptionChangeHandler?
+        onSubscribe subscribeHandler: RedisSubscribeHandler?,
+        onUnsubscribe unsubscribeHandler: RedisUnsubscribeHandler?
     ) -> EventLoopFuture<Void> {
         return self._subscribe(.channels(channels), receiver, subscribeHandler, unsubscribeHandler, eventLoop, logger)
     }
@@ -392,8 +392,8 @@ extension RedisConnection {
         eventLoop: EventLoop? = nil,
         logger: Logger? = nil,
         messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver,
-        onSubscribe subscribeHandler: RedisSubscriptionChangeHandler? = nil,
-        onUnsubscribe unsubscribeHandler: RedisSubscriptionChangeHandler? = nil
+        onSubscribe subscribeHandler: RedisSubscribeHandler? = nil,
+        onUnsubscribe unsubscribeHandler: RedisUnsubscribeHandler? = nil
     ) -> EventLoopFuture<Void> {
         return self._subscribe(.patterns(patterns), receiver, subscribeHandler, unsubscribeHandler, eventLoop, logger)
     }
@@ -401,8 +401,8 @@ extension RedisConnection {
     private func _subscribe(
         _ target: RedisSubscriptionTarget,
         _ receiver: @escaping RedisSubscriptionMessageReceiver,
-        _ onSubscribe: RedisSubscriptionChangeHandler?,
-        _ onUnsubscribe: RedisSubscriptionChangeHandler?,
+        _ onSubscribe: RedisSubscribeHandler?,
+        _ onUnsubscribe: RedisUnsubscribeHandler?,
         _ eventLoop: EventLoop?,
         _ logger: Logger?
     ) -> EventLoopFuture<Void> {

Sources/RediStack/RedisConnectionPool.swift

@@ -363,8 +363,8 @@ extension RedisConnectionPool: RedisClient {
         eventLoop: EventLoop? = nil,
         logger: Logger? = nil,
         messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver,
-        onSubscribe subscribeHandler: RedisSubscriptionChangeHandler?,
-        onUnsubscribe unsubscribeHandler: RedisSubscriptionChangeHandler?
+        onSubscribe subscribeHandler: RedisSubscribeHandler?,
+        onUnsubscribe unsubscribeHandler: RedisUnsubscribeHandler?
     ) -> EventLoopFuture<Void> {
         return self._subscribe(
             using: {
@@ -388,8 +388,8 @@ extension RedisConnectionPool: RedisClient {
         eventLoop: EventLoop? = nil,
         logger: Logger? = nil,
         messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver,
-        onSubscribe subscribeHandler: RedisSubscriptionChangeHandler?,
-        onUnsubscribe unsubscribeHandler: RedisSubscriptionChangeHandler?
+        onSubscribe subscribeHandler: RedisSubscribeHandler?,
+        onUnsubscribe unsubscribeHandler: RedisUnsubscribeHandler?
     ) -> EventLoopFuture<Void> {
         return self._subscribe(
             using: {
@@ -433,8 +433,8 @@ extension RedisConnectionPool: RedisClient {
     }
 
     private func _subscribe(
-        using operation: @escaping (RedisConnection, @escaping RedisSubscriptionChangeHandler, Logger) -> EventLoopFuture<Void>,
-        onUnsubscribe unsubscribeHandler: RedisSubscriptionChangeHandler?,
+        using operation: @escaping (RedisConnection, @escaping RedisUnsubscribeHandler, Logger) -> EventLoopFuture<Void>,
+        onUnsubscribe unsubscribeHandler: RedisUnsubscribeHandler?,
         eventLoop: EventLoop?,
         taskLogger: Logger?
     ) -> EventLoopFuture<Void> {
@@ -446,11 +446,11 @@ extension RedisConnectionPool: RedisClient {
                     self.pubsubConnection = connection
                 }
                 
-                let onUnsubscribe: RedisSubscriptionChangeHandler = { channelName, subCount in
-                    defer { unsubscribeHandler?(channelName, subCount) }
+                let onUnsubscribe: RedisUnsubscribeHandler = { subscriptionDetails, error in
+                    defer { unsubscribeHandler?(subscriptionDetails, error) }
                     
                     guard
-                        subCount == 0,
+                        subscriptionDetails.currentSubscriptionCount == 0,
                         let connection = self.pubsubConnection
                     else { return }
 

Sources/RediStack/RedisLogging.swift

@@ -109,12 +109,12 @@ internal struct CustomLoggerRedisClient<Client: RedisClient>: RedisClient {
         return self.client.punsubscribe(from: patterns, eventLoop: eventLoop, logger: logger)
     }
 
-    internal func subscribe(to channels: [RedisChannelName], eventLoop: EventLoop?, logger: Logger?, messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver, onSubscribe subscribeHandler: RedisSubscriptionChangeHandler?, onUnsubscribe unsubscribeHandler: RedisSubscriptionChangeHandler?) -> EventLoopFuture<Void> {
+    internal func subscribe(to channels: [RedisChannelName], eventLoop: EventLoop?, logger: Logger?, messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver, onSubscribe subscribeHandler: RedisSubscribeHandler?, onUnsubscribe unsubscribeHandler: RedisUnsubscribeHandler?) -> EventLoopFuture<Void> {
         let logger = logger ?? self.defaultLogger
         return self.client.subscribe(to: channels, eventLoop: eventLoop, logger: logger, messageReceiver: receiver, onSubscribe: subscribeHandler, onUnsubscribe: unsubscribeHandler)
     }
 
-    internal func psubscribe(to patterns: [String], eventLoop: EventLoop?, logger: Logger?, messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver, onSubscribe subscribeHandler: RedisSubscriptionChangeHandler?, onUnsubscribe unsubscribeHandler: RedisSubscriptionChangeHandler?) -> EventLoopFuture<Void> {
+    internal func psubscribe(to patterns: [String], eventLoop: EventLoop?, logger: Logger?, messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver, onSubscribe subscribeHandler: RedisSubscribeHandler?, onUnsubscribe unsubscribeHandler: RedisUnsubscribeHandler?) -> EventLoopFuture<Void> {
         let logger = logger ?? self.defaultLogger
         return self.client.psubscribe(to: patterns, eventLoop: eventLoop, logger: logger, messageReceiver: receiver, onSubscribe: subscribeHandler, onUnsubscribe: unsubscribeHandler)
     }

Tests/RediStackIntegrationTests/Commands/PubSubCommandsTests.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the RediStack open source project
 //
-// Copyright (c) 2020 RediStack project authors
+// Copyright (c) 2020-2022 RediStack project authors
 // Licensed under Apache License v2.0
 //
 // See LICENSE.txt for license information
@@ -12,7 +12,9 @@
 //
 //===----------------------------------------------------------------------===//
 
-import RediStack
+import NIO
+import NIOEmbedded
+@testable import RediStack
 import RediStackTestUtils
 import XCTest
 
@@ -40,8 +42,12 @@ final class RedisPubSubCommandsTests: RediStackIntegrationTestCase {
                 guard $0 == #function, $1 == 1 else { return }
                 subscribeExpectation.fulfill()
             },
-            onUnsubscribe: {
-                guard $0 == #function, $1 == 0 else { return }
+            onUnsubscribe: { details, error in
+                guard
+                    error == nil,
+                    details.subscriptionKey == #function,
+                    details.currentSubscriptionCount == 0
+                else { return }
                 unsubscribeExpectation.fulfill()
             }
         ).wait()
@@ -294,8 +300,12 @@ final class RedisPubSubCommandsPoolTests: RediStackConnectionPoolIntegrationTest
                 guard $0 == #function, $1 == 1 else { return }
                 subscribeExpectation.fulfill()
             },
-            onUnsubscribe: {
-                guard $0 == #function, $1 == 0 else { return }
+            onUnsubscribe: { details, error in
+                guard
+                    error == nil,
+                    details.subscriptionKey == #function,
+                    details.currentSubscriptionCount == 0
+                else { return }
                 unsubscribeExpectation.fulfill()
             }
         ).wait()
@@ -341,3 +351,47 @@ final class RedisPubSubCommandsPoolTests: RediStackConnectionPoolIntegrationTest
         XCTAssertEqual(self.pool.leasedConnectionCount, 0)
     }
 }
+
+// MARK: - #103 tests
+
+extension RedisPubSubCommandsTests {
+    func test_pubsub_calls_unsubscribe_whenUnexpectedClose() throws {
+        let channel = EmbeddedChannel()
+        try channel
+            .addBaseRedisHandlers()
+            .wait()
+
+        let subscribeExpectation = self.expectation(description: "should see subscribe")
+        let unsubscribeExpectation = self.expectation(description: "should see unsubscribe")
+
+        let connection = RedisConnection(configuredRESPChannel: channel, defaultLogger: .init(label: ""))
+        let subscribeFuture = connection
+            .subscribe(
+                to: [.init(#function)],
+                messageReceiver: { _, _ in },
+                onSubscribe: { _, _ in subscribeExpectation.fulfill() },
+                onUnsubscribe: { _, error in
+                    guard error != nil else { return }
+                    unsubscribeExpectation.fulfill()
+                }
+            )
+
+        // mimics a successful subscription response from the server
+        let allocator = ByteBufferAllocator()
+        var buffer = allocator.buffer(capacity: 300)
+        buffer.writeRESPValue(.array([
+            .init(bulk: "subscribe"),
+            .init(bulk: "\(#function)"),
+            .integer(1)
+        ]))
+        try channel.writeInbound(buffer)
+
+        // lets the initial subscription work finish
+        try subscribeFuture.wait()
+
+        // 'unexpected' close, should trigger expectations
+        try channel.close().wait()
+
+        self.waitForExpectations(timeout: 0.5)
+    }
+}