Skip to content

Commit 5c27005

Browse files
Merge pull request #15 from d-exclaimation/async-await
Adding async await to DataLoader
2 parents b235b88 + 6a9ddee commit 5c27005

File tree

3 files changed

+183
-3
lines changed

3 files changed

+183
-3
lines changed

.github/workflows/test.yml

+18-3
Original file line numberDiff line numberDiff line change
@@ -6,15 +6,30 @@ on:
66
pull_request:
77
branches: [ master ]
88
jobs:
9-
build:
9+
linux-build:
1010
name: Build and test on ${{ matrix.os }}
1111
runs-on: ${{ matrix.os }}
1212
strategy:
1313
matrix:
14-
os: [macos-latest, ubuntu-latest]
14+
os: [ubuntu-latest]
1515
steps:
1616
- uses: actions/checkout@v2
17-
- uses: fwal/setup-swift@v1
17+
- uses: swift-actions/setup-swift@v1
18+
- name: Build
19+
run: swift build
20+
- name: Run tests
21+
run: swift test
22+
macos-build:
23+
name: Build and test on ${{ matrix.os }}
24+
runs-on: ${{ matrix.os }}
25+
strategy:
26+
matrix:
27+
os: [macos-latest]
28+
steps:
29+
- uses: actions/checkout@v2
30+
- uses: maxim-lobanov/setup-xcode@v1
31+
with:
32+
xcode-version: latest-stable
1833
- name: Build
1934
run: swift build
2035
- name: Run tests

Sources/DataLoader/DataLoader.swift

