Add NIORedisPipeline and RedisPipeline to allow chaining requests and receiving a single response

Author: Mordil

Sources/NIORedis/NIORedisConnection.swift

@@ -31,7 +31,7 @@ public final class NIORedisConnection {
     /// Executes the desired command with the specified arguments.
     /// - Important: All arguments should be in `.bulkString` format.
     public func command(_ command: String, _ arguments: [RedisData] = []) -> EventLoopFuture<RedisData> {
-        return send(.array([RedisData(bulk: command)] + arguments))
+        return _send(.array([RedisData(bulk: command)] + arguments))
             .thenThrowing { response in
                 switch response {
                 case let .error(error): throw error
@@ -40,7 +40,12 @@ public final class NIORedisConnection {
             }
     }
 
-    private func send(_ message: RedisData) -> EventLoopFuture<RedisData> {
+    /// Creates a `NIORedisPipeline` for executing a batch of commands.
+    public func makePipeline() -> NIORedisPipeline {
+        return .init(using: self)
+    }
+
+    func _send(_ message: RedisData) -> EventLoopFuture<RedisData> {
         // ensure the connection is still open
         guard !isClosed.load() else { return eventLoop.makeFailedFuture(error: RedisError.connectionClosed) }
 

Sources/NIORedis/NIORedisPipeline.swift

@@ -0,0 +1,86 @@
+import Foundation
+import NIO
+
+/// An object that provides a mechanism to "pipeline" multiple Redis commands in sequence, providing an aggregate response
+/// of all the Redis responses for each individual command.
+///
+///     let results = connection.makePipeline()
+///         .enqueue(command: "SET", arguments: ["my_key", 3])
+///         .enqueue(command: "INCR", arguments: ["my_key"])
+///         .execute()
+///     // results == Future<[RedisData]>
+///     // results[0].string == Optional("OK")
+///     // results[1].int == Optional(4)
+/// - Important: The larger the pipeline queue, the more memory both NIORedis and Redis will use.
+/// See https://redis.io/topics/pipelining#redis-pipelining
+public final class NIORedisPipeline {
+    /// The client to execute the commands on.
+    private let connection: NIORedisConnection
+    private let encoder: RedisDataEncoder = .init()
+
+    /// The queue of complete, encoded commands to execute.
+    private var queue: [RedisData]
+    private var messageCount: Int
+
+    /// Creates a new pipeline queue using the provided `NIORedisConnection`.
+    /// - Parameter using: The connection to execute the commands on.
+    public init(using connection: NIORedisConnection) {
+        self.connection = connection
+        self.queue = []
+        self.messageCount = 0
+    }
+
+    /// Queues the provided command and arguments to be executed when `execute()` is invoked.
+    /// - Parameters:
+    ///     - command: The command to execute. See https://redis.io/commands
+    ///     - arguments: The arguments, if any, to send with the command.
+    /// - Returns: A self-reference to this `NIORedisPipeline` instance for chaining commands.
+    @discardableResult
+    public func enqueue(command: String, arguments: [RedisDataConvertible] = []) throws -> NIORedisPipeline {
+        let args = try arguments.map { try $0.convertToRedisData() }
+
+        queue.append(.array([RedisData(bulk: command)] + args))
+
+        return self
+    }
+
+    /// Flushes the queue, sending all of the commands to Redis in the same order as they were enqueued.
+    /// - Important: If any of the commands fail, the remaining commands will not execute and the `EventLoopFuture` will fail.
+    /// - Returns: A `EventLoopFuture` that resolves the `RedisData` responses, in the same order as the command queue.
+    public func execute() -> EventLoopFuture<[RedisData]> {
+        let promise = connection.eventLoop.makePromise(of: [RedisData].self)
+
+        var results = [RedisData]()
+        var iterator = queue.makeIterator()
+
+        // recursive internal method for chaining each request and
+        // attaching callbacks for failing or ultimately succeeding
+        func handle(_ command: RedisData) {
+            let future = connection._send(command)
+            future.whenSuccess { response in
+                switch response {
+                case let .error(error): promise.fail(error: error)
+                default:
+                    results.append(response)
+
+                    if let next = iterator.next() {
+                        handle(next)
+                    } else {
+                        promise.succeed(result: results)
+                    }
+                }
+            }
+            future.whenFailure { promise.fail(error: $0) }
+        }
+
+        if let first = iterator.next() {
+            handle(first)
+        } else {
+            promise.succeed(result: [])
+        }
+
+        promise.futureResult.whenComplete { self.queue = [] }
+
+        return promise.futureResult
+    }
+}

Sources/Redis/RedisConnection.swift

@@ -2,22 +2,28 @@ import Foundation
 import NIORedis
 
 public final class RedisConnection {
-    private let driverConnection: NIORedisConnection
+    let _driverConnection: NIORedisConnection
+
     private let queue: DispatchQueue
 
-    deinit { driverConnection.close() }
+    deinit { _driverConnection.close() }
 
     init(driver: NIORedisConnection, callbackQueue: DispatchQueue) {
-        self.driverConnection = driver
+        self._driverConnection = driver
         self.queue = callbackQueue
     }
 
+    /// Creates a `RedisPipeline` for executing a batch of commands.
+    public func makePipeline(callbackQueue: DispatchQueue? = nil) -> RedisPipeline {
+        return .init(using: self, callbackQueue: callbackQueue ?? queue)
+    }
+
     public func get(
         _ key: String,
         _ callback: @escaping (Result<String?, Error>
     ) -> Void) {
         // TODO: Make this a generic method to avoid copy/paste
-        driverConnection.get(key)
+        _driverConnection.get(key)
             .map { result in
                 self.queue.async { callback(.success(result)) }
             }

Sources/Redis/RedisPipeline.swift

@@ -0,0 +1,52 @@
+import Foundation
+import NIORedis
+
+/// An object that provides a mechanism to "pipeline" multiple Redis commands in sequence, providing an aggregate response
+/// of all the Redis responses for each individual command.
+///
+///     connection.makePipeline()
+///         .enqueue(command: "SET", arguments: ["my_key", 3])
+///         .enqueue(command: "INCR", arguments: ["my_key"])
+///         .execute { results in
+///             // results[0].string == Optional("OK")
+///             // results[1].int == Optional(4)
+///         }
+/// - Important: The larger the pipeline queue, the more memory both the Redis driver and Redis server will use.
+/// See https://redis.io/topics/pipelining#redis-pipelining
+public final class RedisPipeline {
+    private let _driverPipeline: NIORedisPipeline
+    private let queue: DispatchQueue
+
+    /// Creates a new pipeline queue using the provided `RedisConnection`, executing callbacks on the provided `DispatchQueue`.
+    /// - Parameters:
+    ///     - using: The connection to execute the commands on.
+    ///     - callbackQueue: The queue to execute all callbacks on.
+    public init(using connection: RedisConnection, callbackQueue: DispatchQueue) {
+        self._driverPipeline = NIORedisPipeline(using: connection._driverConnection)
+        self.queue = callbackQueue
+    }
+
+    /// Queues the provided command and arguments to be executed when `execute()` is invoked.
+    /// - Parameters:
+    ///     - command: The command to execute. See https://redis.io/commands
+    ///     - arguments: The arguments, if any, to send with the command.
+    /// - Returns: A self-reference to this `RedisPipeline` instance for chaining commands.
+    @discardableResult
+    public func enqueue(command: String, arguments: [RedisDataConvertible] = []) throws -> RedisPipeline {
+        try _driverPipeline.enqueue(command: command, arguments: arguments)
+        return self
+    }
+
+    /// Flushes the queue, sending all of the commands to Redis in the same order as they were enqueued.
+    /// - Important: If any of the commands fail, the remaining commands will not execute and the callback will receive a failure.
+    /// - Parameter callback: The callback to receive the results of the pipeline of commands, or an error if thrown.
+    public func execute(_ callback: @escaping (Result<[RedisData], Error>) -> Void) {
+        _driverPipeline.execute()
+            .map { results in
+                self.queue.async { callback(.success(results)) }
+            }
+            .whenFailure { error in
+                self.queue.async { callback(.failure(error)) }
+            }
+    }
+}

Tests/NIORedisTests/NIORedisPipelineTests.swift

@@ -0,0 +1,83 @@
+@testable import NIORedis
+import XCTest
+
+final class NIORedisPipelineTests: XCTestCase {
+    private var redis: NIORedis!
+    private var connection: NIORedisConnection!
+
+    override func setUp() {
+        let redis = NIORedis(executionModel: .spawnThreads(2))
+
+        guard let connection = try? redis.makeConnection().wait() else {
+            return XCTFail("Failed to create connection!")
+        }
+
+        self.redis = redis
+        self.connection = connection
+    }
+
+    override func tearDown() {
+        _ = try? connection.command("FLUSHALL").wait()
+        connection.close()
+        try? redis.terminate()
+    }
+
+    func test_enqueue() {
+        let pipeline = connection.makePipeline()
+
+        XCTAssertNoThrow(try pipeline.enqueue(command: "PING"))
+        XCTAssertNoThrow(try pipeline.enqueue(command: "SET", arguments: ["KEY", 3]))
+    }
+
+    func test_executeFails() throws {
+        let pipeline = try connection.makePipeline()
+            .enqueue(command: "GET")
+
+        XCTAssertThrowsError(try pipeline.execute().wait())
+    }
+
+    func test_singleCommand() throws {
+        let results = try connection.makePipeline()
+            .enqueue(command: "PING")
+            .execute()
+            .wait()
+
+        XCTAssertEqual(results[0].string, "PONG")
+    }
+
+    func test_multipleCommands() throws {
+        let results = try connection.makePipeline()
+            .enqueue(command: "PING")
+            .enqueue(command: "SET", arguments: ["my_key", 3])
+            .enqueue(command: "GET", arguments: ["my_key"])
+            .execute()
+            .wait()
+
+        XCTAssertEqual(results[0].string, "PONG")
+        XCTAssertEqual(results[1].string, "OK")
+        XCTAssertEqual(results[2].data, "3".convertedToData())
+    }
+
+    func test_executeIsOrdered() throws {
+        let results = try connection.makePipeline()
+            .enqueue(command: "SET", arguments: ["key", 1])
+            .enqueue(command: "INCR", arguments: ["key"])
+            .enqueue(command: "DECR", arguments: ["key"])
+            .enqueue(command: "INCRBY", arguments: ["key", 15])
+            .execute()
+            .wait()
+
+        XCTAssertEqual(results[0].string, "OK")
+        XCTAssertEqual(results[1].int, 2)
+        XCTAssertEqual(results[2].int, 1)
+        XCTAssertEqual(results[3].int, 16)
+    }
+
+    static var allTests = [
+        ("test_enqueue", test_enqueue),
+        ("test_executeFails", test_executeFails),
+        ("test_singleCommand", test_singleCommand),
+        ("test_multipleCommands", test_multipleCommands),
+        ("test_executeIsOrdered", test_executeIsOrdered),
+    ]
+}

Tests/NIORedisTests/XCTestManifests.swift

@@ -10,6 +10,7 @@ public func allTests() -> [XCTestCaseEntry] {
         testCase(RedisDataEncoderTests.allTests),
         testCase(RedisDataEncoderParsingTests.allTests),
         testCase(BasicCommandsTests.allTests),
+        testCase(NIORedisPipelineTests.allTests)
     ]
 }
 #endif