Skip to content

Commit 1168ed0

Browse files
LukasaLukasa
committed
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.
1 parent 338a6f4 commit 1168ed0

File tree

5 files changed

+168
-3
lines changed

5 files changed

+168
-3
lines changed

‎Package.swift

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,15 +25,17 @@ let package = Package(
2525
dependencies: [
2626
.package(url: "https://github.com/apple/swift-log.git", from: "1.0.0"),
2727
.package(url: "https://github.com/apple/swift-metrics.git", "1.0.0" ..< "3.0.0"),
28-
.package(url: "https://github.com/apple/swift-nio.git", from: "2.0.0")
28+
.package(url: "https://github.com/apple/swift-nio.git", from: "2.18.0"),
29+
.package(url: "https://github.com/apple/swift-service-discovery", from: "1.0.0"),
2930
],
3031
targets: [
3132
.target(
3233
name: "RediStack",
3334
dependencies: [
3435
.product(name: "NIO", package: "swift-nio"),
3536
.product(name: "Logging", package: "swift-log"),
36-
.product(name: "Metrics", package: "swift-metrics")
37+
.product(name: "Metrics", package: "swift-metrics"),
38+
.product(name: "ServiceDiscovery", package: "swift-service-discovery")
3739
]
3840
),
3941
.testTarget(
@@ -66,7 +68,8 @@ let package = Package(
6668
name: "RediStackIntegrationTests",
6769
dependencies: [
6870
"RediStack", "RediStackTestUtils",
69-
.product(name: "NIO", package: "swift-nio")
71+
.product(name: "NIO", package: "swift-nio"),
72+
.product(name: "ServiceDiscovery", package: "swift-service-discovery")
7073
]
7174
)
7275
]

‎Sources/RediStack/RedisConnectionPool.swift

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import struct Foundation.UUID
1515
import NIO
1616
import NIOConcurrencyHelpers
1717
import Logging
18+
import ServiceDiscovery
1819

1920
/// A `RedisConnectionPool` is an implementation of `RedisClient` backed by a pool of connections to Redis,
2021
/// rather than a single one.
@@ -50,6 +51,15 @@ public class RedisConnectionPool {
5051
// This array buffers any request for a connection that cannot be succeeded right away in the case where we have no target.
5152
// We never allow this to get larger than a specific bound, to resist DoS attacks. Past that bound we will fast-fail.
5253
private var requestsForConnections: [EventLoopPromise<RedisConnection>] = []
54+
// This is var because if we're using service discovery, we don't start doing that until activate is called.
55+
private var cancellationToken: CancellationToken? {
56+
willSet {
57+
guard let token = self.cancellationToken, !token.isCancelled, token !== newValue else {
58+
return
59+
}
60+
token.cancel()
61+
}
62+
}
5363

5464
/// The maximum number of connection requests we'll buffer in `requestsForConnections` before we start fast-failing. These
5565
/// 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 {
8898
}
8999
}
90100

101+
// MARK: Alternative initializers
102+
extension RedisConnectionPool {
103+
/// Constructs a `RedisConnectionPool` that updates its addresses based on information from
104+
/// service discovery.
105+
///
106+
/// This constructor behaves similarly to the regular constructor. However, it also activates the
107+
/// connection pool before returning it to the user. This is necessary because the act of subscribing
108+
/// to service discovery forms a reference cycle between the service discovery instance and the
109+
/// `RedisConnectionPool`. Pools constructed via this constructor _must_ always have `close` called
110+
/// on them.
111+
///
112+
/// Pools created via this constructor will be auto-closed when the service discovery instance is completed for
113+
/// any reason, including on error. Users should still always call `close` in their own code during teardown.
114+
public static func activatedServiceDiscoveryPool<Discovery: ServiceDiscovery>(
115+
service: Discovery.Service,
116+
discovery: Discovery,
117+
configuration: Configuration,
118+
boundEventLoop: EventLoop,
119+
logger: Logger? = nil
120+
) -> RedisConnectionPool where Discovery.Instance == SocketAddress {
121+
let pool = RedisConnectionPool(configuration: configuration, boundEventLoop: boundEventLoop)
122+
pool.beginSubscribingToServiceDiscovery(service: service, discovery: discovery, logger: logger)
123+
pool.activate(logger: logger)
124+
return pool
125+
}
126+
}
127+
91128
// MARK: General helpers.
92129
extension RedisConnectionPool {
93130
/// Starts the connection pool.
@@ -122,6 +159,9 @@ extension RedisConnectionPool {
122159
for request in self.requestsForConnections {
123160
request.fail(RedisConnectionPoolError.poolClosed)
124161
}
162+
163+
// This cancels service discovery.
164+
self.cancellationToken = nil
125165
}
126166
}
127167

@@ -247,6 +287,36 @@ extension RedisConnectionPool {
247287
logger[metadataKey: RedisLogging.MetadataKeys.connectionPoolID] = "\(self.id)"
248288
return logger
249289
}
290+
291+
/// A private helper function used for the service discovery constructor.
292+
private func beginSubscribingToServiceDiscovery<Discovery: ServiceDiscovery>(
293+
service: Discovery.Service,
294+
discovery: Discovery,
295+
logger: Logger?
296+
) where Discovery.Instance == SocketAddress {
297+
self.loop.execute {
298+
let logger = self.prepareLoggerForUse(logger)
299+
300+
self.cancellationToken = discovery.subscribe(
301+
to: service,
302+
onNext: { result in
303+
// This closure may execute on any thread.
304+
self.loop.execute {
305+
switch result {
306+
case .success(let targets):
307+
self.updateConnectionAddresses(targets, logger: logger)
308+
case .failure(let error):
309+
logger.error("Service discovery error", metadata: [RedisLogging.MetadataKeys.error: "\(error)"])
310+
}
311+
}
312+
},
313+
onComplete: { (_: CompletionReason) in
314+
// We don't really care about the reason, we just want to brick this client.
315+
self.close(logger: logger)
316+
}
317+
)
318+
}
319+
}
250320
}
251321

