Add support for service discovery.

The newly-released Service Discovery framework gives us the interesting
opportunity to make RediStack aware of complex service discovery tools.
This patch supplies a simple adaptor to integrat Service Discovery with
RediStack's pooled client, allowing users to work with arbitrary service
discovery systems.

Author: Lukasa

Package.swift

@@ -25,15 +25,17 @@ let package = Package(
     dependencies: [
         .package(url: "https://github.com/apple/swift-log.git", from: "1.0.0"),
         .package(url: "https://github.com/apple/swift-metrics.git", "1.0.0" ..< "3.0.0"),
-        .package(url: "https://github.com/apple/swift-nio.git", from: "2.0.0")
+        .package(url: "https://github.com/apple/swift-nio.git", from: "2.18.0"),
+        .package(url: "https://github.com/apple/swift-service-discovery", from: "1.0.0"),
     ],
     targets: [
         .target(
             name: "RediStack",
             dependencies: [
                 .product(name: "NIO", package: "swift-nio"),
                 .product(name: "Logging", package: "swift-log"),
-                .product(name: "Metrics", package: "swift-metrics")
+                .product(name: "Metrics", package: "swift-metrics"),
+                .product(name: "ServiceDiscovery", package: "swift-service-discovery")
             ]
         ),
         .testTarget(
@@ -66,7 +68,8 @@ let package = Package(
             name: "RediStackIntegrationTests",
             dependencies: [
                 "RediStack", "RediStackTestUtils",
-                .product(name: "NIO", package: "swift-nio")
+                .product(name: "NIO", package: "swift-nio"),
+                .product(name: "ServiceDiscovery", package: "swift-service-discovery")
             ]
         )
     ]

Sources/RediStack/RedisConnectionPool.swift

@@ -15,6 +15,7 @@ import struct Foundation.UUID
 import NIO
 import NIOConcurrencyHelpers
 import Logging
+import ServiceDiscovery
 
 /// A `RedisConnectionPool` is an implementation of `RedisClient` backed by a pool of connections to Redis,
 /// rather than a single one.
@@ -50,6 +51,15 @@ public class RedisConnectionPool {
     // This array buffers any request for a connection that cannot be succeeded right away in the case where we have no target.
     // We never allow this to get larger than a specific bound, to resist DoS attacks. Past that bound we will fast-fail.
     private var requestsForConnections: [EventLoopPromise<RedisConnection>] = []
+    // This is var because if we're using service discovery, we don't start doing that until activate is called.
+    private var cancellationToken: CancellationToken? {
+        willSet {
+            guard let token = self.cancellationToken, !token.isCancelled, token !== newValue else {
+                return
+            }
+            token.cancel()
+        }
+    }
 
     /// The maximum number of connection requests we'll buffer in `requestsForConnections` before we start fast-failing. These
     /// are buffered only when there are no available addresses to connect to, so in practice it's highly unlikely this will be
@@ -88,6 +98,33 @@ public class RedisConnectionPool {
     }
 }
 
+// MARK: Alternative initializers
+extension RedisConnectionPool {
+    /// Constructs a `RedisConnectionPool` that updates its addresses based on information from
+    /// service discovery.
+    ///
+    /// This constructor behaves similarly to the regular constructor. However, it also activates the
+    /// connection pool before returning it to the user. This is necessary because the act of subscribing
+    /// to service discovery forms a reference cycle between the service discovery instance and the
+    /// `RedisConnectionPool`. Pools constructed via this constructor _must_ always have `close` called
+    /// on them.
+    ///
+    /// Pools created via this constructor will be auto-closed when the service discovery instance is completed for
+    /// any reason, including on error. Users should still always call `close` in their own code during teardown.
+    public static func activatedServiceDiscoveryPool<Discovery: ServiceDiscovery>(
+        service: Discovery.Service,
+        discovery: Discovery,
+        configuration: Configuration,
+        boundEventLoop: EventLoop,
+        logger: Logger? = nil
+    ) -> RedisConnectionPool where Discovery.Instance == SocketAddress {
+        let pool = RedisConnectionPool(configuration: configuration, boundEventLoop: boundEventLoop)
+        pool.beginSubscribingToServiceDiscovery(service: service, discovery: discovery, logger: logger)
+        pool.activate(logger: logger)
+        return pool
+    }
+}
+
 // MARK: General helpers.
 extension RedisConnectionPool {
     /// Starts the connection pool.
@@ -122,6 +159,9 @@ extension RedisConnectionPool {
             for request in self.requestsForConnections {
                 request.fail(RedisConnectionPoolError.poolClosed)
             }
+
+            // This cancels service discovery.
+            self.cancellationToken = nil
         }
     }
 
@@ -247,6 +287,36 @@ extension RedisConnectionPool {
         logger[metadataKey: RedisLogging.MetadataKeys.connectionPoolID] = "\(self.id)"
         return logger
     }
+
+    /// A private helper function used for the service discovery constructor.
+    private func beginSubscribingToServiceDiscovery<Discovery: ServiceDiscovery>(
+        service: Discovery.Service,
+        discovery: Discovery,
+        logger: Logger?
+    ) where Discovery.Instance == SocketAddress {
+        self.loop.execute {
+            let logger = self.prepareLoggerForUse(logger)
+
+            self.cancellationToken = discovery.subscribe(
+                to: service,
+                onNext: { result in
+                    // This closure may execute on any thread.
+                    self.loop.execute {
+                        switch result {
+                        case .success(let targets):
+                            self.updateConnectionAddresses(targets, logger: logger)
+                        case .failure(let error):
+                            logger.error("Service discovery error", metadata: [RedisLogging.MetadataKeys.error: "\(error)"])
+                        }
+                    }
+                },
+                onComplete: { (_: CompletionReason) in
+                    // We don't really care about the reason, we just want to brick this client.
+                    self.close(logger: logger)
+                }
+            )
+        }
+    }
 }
 
 // MARK: RedisClient conformance

Sources/RediStack/RedisLogging.swift

@@ -21,6 +21,7 @@ public enum RedisLogging {
     public struct Labels {
         public static var connection: String { "RediStack.RedisConnection" }
         public static var connectionPool: String { "RediStack.RedisConnectionPool" }
+        public static var serviceDiscovery: String { "RediStack.RedisServiceDiscoveryClient" }
     }
     /// The key values used in RediStack for storing `Logging.Logger.Metadata` in log messages.
     ///
@@ -53,6 +54,7 @@ public enum RedisLogging {
 
     public static let baseConnectionLogger = Logger(label: Labels.connection)
     public static let baseConnectionPoolLogger = Logger(label: Labels.connectionPool)
+    public static let baseServiceDiscoveryLogger = Logger(label: Labels.serviceDiscovery)
 }
 
 // MARK: Logger integration

Tests/RediStackIntegrationTests/RedisLoggingTests.swift

@@ -13,6 +13,8 @@
 //===----------------------------------------------------------------------===//
 
 import Logging
+import ServiceDiscovery
+import NIO
 import RediStack
 import RediStackTestUtils
 import XCTest
@@ -67,6 +69,38 @@ final class RedisLoggingTests: RediStackIntegrationTestCase {
             .string(pool.id.uuidString)
         )
     }
+
+    func test_serviceDiscoveryMetadata() throws {
+        let handler = TestLogHandler()
+        let logger = Logger(label: #function, factory: { _ in return handler })
+        let hosts = InMemoryServiceDiscovery<String, SocketAddress>(configuration: .init())
+        let config = RedisConnectionPool.Configuration(
+            initialServerConnectionAddresses: [],
+            maximumConnectionCount: .maximumActiveConnections(1),
+            connectionFactoryConfiguration: .init(connectionPassword: self.redisPassword)
+        )
+        let client = RedisConnectionPool.activatedServiceDiscoveryPool(
+            service: "default.local",
+            discovery: hosts,
+            configuration: config,
+            boundEventLoop: self.connection.eventLoop)
+        defer {
+            client.close()
+        }
+
+        let address = try SocketAddress.makeAddressResolvingHost(self.redisHostname, port: self.redisPort)
+        hosts.register("default.local", instances: [address])
+
+        _ = try client
+            .logging(to: logger)
+            .ping()
+            .wait()
+        XCTAssertTrue(handler.metadata.keys.contains(RedisLogging.MetadataKeys.connectionID))
+        XCTAssertEqual(
+            handler.metadata[RedisLogging.MetadataKeys.connectionPoolID],
+            .string(client.id.uuidString)
+        )
+    }
 }
 
 final class TestLogHandler: LogHandler {

Tests/RediStackIntegrationTests/RedisServiceDiscoveryTests.swift

@@ -0,0 +1,56 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the RediStack open source project
+//
+// Copyright (c) 2020 RediStack project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of RediStack project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import ServiceDiscovery
+import NIO
+import Logging
+@testable import RediStack
+import RediStackTestUtils
+import XCTest
+
+final class RedisServiceDiscoveryTests: RediStackConnectionPoolIntegrationTestCase {
+    func test_basicServiceDiscovery() throws {
+        let hosts = InMemoryServiceDiscovery<String, SocketAddress>(configuration: .init())
+        let config = RedisConnectionPool.Configuration(
+            initialServerConnectionAddresses: [],
+            maximumConnectionCount: .maximumActiveConnections(5),
+            connectionFactoryConfiguration: .init(connectionPassword: self.redisPassword),
+            minimumConnectionCount: 1
+        )
+        let client = RedisConnectionPool.activatedServiceDiscoveryPool(
+            service: "default.local",
+            discovery: hosts,
+            configuration: config,
+            boundEventLoop: self.eventLoopGroup.next()
+        )
+        defer {
+            client.close()
+        }
+
+        let address = try SocketAddress.makeAddressResolvingHost(self.redisHostname, port: self.redisPort)
+        hosts.register("default.local", instances: [address])
+        hosts.register("another.local", instances: [])
+
+        // Now we try to make a bunch of requests.
+        // We're going to insert a bunch of elements into a set, and then when all is done confirm that every
+        // element exists.
+        let operations = (0..<50).map { number in
+            client.send(.sadd([number], to: #function))
+        }
+        let results = try EventLoopFuture<Int>.whenAllSucceed(operations, on: self.eventLoopGroup.next()).wait()
+        XCTAssertEqual(results, Array(repeating: 1, count: 50))
+        let whatRedisThinks = try client.send(.smembers(of: #function)).wait()
+        XCTAssertEqual(whatRedisThinks.compactMap { $0.int }.sorted(), Array(0..<50))
+    }
+}