Skip to content

Commit 733a40b

Browse files
kimdvkimdv
committed
Initial commit of DataLoader
1 parent 81a2b33 commit 733a40b

File tree

9 files changed

+772
-0
lines changed

9 files changed

+772
-0
lines changed

‎Package.resolved

Lines changed: 25 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Package.swift

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
// swift-tools-version:4.1
2+
// The swift-tools-version declares the minimum version of Swift required to build this package.
3+
4+
import PackageDescription
5+
6+
let package = Package(
7+
name: "SwiftDataLoader",
8+
products: [
9+
.library(
10+
name: "SwiftDataLoader",
11+
targets: ["SwiftDataLoader"]),
12+
],
13+
dependencies: [
14+
.package(url: "https://github.com/apple/swift-nio.git", from: "1.8.0"),
15+
],
16+
targets: [
17+
.target(
18+
name: "SwiftDataLoader",
19+
dependencies: ["NIO"]),
20+
.testTarget(
21+
name: "SwiftDataLoaderTests",
22+
dependencies: ["SwiftDataLoader"]),
23+
]
24+
)

‎Sources/SwiftDataLoader/DataLoader.swift

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
//
2+
// DataLoader.swift
3+
// App
4+
//
5+
// Created by Kim de Vos on 01/06/2018.
6+
//
7+
import NIO
8+
9+
public enum DataLoaderFutureValue<T> {
10+
case success(T)
11+
case failure(Error)
12+
}
13+
14+
public typealias BatchLoadFunction<Key, Value> = (_ keys: [Key]) throws -> EventLoopFuture<[DataLoaderFutureValue<Value>]>
15+
16+
// Private
17+
private typealias LoaderQueue<Key, Value> = Array<(key: Key, promise: EventLoopPromise<Value>)>
18+
19+
final public class DataLoader<Key: Hashable, Value> {
20+
21+
private let batchLoadFunction: BatchLoadFunction<Key, Value>
22+
private let options: DataLoaderOptions<Key, Value>
23+
24+
private var futureCache = [Key: EventLoopFuture<Value>]()
25+
private var queue = LoaderQueue<Key, Value>()
26+
27+
public init(options: DataLoaderOptions<Key, Value> = DataLoaderOptions(), batchLoadFunction: @escaping BatchLoadFunction<Key, Value>) {
28+
self.options = options
29+
self.batchLoadFunction = batchLoadFunction
30+
}
31+
32+
33+
/// Loads a key, returning a `Promise` for the value represented by that key.
34+
public func load(key: Key, on eventLoop: EventLoopGroup) throws -> EventLoopFuture<Value> {
35+
let cacheKey = options.cacheKeyFunction?(key) ?? key
36+
37+
if options.cachingEnabled, let cachedFuture = futureCache[cacheKey] {
38+
return cachedFuture
39+
}
40+
41+
let promise: EventLoopPromise<Value> = eventLoop.next().newPromise()
42+
43+
if options.batchingEnabled {
44+
queue.append((key: key, promise: promise))
45+
} else {
46+
_ = try batchLoadFunction([key]).map { results in
47+
if results.isEmpty {
48+
promise.fail(error: DataLoaderError.noValueForKey("Did not return value for key: \(key)"))
49+
} else {
50+
let result = results[0]
51+
switch result {
52+
case .success(let value): promise.succeed(result: value)
53+
case .failure(let error): promise.fail(error: error)
54+
}
55+
}
56+
}
57+
}
58+
59+
let future = promise.futureResult
60+
61+
if options.cachingEnabled {
62+
futureCache[cacheKey] = future
63+
}
64+
65+
return future
66+
}
67+
68+
public func loadMany(keys: [Key], on eventLoop: EventLoopGroup) throws -> EventLoopFuture<[Value]> {
69+
guard !keys.isEmpty else { return eventLoop.next().newSucceededFuture(result: []) }
70+
71+
let promise: EventLoopPromise<[Value]> = eventLoop.next().newPromise()
72+
73+
var result = [Value]()
74+
75+
let futures = try keys.map { try load(key: $0, on: eventLoop) }
76+
77+
for future in futures {
78+
_ = future.map { value in
79+
result.append(value)
80+
81+
if result.count == keys.count {
82+
promise.succeed(result: result)
83+
}
84+
}
85+
}
86+
87+
return promise.futureResult
88+
}
89+
90+
func clear(key: Key) -> DataLoader<Key, Value> {
91+
let cacheKey = options.cacheKeyFunction?(key) ?? key
92+
futureCache.removeValue(forKey: cacheKey)
93+
return self
94+
}
95+
96+
func clearAll() -> DataLoader<Key, Value> {
97+
futureCache.removeAll()
98+
return self
99+
}
100+
101+
func prime(key: Key, value: Value, on eventLoop: EventLoopGroup) -> DataLoader<Key, Value> {
102+
let cacheKey = options.cacheKeyFunction?(key) ?? key
103+
104+
if futureCache[cacheKey] == nil {
105+
let promise: EventLoopPromise<Value> = eventLoop.next().newPromise()
106+
promise.succeed(result: value)
107+
108+
futureCache[cacheKey] = promise.futureResult
109+
}
110+
111+
return self
112+
}
113+
114+
// MARK: - Private
115+
private func dispatchQueueBatch(queue: LoaderQueue<Key, Value>, on eventLoop: EventLoopGroup) throws { //}-> EventLoopFuture<[Value]> {
116+
let keys = queue.map { $0.key }
117+
118+
if keys.isEmpty {
119+
return //eventLoop.next().newSucceededFuture(result: [])
120+
}
121+
122+
// Step through the values, resolving or rejecting each Promise in the
123+
// loaded queue.
124+
_ = try batchLoadFunction(keys)
125+
.thenThrowing { values in
126+
if values.count != keys.count {
127+
throw DataLoaderError.typeError("The function did not return an array of the same length as the array of keys. \nKeys count: \(keys.count)\nValues count: \(values.count)")
128+
}
129+
130+
for entry in queue.enumerated() {
131+
let result = values[entry.offset]
132+
133+
switch result {
134+
case .failure(let error): entry.element.promise.fail(error: error)
135+
case .success(let value): entry.element.promise.succeed(result: value)
136+
}
137+
}
138+
}
139+
.mapIfError{ error in
140+
self.failedDispatch(queue: queue, error: error)
141+
}
142+
}
143+
144+
public func dispatchQueue(on eventLoop: EventLoopGroup) throws {
145+
// Take the current loader queue, replacing it with an empty queue.
146+
let queue = self.queue
147+
self.queue = []
148+
149+
// If a maxBatchSize was provided and the queue is longer, then segment the
150+
// queue into multiple batches, otherwise treat the queue as a single batch.
151+
if let maxBatchSize = options.maxBatchSize, maxBatchSize > 0 && maxBatchSize < queue.count {
152+
for i in 0...(queue.count / maxBatchSize) {
153+
let startIndex = i * maxBatchSize
154+
let endIndex = (i + 1) * maxBatchSize
155+
let slicedQueue = queue[startIndex..<min(endIndex, queue.count)]
156+
try dispatchQueueBatch(queue: Array(slicedQueue), on: eventLoop)
157+
}
158+
} else {
159+
try dispatchQueueBatch(queue: queue, on: eventLoop)
160+
}
161+
}
162+
163+
private func failedDispatch(queue: LoaderQueue<Key, Value>, error: Error) {
164+
queue.forEach { (key, promise) in
165+
_ = clear(key: key)
166+
promise.fail(error: error)
167+
}
168+
}
169+
}