252322
// MARK: RedisClient conformance

‎Sources/RediStack/RedisLogging.swift

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ public enum RedisLogging {
2121
public struct Labels {
2222
public static var connection: String { "RediStack.RedisConnection" }
2323
public static var connectionPool: String { "RediStack.RedisConnectionPool" }
24+
public static var serviceDiscovery: String { "RediStack.RedisServiceDiscoveryClient" }
2425
}
2526
/// The key values used in RediStack for storing `Logging.Logger.Metadata` in log messages.
2627
///
@@ -53,6 +54,7 @@ public enum RedisLogging {
5354

5455
public static let baseConnectionLogger = Logger(label: Labels.connection)
5556
public static let baseConnectionPoolLogger = Logger(label: Labels.connectionPool)
57+
public static let baseServiceDiscoveryLogger = Logger(label: Labels.serviceDiscovery)
5658
}
5759

5860
// MARK: Logger integration

‎Tests/RediStackIntegrationTests/RedisLoggingTests.swift

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313
//===----------------------------------------------------------------------===//
1414

1515
import Logging
16+
import ServiceDiscovery
17+
import NIO
1618
import RediStack
1719
import RediStackTestUtils
1820
import XCTest
@@ -67,6 +69,38 @@ final class RedisLoggingTests: RediStackIntegrationTestCase {
6769
.string(pool.id.uuidString)
6870
)
6971
}
72+
73+
func test_serviceDiscoveryMetadata() throws {
74+
let handler = TestLogHandler()
75+
let logger = Logger(label: #function, factory: { _ in return handler })
76+
let hosts = InMemoryServiceDiscovery<String, SocketAddress>(configuration: .init())
77+
let config = RedisConnectionPool.Configuration(
78+
initialServerConnectionAddresses: [],
79+
maximumConnectionCount: .maximumActiveConnections(1),
80+
connectionFactoryConfiguration: .init(connectionPassword: self.redisPassword)
81+
)
82+
let client = RedisConnectionPool.activatedServiceDiscoveryPool(
83+
service: "default.local",
84+
discovery: hosts,
85+
configuration: config,
86+
boundEventLoop: self.connection.eventLoop)
87+
defer {
88+
client.close()
89+
}
90+
91+
let address = try SocketAddress.makeAddressResolvingHost(self.redisHostname, port: self.redisPort)
92+
hosts.register("default.local", instances: [address])
93+
94+
_ = try client
95+
.logging(to: logger)
96+
.ping()
97+
.wait()
98+
XCTAssertTrue(handler.metadata.keys.contains(RedisLogging.MetadataKeys.connectionID))
99+
XCTAssertEqual(
100+
handler.metadata[RedisLogging.MetadataKeys.connectionPoolID],
101+
.string(client.id.uuidString)
102+
)
103+
}
70104
}
71105

72106
final class TestLogHandler: LogHandler {

‎Tests/RediStackIntegrationTests/RedisServiceDiscoveryTests.swift

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
//===----------------------------------------------------------------------===//
2+
//
3+
// This source file is part of the RediStack open source project
4+
//
5+
// Copyright (c) 2020 RediStack project authors
6+
// Licensed under Apache License v2.0
7+
//
8+
// See LICENSE.txt for license information
9+
// See CONTRIBUTORS.txt for the list of RediStack project authors
10+
//
11+
// SPDX-License-Identifier: Apache-2.0
12+
//
13+
//===----------------------------------------------------------------------===//
14+
15+
import ServiceDiscovery
16+
import NIO
17+
import Logging
18+
@testable import RediStack
19+
import RediStackTestUtils
20+
import XCTest
21+
22+
final class RedisServiceDiscoveryTests: RediStackConnectionPoolIntegrationTestCase {
23+
func test_basicServiceDiscovery() throws {
24+
let hosts = InMemoryServiceDiscovery<String, SocketAddress>(configuration: .init())
25+
let config = RedisConnectionPool.Configuration(
26+
initialServerConnectionAddresses: [],
27+
maximumConnectionCount: .maximumActiveConnections(5),
28+
connectionFactoryConfiguration: .init(connectionPassword: self.redisPassword),
29+
minimumConnectionCount: 1
30+
)
31+
let client = RedisConnectionPool.activatedServiceDiscoveryPool(
32+
service: "default.local",
33+
discovery: hosts,
34+
configuration: config,
35+
boundEventLoop: self.eventLoopGroup.next()
36+
)
37+
defer {
38+
client.close()
39+
}
40+
41+
let address = try SocketAddress.makeAddressResolvingHost(self.redisHostname, port: self.redisPort)
42+
hosts.register("default.local", instances: [address])
43+
hosts.register("another.local", instances: [])
44+
45+
// Now we try to make a bunch of requests.
46+
// We're going to insert a bunch of elements into a set, and then when all is done confirm that every
47+
// element exists.
48+
let operations = (0..<50).map { number in
49+
client.send(.sadd([number], to: #function))
50+
}
51+
let results = try EventLoopFuture<Int>.whenAllSucceed(operations, on: self.eventLoopGroup.next()).wait()
52+
XCTAssertEqual(results, Array(repeating: 1, count: 50))
53+
let whatRedisThinks = try client.send(.smembers(of: #function)).wait()
54+
XCTAssertEqual(whatRedisThinks.compactMap { $0.int }.sorted(), Array(0..<50))
55+
}
56+
}

0 commit comments

Comments
 (0)