From 7ed140732ef579674529b439b87f5cd7f39bdbc7 Mon Sep 17 00:00:00 2001 From: Nathan Harris Date: Sun, 17 Apr 2022 23:30:44 -0500 Subject: [PATCH] #103 -- Provide error in unsubscribe pub/sub handler if available --- .../ChannelHandlers/RedisPubSubHandler.swift | 47 ++++++++----- Sources/RediStack/RedisClient.swift | 24 +++---- Sources/RediStack/RedisConnection.swift | 12 ++-- Sources/RediStack/RedisConnectionPool.swift | 18 ++--- Sources/RediStack/RedisLogging.swift | 4 +- .../Commands/PubSubCommandsTests.swift | 66 +++++++++++++++++-- 6 files changed, 121 insertions(+), 50 deletions(-) diff --git a/Sources/RediStack/ChannelHandlers/RedisPubSubHandler.swift b/Sources/RediStack/ChannelHandlers/RedisPubSubHandler.swift index 1eb3770..d94a77a 100644 --- a/Sources/RediStack/ChannelHandlers/RedisPubSubHandler.swift +++ b/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 { 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 diff --git a/Sources/RediStack/RedisClient.swift b/Sources/RediStack/RedisClient.swift index a4963a9..9e8344b 100644 --- a/Sources/RediStack/RedisClient.swift +++ b/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 /// 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 /// 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 { 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 { 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 { 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 { return self.psubscribe(to: patterns, eventLoop: eventLoop, logger: logger, messageReceiver: receiver, onSubscribe: subscribeHandler, onUnsubscribe: unsubscribeHandler) } diff --git a/Sources/RediStack/RedisConnection.swift b/Sources/RediStack/RedisConnection.swift index a07f9f4..0f70c21 100644 --- a/Sources/RediStack/RedisConnection.swift +++ b/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 { 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 { 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 { diff --git a/Sources/RediStack/RedisConnectionPool.swift b/Sources/RediStack/RedisConnectionPool.swift index 5a9f9c3..5a7a24b 100644 --- a/Sources/RediStack/RedisConnectionPool.swift +++ b/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 { 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 { return self._subscribe( using: { @@ -433,8 +433,8 @@ extension RedisConnectionPool: RedisClient { } private func _subscribe( - using operation: @escaping (RedisConnection, @escaping RedisSubscriptionChangeHandler, Logger) -> EventLoopFuture, - onUnsubscribe unsubscribeHandler: RedisSubscriptionChangeHandler?, + using operation: @escaping (RedisConnection, @escaping RedisUnsubscribeHandler, Logger) -> EventLoopFuture, + onUnsubscribe unsubscribeHandler: RedisUnsubscribeHandler?, eventLoop: EventLoop?, taskLogger: Logger? ) -> EventLoopFuture { @@ -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 } diff --git a/Sources/RediStack/RedisLogging.swift b/Sources/RediStack/RedisLogging.swift index ef7fa03..d1a1e3e 100644 --- a/Sources/RediStack/RedisLogging.swift +++ b/Sources/RediStack/RedisLogging.swift @@ -109,12 +109,12 @@ internal struct CustomLoggerRedisClient: 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 { + internal func subscribe(to channels: [RedisChannelName], eventLoop: EventLoop?, logger: Logger?, messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver, onSubscribe subscribeHandler: RedisSubscribeHandler?, onUnsubscribe unsubscribeHandler: RedisUnsubscribeHandler?) -> EventLoopFuture { 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 { + internal func psubscribe(to patterns: [String], eventLoop: EventLoop?, logger: Logger?, messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver, onSubscribe subscribeHandler: RedisSubscribeHandler?, onUnsubscribe unsubscribeHandler: RedisUnsubscribeHandler?) -> EventLoopFuture { let logger = logger ?? self.defaultLogger return self.client.psubscribe(to: patterns, eventLoop: eventLoop, logger: logger, messageReceiver: receiver, onSubscribe: subscribeHandler, onUnsubscribe: unsubscribeHandler) } diff --git a/Tests/RediStackIntegrationTests/Commands/PubSubCommandsTests.swift b/Tests/RediStackIntegrationTests/Commands/PubSubCommandsTests.swift index d7f9f27..9159581 100644 --- a/Tests/RediStackIntegrationTests/Commands/PubSubCommandsTests.swift +++ b/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) + } +}