‎Sources/SwiftDataLoader/DataLoaderError.swift

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
//
2+
// DataLoaderError.swift
3+
// CNIOAtomics
4+
//
5+
// Created by Kim de Vos on 02/06/2018.
6+
//
7+
8+
import Foundation
9+
10+
public enum DataLoaderError: Error {
11+
case typeError(String)
12+
case noValueForKey(String)
13+
}

‎Sources/SwiftDataLoader/DataLoaderOptions.swift

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
//
2+
// DataLoaderOptions.swift
3+
// CNIOAtomics
4+
//
5+
// Created by Kim de Vos on 02/06/2018.
6+
//
7+
8+
public struct DataLoaderOptions<Key: Hashable, Value> {
9+
public let batchingEnabled: Bool
10+
public let cachingEnabled: Bool
11+
public let cacheMap: [Key: Value]
12+
public let maxBatchSize: Int?
13+
public let cacheKeyFunction: ((Key) -> Key)?
14+
15+
public init(batchingEnabled: Bool = true,
16+
cachingEnabled: Bool = true,
17+
maxBatchSize: Int? = nil,
18+
cacheMap: [Key: Value] = [:],
19+
cacheKeyFunction: ((Key) -> Key)? = nil) {
20+
self.batchingEnabled = batchingEnabled
21+
self.cachingEnabled = cachingEnabled
22+
self.maxBatchSize = maxBatchSize
23+
self.cacheMap = cacheMap
24+
self.cacheKeyFunction = cacheKeyFunction
25+
}
26+
}

‎Tests/LinuxMain.swift

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
import XCTest
2+
3+
import DataLoaderTests
4+
5+
var tests = [XCTestCaseEntry]()
6+
tests += DataLoaderTests.allTests()
7+
tests += DataLoaderAbuseTests.allTests()
8+
XCTMain(tests)

