Add onUnexpectedConnectionClose callback to pool

Backports `onUnexpectedConnectionClose` on the pool to 1.x

Author: marius-se

Sources/RediStack/Configuration.swift

@@ -27,20 +27,20 @@ extension RedisConnection.Configuration {
         var localizedDescription: String { self.kind.localizedDescription }
 
         private let kind: Kind
-        
+
         private init(_ kind: Kind) { self.kind = kind }
-        
+
         public static func ==(lhs: ValidationError, rhs: ValidationError) -> Bool {
             return lhs.kind == rhs.kind
         }
-        
+
         private enum Kind: LocalizedError {
             case invalidURLString
             case missingURLScheme
             case invalidURLScheme
             case missingURLHost
             case outOfBoundsDatabaseID
-            
+
             var localizedDescription: String {
                 let message: String = {
                     switch self {
@@ -66,7 +66,7 @@ extension RedisConnection {
         ///
         /// See [https://redis.io/topics/quickstart](https://redis.io/topics/quickstart)
         public static var defaultPort = 6379
-        
+
         internal static let defaultLogger = Logger.redisBaseConnectionLogger
 
         /// The hostname of the connection address. If the address is a Unix socket, then it will be `nil`.
@@ -85,9 +85,9 @@ extension RedisConnection {
         public let initialDatabase: Int?
         /// The logger prototype that will be used by the connection by default when generating logs.
         public let defaultLogger: Logger
-        
+
         internal let address: SocketAddress
-        
+
         /// Creates a new connection configuration with the provided details.
         /// - Parameters:
         ///     - address: The socket address information to use for creating the Redis connection.
@@ -106,7 +106,7 @@ extension RedisConnection {
             if initialDatabase != nil && initialDatabase! < 0 {
                 throw ValidationError.outOfBoundsDatabaseID
             }
-            
+
             self.address = address
             self.password = password
             self.initialDatabase = initialDatabase
@@ -182,9 +182,9 @@ extension RedisConnection {
             try Self.validateRedisURL(url)
 
             guard let host = url.host, !host.isEmpty else { throw ValidationError.missingURLHost }
-            
+
             let databaseID = Int(url.lastPathComponent)
-            
+
             try self.init(
                 address: try .makeAddressResolvingHost(host, port: url.port ?? Self.defaultPort),
                 password: url.password,
@@ -219,7 +219,7 @@ extension RedisConnectionPool {
         public let connectionInitialDatabase: Int?
         /// The pre-configured TCP client for connections to use.
         public let tcpClient: ClientBootstrap?
-    
+
         /// Creates a new connection factory configuration with the provided options.
         /// - Parameters:
         ///     - connectionInitialDatabase: The optional database index to initially connect to. The default is `nil`.
@@ -255,11 +255,13 @@ extension RedisConnectionPool {
         public let maximumConnectionCount: RedisConnectionPoolSize
         /// The configuration object that controls the connection retry behavior.
         public let connectionRetryConfiguration: (backoff: (initialDelay: TimeAmount, factor: Float32), timeout: TimeAmount)
+        /// Called when a connection in the pool is closed unexpectedly.
+        public let onUnexpectedConnectionClose: ((RedisConnection) -> Void)?
         // these need to be var so they can be updated by the pool in some cases
         public internal(set) var factoryConfiguration: ConnectionFactoryConfiguration
         /// The logger prototype that will be used by the connection pool by default when generating logs.
         public internal(set) var poolDefaultLogger: Logger
-        
+
         /// Creates a new connection configuration with the provided options.
         /// - Parameters:
         ///     - initialServerConnectionAddresses: The set of Redis servers to which this pool is initially willing to connect.
@@ -284,6 +286,7 @@ extension RedisConnectionPool {
             connectionBackoffFactor: Float32 = 2,
             initialConnectionBackoffDelay: TimeAmount = .milliseconds(100),
             connectionRetryTimeout: TimeAmount? = .seconds(60),
+            onUnexpectedConnectionClose: ((RedisConnection) -> Void)? = nil,
             poolDefaultLogger: Logger? = nil
         ) {
             self.initialConnectionAddresses = initialServerConnectionAddresses
@@ -294,6 +297,7 @@ extension RedisConnectionPool {
                 (initialConnectionBackoffDelay, connectionBackoffFactor),
                 connectionRetryTimeout ?? .milliseconds(10) // always default to a baseline 10ms
             )
+            self.onUnexpectedConnectionClose = onUnexpectedConnectionClose
             self.poolDefaultLogger = poolDefaultLogger ?? .redisBaseConnectionPoolLogger
         }
     }

Sources/RediStack/RedisConnection.swift

@@ -21,7 +21,7 @@ import NIO
 import NIOConcurrencyHelpers
 
 extension RedisConnection {
-    
+
     /// Creates a new connection with provided configuration and sychronization objects.
     ///
     /// If you would like to specialize the `NIO.ClientBootstrap` that the connection communicates on, override the default by passing it in as `configuredTCPClient`.
@@ -55,7 +55,7 @@ extension RedisConnection {
         configuredTCPClient client: ClientBootstrap? = nil
     ) -> EventLoopFuture<RedisConnection> {
         let client = client ?? .makeRedisTCPClient(group: eventLoop)
-        
+
         var future = client
             .connect(to: config.address)
             .map { return RedisConnection(configuredRESPChannel: $0, context: config.defaultLogger) }
@@ -73,7 +73,7 @@ extension RedisConnection {
                 return connection.select(database: database).map { connection }
             }
         }
-        
+
         return future
     }
 }
@@ -157,14 +157,14 @@ public final class RedisConnection: RedisClient, RedisClientWithUserContext {
         get { return _stateLock.withLock { self._state } }
         set(newValue) { _stateLock.withLockVoid { self._state = newValue } }
     }
-    
+
     deinit {
         if isConnected {
             assertionFailure("close() was not called before deinit!")
             self.logger.warning("connection was not properly shutdown before deinit")
         }
     }
-    
+
     internal init(configuredRESPChannel: Channel, context: Context) {
         self.channel = configuredRESPChannel
         // there is a mix of verbiage here as the API is forward thinking towards "baggage context"
@@ -176,14 +176,14 @@ public final class RedisConnection: RedisClient, RedisClientWithUserContext {
 
         RedisMetrics.activeConnectionCount.increment()
         RedisMetrics.totalConnectionCount.increment()
-        
+
         // attach a callback to the channel to capture situations where the channel might be closed out from under
         // the connection
         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 }
-            
+
             self.state = .closed
             self.logger.error("connection was closed unexpectedly")
             RedisMetrics.activeConnectionCount.decrement()
@@ -192,13 +192,13 @@ public final class RedisConnection: RedisClient, RedisClientWithUserContext {
 
         self.logger.trace("connection created")
     }
-    
+
     internal enum ConnectionState {
         case open
         case pubsub(RedisPubSubHandler)
         case shuttingDown
         case closed
-        
+
         var isConnected: Bool {
             switch self {
             case .open, .pubsub: return true
@@ -242,40 +242,40 @@ extension RedisConnection {
             return self.channel.eventLoop.makeFailedFuture(error)
         }
         logger.trace("received command request")
-        
+
         logger.debug("sending command", metadata: [
             RedisLogging.MetadataKeys.commandKeyword: "\(command)",
             RedisLogging.MetadataKeys.commandArguments: "\(arguments)"
         ])
-        
+
         var message: [RESPValue] = [.init(bulk: command)]
         message.append(contentsOf: arguments)
-        
+
         let promise = channel.eventLoop.makePromise(of: RESPValue.self)
         let command = RedisCommand(
             message: .array(message),
             responsePromise: promise
         )
-        
+
         let startTime = DispatchTime.now().uptimeNanoseconds
         promise.futureResult.whenComplete { result in
             let duration = DispatchTime.now().uptimeNanoseconds - startTime
             RedisMetrics.commandRoundTripTime.recordNanoseconds(duration)
-            
+
             // log data based on the result
             switch result {
             case let .failure(error):
                 logger.error("command failed", metadata: [
                     RedisLogging.MetadataKeys.error: "\(error.localizedDescription)"
                 ])
-                
+
             case let .success(value):
                 logger.debug("command succeeded", metadata: [
                     RedisLogging.MetadataKeys.commandResult: "\(value)"
                 ])
             }
         }
-        
+
         defer { logger.trace("command sent") }
 
         if self.sendCommandsImmediately {
@@ -310,10 +310,10 @@ 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
-        
+
         notification.whenFailure {
             logger.error("error while closing connection", metadata: [
                 RedisLogging.MetadataKeys.error: "\($0)"
@@ -324,10 +324,10 @@ extension RedisConnection {
             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> {
@@ -344,22 +344,22 @@ extension RedisConnection {
             .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 {
@@ -395,7 +395,7 @@ extension RedisConnection {
     ) -> EventLoopFuture<Void> {
         return self._subscribe(.channels(channels), receiver, subscribeHandler, unsubscribeHandler, nil)
     }
-    
+
     public func psubscribe(
         to patterns: [String],
         messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver,
@@ -404,7 +404,7 @@ extension RedisConnection {
     ) -> EventLoopFuture<Void> {
         return self._subscribe(.patterns(patterns), receiver, subscribeHandler, unsubscribeHandler, nil)
     }
-    
+
     internal func subscribe(
         to channels: [RedisChannelName],
         messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver,
@@ -414,7 +414,7 @@ extension RedisConnection {
     ) -> EventLoopFuture<Void> {
         return self._subscribe(.channels(channels), receiver, subscribeHandler, unsubscribeHandler, context)
     }
-    
+
     internal func psubscribe(
         to patterns: [String],
         messageReceiver receiver: @escaping RedisSubscriptionMessageReceiver,
@@ -424,7 +424,7 @@ extension RedisConnection {
     ) -> EventLoopFuture<Void> {
         return self._subscribe(.patterns(patterns), receiver, subscribeHandler, unsubscribeHandler, context)
     }
-    
+
     private func _subscribe(
         _ target: RedisSubscriptionTarget,
         _ receiver: @escaping RedisSubscriptionMessageReceiver,
@@ -433,9 +433,9 @@ extension RedisConnection {
         _ logger: Logger?
     ) -> EventLoopFuture<Void> {
         let logger = self.prepareLoggerForUse(logger)
-        
+
         logger.trace("received subscribe request")
-        
+
         // if we're closed, just error out
         guard self.state.isConnected else { return self.eventLoop.makeFailedFuture(RedisClientError.connectionClosed) }
 
@@ -483,7 +483,7 @@ extension RedisConnection {
                     logger.debug("the connection is now in pubsub mode")
                 }
         }
-        
+
         // add the subscription and just ignore the subscription count
         return handler
             .addSubscription(for: target, messageReceiver: receiver, onSubscribe: onSubscribe, onUnsubscribe: onUnsubscribe)
@@ -497,27 +497,27 @@ extension RedisConnection {
     public func unsubscribe(from channels: [RedisChannelName]) -> EventLoopFuture<Void> {
         return self._unsubscribe(.channels(channels), nil)
     }
-    
+
     public func punsubscribe(from patterns: [String]) -> EventLoopFuture<Void> {
         return self._unsubscribe(.patterns(patterns), nil)
     }
-    
+
     internal func unsubscribe(from channels: [RedisChannelName], context: Context?) -> EventLoopFuture<Void> {
         return self._unsubscribe(.channels(channels), context)
     }
-    
+
     internal func punsubscribe(from patterns: [String], context: Context?) -> EventLoopFuture<Void> {
         return self._unsubscribe(.patterns(patterns), context)
     }
-    
+
     private func _unsubscribe(_ target: RedisSubscriptionTarget, _ logger: Logger?) -> EventLoopFuture<Void> {
         let logger = self.prepareLoggerForUse(logger)
-        
+
         logger.trace("received unsubscribe request")
 
         // if we're closed, just error out
         guard self.state.isConnected else { return self.eventLoop.makeFailedFuture(RedisClientError.connectionClosed) }
-        
+
         // if we're not in pubsub mode, then we just succeed as a no-op
         guard case let .pubsub(handler) = self.state else {
             // but we still assert just to give some notification to devs at debug
@@ -526,11 +526,11 @@ extension RedisConnection {
             ])
             return self.eventLoop.makeSucceededFuture(())
         }
-        
+
         logger.trace("removing subscription", metadata: [
             RedisLogging.MetadataKeys.pubsubTarget: "\(target.debugDescription)"
         ])
-        
+
         // remove the subscription
         return handler.removeSubscription(for: target)
             .flatMap {