Allow repeated commands to same connection in pool

Motivation:

Some Redis commands are very connection specific that have impacts on future access that makes it difficult in the current
checkout-use-return cycle that `RedisConnectionPool` uses.

Developers need a way to borrow a specific connection, chain several commands together, and then return the connection to the pool.

Modifications:

- Add: `leaseConnection` method to `RedisConnectionPool` which provides a connection from the pool and returns it after a provided closure's ELF resolves
- Add: `allowSubscriptions` property to `RedisConnection` for controlling the ability to make PubSub subscriptions
- Add: `RedisClientError.pubsubNotAllowed` case for when `RedisConnection.allowSubscriptions` is set to `false` and a subscription was still attempted

Result:

Developers should now have an "escape hatch" with `RedisConnectionPool` to do limited exclusive chains of operations on a specific connection.

Author: Mordil

Sources/RediStack/Connection Pool/ConnectionPool.swift

@@ -366,13 +366,17 @@ extension ConnectionPool {
         // that yet, so double-check. Leave the dead ones there: we'll get them later.
         while let connection = self.availableConnections.popLast() {
             if connection.isConnected {
+                logger.debug("found available connection", metadata: [
+                    RedisLogging.MetadataKeys.connectionID: "\(connection.id)"
+                ])
                 self.leaseConnection(connection, to: waiter)
                 return waiter.futureResult
             }
         }
 
         // Ok, we didn't have any available connections. We're going to have to wait. Set our timeout.
         waiter.scheduleDeadline(loop: self.loop, deadline: deadline) {
+            logger.trace("connection not found in time")
             // The waiter timed out. We're going to fail the promise and remove the waiter.
             waiter.fail(RedisConnectionPoolError.timedOutWaitingForConnection)
 
@@ -385,6 +389,7 @@ extension ConnectionPool {
         // below the max, or the pool is leaky, we can create a new connection. Otherwise, we just have
         // to wait for a connection to come back.
         if self.activeConnectionCount < self.maximumConnectionCount || self.leaky {
+            logger.trace("creating new connection")
             self._createConnection(backoff: self.initialBackoffDelay, startIn: .nanoseconds(0), logger: logger)
         }
 

Sources/RediStack/RedisClient.swift

@@ -249,6 +249,8 @@ public struct RedisClientError: LocalizedError, Equatable, Hashable {
     public static let connectionClosed = RedisClientError(.connectionClosed)
     /// A race condition was triggered between unsubscribing from the last target while subscribing to a new target.
     public static let subscriptionModeRaceCondition = RedisClientError(.subscriptionModeRaceCondition)
+    /// A connection that is not authorized for PubSub subscriptions attempted to create a subscription.
+    public static let pubsubNotAllowed = RedisClientError(.pubsubNotAllowed)
 
     /// Conversion from `RESPValue` to the specified type failed.
     ///
@@ -271,6 +273,7 @@ public struct RedisClientError: LocalizedError, Equatable, Hashable {
         case let .failedRESPConversion(type): message = "failed to convert RESP to \(type)"
         case let .assertionFailure(text): message = text
         case .subscriptionModeRaceCondition: message = "received request to subscribe after subscription mode has ended"
+        case .pubsubNotAllowed: message = "connection attempted to create a PubSub subscription"
         }
         return "(RediStack) \(message)"
     }
@@ -281,6 +284,7 @@ public struct RedisClientError: LocalizedError, Equatable, Hashable {
         case .failedRESPConversion: return "Ensure that the data type being requested is actually what's being returned. If you see this error and are not sure why, capture the original RESPValue string sent from Redis to add to your bug report."
         case .assertionFailure: return "This error should in theory never happen. If you trigger this error, capture the original RESPValue string sent from Redis along with the command and arguments that you sent to Redis to add to your bug report."
         case .subscriptionModeRaceCondition: return "This is a race condition where the PubSub handler was removed after a subscription was being added, but before it was committed. This can be solved by just retrying the subscription."
+        case .pubsubNotAllowed: return "When connections are managed by a pool, they are not allowed to create PubSub subscriptions on their own. Use the appropriate PubSub commands on the connection pool itself. If the connection is not managed by a pool, this is a bug and should be reported."
         }
     }
 
@@ -303,5 +307,6 @@ public struct RedisClientError: LocalizedError, Equatable, Hashable {
         case failedRESPConversion(to: Any.Type)
         case assertionFailure(message: String)
         case subscriptionModeRaceCondition
+        case pubsubNotAllowed
     }
 }

Sources/RediStack/RedisConnection.swift

@@ -125,12 +125,29 @@ public final class RedisConnection: RedisClient, RedisClientWithUserContext {
             autoflush.store(newValue)
         }
     }