‎Tests/SwiftDataLoaderTests/DataLoaderAbuseTests.swift

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
//
2+
// DataLoaderAbuseTests.swift
3+
// DataLoaderTests
4+
//
5+
// Created by Kim de Vos on 03/06/2018.
6+
//
7+
8+
import XCTest
9+
import NIO
10+
11+
@testable import SwiftDataLoader
12+
13+
/// Provides descriptive error messages for API abuse
14+
class DataLoaderAbuseTests: XCTestCase {
15+
16+
func testFuntionWithNoValues() throws {
17+
let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
18+
defer {
19+
XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully())
20+
}
21+
22+
let identityLoader = DataLoader<Int, Int>(options: DataLoaderOptions(batchingEnabled: false)) { keys in
23+
eventLoopGroup.next().newSucceededFuture(result: [])
24+
}
25+
26+
let value = try identityLoader.load(key: 1, on: eventLoopGroup)
27+
28+
XCTAssertNoThrow(try identityLoader.dispatchQueue(on: eventLoopGroup))
29+
30+
XCTAssertThrowsError(try value.wait(), "Did not return value for key: 1")
31+
}
32+
33+
func testBatchFuntionMustPromiseAnArrayOfCorrectLength() throws {
34+
let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
35+
defer {
36+
XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully())
37+
}
38+
39+
let identityLoader = DataLoader<Int, Int>(options: DataLoaderOptions()) { keys in
40+
eventLoopGroup.next().newSucceededFuture(result: [])
41+
}
42+
43+
let value = try identityLoader.load(key: 1, on: eventLoopGroup)
44+
45+
XCTAssertNoThrow(try identityLoader.dispatchQueue(on: eventLoopGroup))
46+
47+
XCTAssertThrowsError(try value.wait(), "The function did not return an array of the same length as the array of keys. \nKeys count: 1\nValues count: 0")
48+
}
49+
50+
func testBatchFuntionWithSomeValues() throws {
51+
let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
52+
defer {
53+
XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully())
54+
}
55+
56+
let identityLoader = DataLoader<Int, Int>(options: DataLoaderOptions()) { keys in
57+
var results = [DataLoaderFutureValue<Int>]()
58+
59+
for key in keys {
60+
if key == 1 {
61+
results.append(DataLoaderFutureValue.success(key))
62+
} else {
63+
results.append(DataLoaderFutureValue.failure("Test error"))
64+
}
65+
}
66+
67+
return eventLoopGroup.next().newSucceededFuture(result: results)
68+
}
69+
70+
let value1 = try identityLoader.load(key: 1, on: eventLoopGroup)
71+
let value2 = try identityLoader.load(key: 2, on: eventLoopGroup)
72+
73+
XCTAssertNoThrow(try identityLoader.dispatchQueue(on: eventLoopGroup))
74+
75+
XCTAssertThrowsError(try value2.wait())
76+
77+
XCTAssertTrue(try value1.wait() == 1)
78+
}
79+
80+
func testFuntionWithSomeValues() throws {
81+
let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
82+
defer {
83+
XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully())
84+
}
85+
86+
let identityLoader = DataLoader<Int, Int>(options: DataLoaderOptions(batchingEnabled: false)) { keys in
87+
var results = [DataLoaderFutureValue<Int>]()
88+
89+
for key in keys {
90+
if key == 1 {
91+
results.append(DataLoaderFutureValue.success(key))
92+
} else {
93+
results.append(DataLoaderFutureValue.failure("Test error"))
94+
}
95+
}
96+
97+
return eventLoopGroup.next().newSucceededFuture(result: results)
98+
}
99+
100+
let value1 = try identityLoader.load(key: 1, on: eventLoopGroup)
101+
let value2 = try identityLoader.load(key: 2, on: eventLoopGroup)
102+
103+
XCTAssertNoThrow(try identityLoader.dispatchQueue(on: eventLoopGroup))
104+
105+
XCTAssertThrowsError(try value2.wait())
106+
107+
XCTAssertTrue(try value1.wait() == 1)
108+
}
109+
110+
static var allTests = [
111+
("testFuntionWithNoValues", testFuntionWithNoValues),
112+
("testBatchFuntionMustPromiseAnArrayOfCorrectLength", testBatchFuntionMustPromiseAnArrayOfCorrectLength),
113+
("testBatchFuntionWithSomeValues", testBatchFuntionWithSomeValues),
114+
("testFuntionWithSomeValues", testFuntionWithSomeValues)
115+
]
116+
117+
}
118+
119+
extension String: Error { }

0 commit comments

Comments
 (0)