+48
Original file line numberDiff line numberDiff line change
@@ -217,3 +217,51 @@ final public class DataLoader<Key: Hashable, Value> {
217217
}
218218
}
219219
}
220+
221+
#if compiler(>=5.5) && canImport(_Concurrency)
222+
223+
/// Batch load function using async await
224+
public typealias ConcurrentBatchLoadFunction<Key, Value> = @Sendable (_ keys: [Key]) async throws -> [DataLoaderFutureValue<Value>]
225+
226+
public extension DataLoader {
227+
@available(macOS 12, iOS 15, watchOS 8, tvOS 15, *)
228+
convenience init(
229+
on eventLoop: EventLoop,
230+
options: DataLoaderOptions<Key, Value> = DataLoaderOptions(),
231+
throwing asyncThrowingLoadFunction: @escaping ConcurrentBatchLoadFunction<Key, Value>
232+
) {
233+
self.init(options: options, batchLoadFunction: { keys in
234+
let promise = eventLoop.next().makePromise(of: [DataLoaderFutureValue<Value>].self)
235+
promise.completeWithTask {
236+
try await asyncThrowingLoadFunction(keys)
237+
}
238+
return promise.futureResult
239+
})
240+
}
241+
242+
/// Asynchronously loads a key, returning the value represented by that key.
243+
@available(macOS 12, iOS 15, watchOS 8, tvOS 15, *)
244+
func load(key: Key, on eventLoopGroup: EventLoopGroup) async throws -> Value {
245+
try await load(key: key, on: eventLoopGroup).get()
246+
}
247+
248+
/// Asynchronously loads multiple keys, promising an array of values:
249+
///
250+
/// ```
251+
/// let aAndB = try await myLoader.loadMany(keys: [ "a", "b" ], on: eventLoopGroup)
252+
/// ```
253+
///
254+
/// This is equivalent to the more verbose:
255+
///
256+
/// ```
257+
/// async let a = myLoader.load(key: "a", on: eventLoopGroup)
258+
/// async let b = myLoader.load(key: "b", on: eventLoopGroup)
259+
/// let aAndB = try await a + b
260+
/// ```
261+
@available(macOS 12, iOS 15, watchOS 8, tvOS 15, *)
262+
func loadMany(keys: [Key], on eventLoopGroup: EventLoopGroup) async throws -> [Value] {
263+
try await loadMany(keys: keys, on: eventLoopGroup).get()
264+
}
265+
}
266+
267+
#endif
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
import XCTest
2+
import NIO
3+
4+
@testable import DataLoader
5+
6+
#if compiler(>=5.5) && canImport(_Concurrency)
7+
8+
@available(macOS 12, iOS 15, watchOS 8, tvOS 15, *)
9+
actor Concurrent<T> {
10+
var wrappedValue: T
11+
12+
func nonmutating<Returned>(_ action: (T) throws -> Returned) async rethrows -> Returned {
13+
try action(wrappedValue)
14+
}
15+
16+
func mutating<Returned>(_ action: (inout T) throws -> Returned) async rethrows -> Returned {
17+
try action(&wrappedValue)
18+
}
19+
20+
init(_ value: T) {
21+
self.wrappedValue = value
22+
}
23+
}
24+
25+
26+
/// Primary API
27+
@available(macOS 12, iOS 15, watchOS 8, tvOS 15, *)
28+
final class DataLoaderAsyncTests: XCTestCase {
29+
30+
/// Builds a really really simple data loader with async await
31+
func testReallyReallySimpleDataLoader() async throws {
32+
let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
33+
defer {
34+
XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully())
35+
}
36+
37+
let identityLoader = DataLoader<Int, Int>(
38+
on: eventLoopGroup.next(),
39+
options: DataLoaderOptions(batchingEnabled: false)
40+
) { keys async in
41+
let task = Task {
42+
keys.map { DataLoaderFutureValue.success($0) }
43+
}
44+
return await task.value
45+
}
46+
47+
let value = try await identityLoader.load(key: 1, on: eventLoopGroup)
48+
49+
XCTAssertEqual(value, 1)
50+
}
51+
52+
/// Supports loading multiple keys in one call
53+
func testLoadingMultipleKeys() async throws {
54+
let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
55+
defer {
56+
XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully())
57+
}
58+
59+
let identityLoader = DataLoader<Int, Int>(on: eventLoopGroup.next()) { keys in
60+
let task = Task {
61+
keys.map { DataLoaderFutureValue.success($0) }
62+
}
63+
return await task.value
64+
}
65+
66+
let values = try await identityLoader.loadMany(keys: [1, 2], on: eventLoopGroup)
67+
68+
XCTAssertEqual(values, [1,2])
69+
70+
let empty = try await identityLoader.loadMany(keys: [], on: eventLoopGroup)
71+
72+
XCTAssertTrue(empty.isEmpty)
73+
}
74+
75+
// Batches multiple requests
76+
func testMultipleRequests() async throws {
77+
let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
78+
defer {
79+
XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully())
80+
}
81+
82+
let loadCalls = Concurrent<[[Int]]>([])
83+
84+
let identityLoader = DataLoader<Int, Int>(
85+
on: eventLoopGroup.next(),
86+
options: DataLoaderOptions(
87+
batchingEnabled: true,
88+
executionPeriod: nil
89+
)
90+
) { keys in
91+
await loadCalls.mutating { $0.append(keys) }
92+
let task = Task {
93+
keys.map { DataLoaderFutureValue.success($0) }
94+
}
95+
return await task.value
96+
}
97+
98+
async let value1 = identityLoader.load(key: 1, on: eventLoopGroup)
99+
async let value2 = identityLoader.load(key: 2, on: eventLoopGroup)
100+
101+
/// Have to wait for a split second because Tasks may not be executed before this statement
102+
try await Task.sleep(nanoseconds: 500_000_000)
103+
104+
XCTAssertNoThrow(try identityLoader.execute())
105+
106+
let result1 = try await value1
107+
XCTAssertEqual(result1, 1)
108+
let result2 = try await value2
109+
XCTAssertEqual(result2, 2)
110+
111+
let calls = await loadCalls.wrappedValue
112+
XCTAssertEqual(calls.count, 1)
113+
XCTAssertEqual(calls.map { $0.sorted() }, [[1, 2]])
114+
}
115+
}
116+
117+
#endif

0 commit comments

Comments
 (0)