-    
+    /// Controls the permission of the connection to be able to have PubSub subscriptions or not.
+    ///
+    /// When set to `true`, this connection is allowed to create subscriptions.
+    /// When set to `false`, this connection is not allowed to create subscriptions. Any potentially existing subscriptions will be removed.
+    public var allowSubscriptions: Bool {
+        get { self.allowPubSub.load() }
+        set(newValue) {
+            self.allowPubSub.store(newValue)
+            // TODO: Re-enable after [p]unsubscribe from all is fixed
+//            guard self.isConnected else { return }
+//            _ = EventLoopFuture<Void>.whenAllComplete([
+//                self.unsubscribe(),
+//                self.punsubscribe()
+//            ], on: self.eventLoop)
+        }
+    }
+
     internal let channel: Channel
     private let systemContext: Context
     private var logger: Logger { self.systemContext }
 
     private let autoflush: NIOAtomic<Bool> = .makeAtomic(value: true)
+    private let allowPubSub: NIOAtomic<Bool> = .makeAtomic(value: true)
     private let _stateLock = Lock()
     private var _state = ConnectionState.open
     private var state: ConnectionState {
@@ -418,6 +435,11 @@ extension RedisConnection {
         // if we're closed, just error out
         guard self.state.isConnected else { return self.eventLoop.makeFailedFuture(RedisClientError.connectionClosed) }
 
+        // if we're not allowed to to subscribe, then fail
+        guard self.allowSubscriptions else {
+            return self.eventLoop.makeFailedFuture(RedisClientError.pubsubNotAllowed)
+        }
+
         logger.trace("adding subscription", metadata: [
             RedisLogging.MetadataKeys.pubsubTarget: "\(target.debugDescription)"
         ])

Sources/RediStack/RedisConnectionPool.swift

@@ -148,6 +148,47 @@ extension RedisConnectionPool {
         }
     }
 
+    /// Provides limited exclusive access to a connection to be used in a user-defined specialized closure of operations.
+    /// - Warning: Attempting to create PubSub subscriptions with connections leased in the closure will result in a failed `NIO.EventLoopFuture`.
+    ///
+    /// `RedisConnectionPool` manages PubSub state and requires exclusive control over creating PubSub subscriptions.
+    /// - Important: This connection **MUST NOT** be stored outside of the closure. It is only available exclusively within the closure.
+    ///
+    /// All operations should be done inside the closure as chained `NIO.EventLoopFuture` callbacks.
+    ///
+    /// For example:
+    /// ```swift
+    /// let countFuture = pool.leaseConnection {
+    ///     $0.logging(to: myLogger)
+    ///         .authorize(with: userPassword)
+    ///         .flatMap { connection.select(database: userDatabase) }
+    ///         .flatMap { connection.increment(counterKey) }
+    /// }
+    /// ```
+    /// - Warning: Some commands change the state of the connection that are not tracked client-side,
+    /// and will not be automatically reset when the connection is returned to the pool.
+    ///
+    /// When the connection is reused from the pool, it will retain this state and may affect future commands executed with it.
+    ///
+    /// For example, if `select(database:)` is used, all future commands made with this connection will be against the selected database.
+    ///
+    /// To protect against future issues, make sure the final commands executed are to reset the connection to it's previous known state.
+    /// - Parameter operation: A closure that receives exclusive access to the provided `RedisConnection` for the lifetime of the closure for specialized Redis command chains.
+    /// - Returns: A `NIO.EventLoopFuture` that resolves the value of the `NIO.EventLoopFuture` in the provided closure operation.
+    @inlinable
+    public func leaseConnection<T>(_ operation: @escaping (RedisConnection) -> EventLoopFuture<T>) -> EventLoopFuture<T> {
+        return self.forwardOperationToConnection(
+            {
+                (connection, returnConnection, context) in
+        
+                return operation(connection)
+                    .always { _ in returnConnection(connection, context) }
+            },
+            preferredConnection: nil,
+            context: nil
+        )
+    }
+
     /// Updates the list of valid connection addresses.
     ///
     /// - Note: This does not invalidate existing connections: as long as those connections continue to stay up, they will be kept by
@@ -179,13 +220,16 @@ extension RedisConnectionPool {
             return targetLoop.makeFailedFuture(RedisConnectionPoolError.noAvailableConnectionTargets)
         }
 
-        return RedisConnection.connect(
+        let connectFuture = RedisConnection.connect(
             to: nextTarget,
             on: targetLoop,
             password: self.connectionPassword,
             logger: self.connectionSystemContext,
             tcpClient: self.connectionTCPClient
         )
+        // disallow subscriptions on all connections by default so that we can enforce our management of PubSub state
+        connectFuture.whenSuccess { $0.allowSubscriptions = false }
+        return connectFuture
     }
 
     private func prepareLoggerForUse(_ logger: Logger?) -> Logger {
@@ -323,7 +367,10 @@ extension RedisConnectionPool: RedisClientWithUserContext {
         return self.forwardOperationToConnection(
             { (connection, returnConnection, context) in
 
-                if self.pubsubConnection == nil { self.pubsubConnection = connection }
+                if self.pubsubConnection == nil {
+                    connection.allowSubscriptions = true // allow pubsub commands which are to come
+                    self.pubsubConnection = connection
+                }
 
                 let onUnsubscribe: RedisSubscriptionChangeHandler = { channelName, subCount in
                     defer { unsubscribeHandler?(channelName, subCount) }
@@ -332,7 +379,8 @@ extension RedisConnectionPool: RedisClientWithUserContext {
                         subCount == 0,
                         let connection = self.pubsubConnection
                     else { return }
-                    
+
+                    connection.allowSubscriptions = false // reset PubSub permissions
                     returnConnection(connection, context)
                     self.pubsubConnection = nil // break ref cycle
                 }
@@ -367,7 +415,8 @@ extension RedisConnectionPool: RedisClientWithUserContext {
         )
     }
 
-    private func forwardOperationToConnection<T>(
+    @usableFromInline
+    internal func forwardOperationToConnection<T>(
         _ operation: @escaping (RedisConnection, @escaping (RedisConnection, Context) -> Void, Context) -> EventLoopFuture<T>,
         preferredConnection: RedisConnection?,
         context: Context?

Sources/RediStack/RedisContext.swift

@@ -20,4 +20,5 @@ import struct Logging.Logger
 // so in order to be "future thinking" we create this typealias and interally refer to this passing of configuration
 // as context
 
+@usableFromInline
 internal typealias Context = Logging.Logger

Sources/RediStackTestUtils/RedisConnectionPoolIntegrationTestCase.swift

@@ -75,14 +75,15 @@ open class RedisConnectionPoolIntegrationTestCase: XCTestCase {
     }
 
     public func makeNewPool(
-        connectionRetryTimeout: TimeAmount? = .seconds(5)
+        connectionRetryTimeout: TimeAmount? = .seconds(5),
+        minimumConnectionCount: Int = 0
     ) throws -> RedisConnectionPool {
         let address = try SocketAddress.makeAddressResolvingHost(self.redisHostname, port: self.redisPort)
         let pool = RedisConnectionPool(
             serverConnectionAddresses: [address],
             loop: self.eventLoopGroup.next(),
             maximumConnectionCount: .maximumActiveConnections(4),
-            minimumConnectionCount: 0,
+            minimumConnectionCount: minimumConnectionCount,
             connectionPassword: self.redisPassword,
             connectionRetryTimeout: connectionRetryTimeout
         )

Tests/RediStackIntegrationTests/Commands/PubSubCommandsTests.swift

@@ -159,7 +159,6 @@ final class RedisPubSubCommandsPoolTests: RediStackConnectionPoolIntegrationTest
 
         let channel = RedisChannelName(#function)
         let pattern = "\(channel.rawValue.dropLast(channel.rawValue.count / 2))*"
-        print(channel, pattern)
 
         try subscriber
             .subscribe(to: channel) { (_, _) in channelMessageExpectation.fulfill() }

Tests/RediStackIntegrationTests/RedisConnectionPoolTests.swift

@@ -43,3 +43,61 @@ final class RedisConnectionPoolTests: RediStackConnectionPoolIntegrationTestCase
         XCTAssertNoThrow(try pool.get(#function).wait())
     }
 }
+
+// MARK: Leasing a connection
+
+extension RedisConnectionPoolTests {
+    func test_borrowedConnectionStillReturnsOnError() throws {
+        enum TestError: Error { case expected }
+
+        let maxConnectionCount = 4
+        let pool = try self.makeNewPool(minimumConnectionCount: maxConnectionCount)
+        defer { pool.close() }
+        _ = try pool.ping().wait()
+
+        let promise = pool.eventLoop.makePromise(of: Void.self)
+        
+        XCTAssertEqual(pool.availableConnectionCount, maxConnectionCount)
+        defer { XCTAssertEqual(pool.availableConnectionCount, maxConnectionCount) }
+
+        let future = pool.leaseConnection { _ in promise.futureResult }
+        
+        promise.fail(TestError.expected)
+        XCTAssertThrowsError(try future.wait()) {
+            XCTAssertTrue($0 is TestError)
+        }
+    }
+    
+    func test_borrowedConnectionClosureHasExclusiveAccess() throws {
+        let maxConnectionCount = 4
+        let pool = try self.makeNewPool(minimumConnectionCount: maxConnectionCount)
+        defer { pool.close() }
+        // populate the connection pool
+        _ = try pool.ping().wait()
+        
+        // assert that we have the max number of connections available,
+        XCTAssertEqual(pool.availableConnectionCount, maxConnectionCount)
+
+        // borrow a connection, asserting that we've taken the connection out of the pool while we do "something" with it
+        // and then assert afterwards that it's back in the pool
+        
+        let promises: [EventLoopPromise<Void>] = [pool.eventLoop.makePromise(), pool.eventLoop.makePromise()]
+        let futures = promises.indices
+            .map { index in
+                return pool
+                    .leaseConnection { connection -> EventLoopFuture<Void> in
+                        XCTAssertTrue(pool.availableConnectionCount < maxConnectionCount)
+                
+                        return promises[index].futureResult
+                    }
+            }
+        
+        promises.forEach { $0.succeed(()) }
+        _ = try EventLoopFuture<Void>
+            .whenAllSucceed(futures, on: pool.eventLoop)
+            .always { _ in
+                XCTAssertEqual(pool.availableConnectionCount, maxConnectionCount)
+            }
+            .wait()
+    }
+}

Tests/RediStackIntegrationTests/RedisConnectionTests.swift

@@ -40,3 +40,43 @@ final class RedisConnectionTests: RediStackIntegrationTestCase {
         }
     }
 }
+
+// MARK: PubSub permissions
+
+extension RedisConnectionTests {
+    func test_subscriptionNotAllowedFails() throws {
+        self.connection.allowSubscriptions = false
+        let subscription = self.connection.subscribe(to: #function) { (_, _) in }
+
+        XCTAssertThrowsError(try subscription.wait()) {
+            guard let error = $0 as? RedisClientError else {
+                XCTFail("unexpected error type: \(type(of: $0))")
+                return
+            }
+            XCTAssertEqual(error, .pubsubNotAllowed)
+        }
+    }
+    
+    // TODO - fix [p]unsubscribe from all and re-enable this unit test
+//    func test_subscriptionPermissionsChanged_endsSubscriptions() throws {
+//        let connection = try self.makeNewConnection()
+//
+//        let channelSubClosedExpectation = self.expectation(description: "channel subscription was closed")
+//        let patternSubClosedExpectation = self.expectation(description: "pattern subscription was closed")
+//
+//        _ = connection.subscribe(
+//            to: #function,
+//            messageReceiver:  { (_, _) in },
+//            onUnsubscribe: { (_, _) in channelSubClosedExpectation.fulfill() }
+//        )
+//        _ = connection.psubscribe(
+//            to: #function,
+//            messageReceiver:  { (_, _) in },
+//            onUnsubscribe: { (_, _) in patternSubClosedExpectation.fulfill() }
+//        )
+//
+//        connection.allowSubscriptions = false
+//
+//        self.waitForExpectations(timeout: 2)
+//    }
+}