diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 0e5e3cd..5887173 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -61,7 +61,7 @@ jobs: runs-on: ubuntu-22.04 strategy: matrix: - swift: ["5.7", "5.8", "5.9", "5.10"] + swift: ["5.9", "5.10"] steps: - uses: swift-actions/setup-swift@v2 with: diff --git a/Package.resolved b/Package.resolved index 55f18b2..7b8fae1 100644 --- a/Package.resolved +++ b/Package.resolved @@ -1,16 +1,68 @@ { - "object": { - "pins": [ - { - "package": "swift-nio", - "repositoryURL": "https://github.com/apple/swift-nio.git", - "state": { - "branch": null, - "revision": "51c3fc2e4a0fcdf4a25089b288dd65b73df1b0ef", - "version": "2.37.0" - } + "pins" : [ + { + "identity" : "async-collections", + "kind" : "remoteSourceControl", + "location" : "https://github.com/adam-fowler/async-collections", + "state" : { + "revision" : "726af96095a19df6b8053ddbaed0a727aa70ccb2", + "version" : "0.1.0" } - ] - }, - "version": 1 + }, + { + "identity" : "swift-algorithms", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-algorithms.git", + "state" : { + "revision" : "f6919dfc309e7f1b56224378b11e28bab5bccc42", + "version" : "1.2.0" + } + }, + { + "identity" : "swift-atomics", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-atomics.git", + "state" : { + "revision" : "cd142fd2f64be2100422d658e7411e39489da985", + "version" : "1.2.0" + } + }, + { + "identity" : "swift-collections", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-collections.git", + "state" : { + "revision" : "671108c96644956dddcd89dd59c203dcdb36cec7", + "version" : "1.1.4" + } + }, + { + "identity" : "swift-nio", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-nio.git", + "state" : { + "revision" : "f7dc3f527576c398709b017584392fb58592e7f5", + "version" : "2.75.0" + } + }, + { + "identity" : "swift-numerics", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-numerics.git", + "state" : { + "revision" : "0a5bc04095a675662cf24757cc0640aa2204253b", + "version" : "1.0.2" + } + }, + { + "identity" : "swift-system", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-system.git", + "state" : { + "revision" : "c8a44d836fe7913603e246acab7c528c2e780168", + "version" : "1.4.0" + } + } + ], + "version" : 2 } diff --git a/Package.swift b/Package.swift index bd7a98e..f67250b 100644 --- a/Package.swift +++ b/Package.swift @@ -1,19 +1,36 @@ -// swift-tools-version:5.0 +// swift-tools-version:5.9 // The swift-tools-version declares the minimum version of Swift required to build this package. import PackageDescription let package = Package( name: "DataLoader", + platforms: [.macOS(.v12), .iOS(.v15), .tvOS(.v15), .watchOS(.v8)], products: [ .library(name: "DataLoader", targets: ["DataLoader"]), + .library(name: "AsyncDataLoader", targets: ["AsyncDataLoader"]), ], dependencies: [ + .package(url: "https://github.com/apple/swift-algorithms.git", from: "1.0.0"), + .package(url: "https://github.com/adam-fowler/async-collections", from: "0.0.1"), .package(url: "https://github.com/apple/swift-nio.git", from: "2.0.0"), ], targets: [ - .target(name: "DataLoader", dependencies: ["NIO", "NIOConcurrencyHelpers"]), + .target( + name: "DataLoader", + dependencies: [ + .product(name: "NIO", package: "swift-nio"), + .product(name: "NIOConcurrencyHelpers", package: "swift-nio"), + ] + ), + .target( + name: "AsyncDataLoader", + dependencies: [ + .product(name: "Algorithms", package: "swift-algorithms"), + .product(name: "AsyncCollections", package: "async-collections"), + ] + ), .testTarget(name: "DataLoaderTests", dependencies: ["DataLoader"]), - ], - swiftLanguageVersions: [.v5] + .testTarget(name: "AsyncDataLoaderTests", dependencies: ["AsyncDataLoader"]), + ] ) diff --git a/Sources/AsyncDataLoader/Channel/Channel.swift b/Sources/AsyncDataLoader/Channel/Channel.swift new file mode 100644 index 0000000..0950bf9 --- /dev/null +++ b/Sources/AsyncDataLoader/Channel/Channel.swift @@ -0,0 +1,55 @@ +actor Channel: Sendable { + private var state = State() +} + +extension Channel { + @discardableResult + func fulfill(_ value: Success) async -> Bool { + if await state.result == nil { + await state.setResult(result: value) + + for waiters in await state.waiters { + waiters.resume(returning: value) + } + + await state.removeAllWaiters() + + return false + } + + return true + } + + @discardableResult + func fail(_ failure: Failure) async -> Bool { + if await state.failure == nil { + await state.setFailure(failure: failure) + + for waiters in await state.waiters { + waiters.resume(throwing: failure) + } + + await state.removeAllWaiters() + + return false + } + + return true + } + + var value: Success { + get async throws { + try await withCheckedThrowingContinuation { continuation in + Task { + if let result = await state.result { + continuation.resume(returning: result) + } else if let failure = await self.state.failure { + continuation.resume(throwing: failure) + } else { + await state.appendWaiters(waiters: continuation) + } + } + } + } + } +} diff --git a/Sources/AsyncDataLoader/Channel/State.swift b/Sources/AsyncDataLoader/Channel/State.swift new file mode 100644 index 0000000..ee5e784 --- /dev/null +++ b/Sources/AsyncDataLoader/Channel/State.swift @@ -0,0 +1,25 @@ +typealias Waiter = CheckedContinuation + +actor State { + var waiters = [Waiter]() + var result: Success? + var failure: Failure? +} + +extension State { + func setResult(result: Success) { + self.result = result + } + + func setFailure(failure: Failure) { + self.failure = failure + } + + func appendWaiters(waiters: Waiter...) { + self.waiters.append(contentsOf: waiters) + } + + func removeAllWaiters() { + waiters.removeAll() + } +} diff --git a/Sources/AsyncDataLoader/DataLoader.swift b/Sources/AsyncDataLoader/DataLoader.swift new file mode 100644 index 0000000..d60b829 --- /dev/null +++ b/Sources/AsyncDataLoader/DataLoader.swift @@ -0,0 +1,229 @@ +import Algorithms +import AsyncCollections + +public enum DataLoaderValue: Sendable { + case success(T) + case failure(Error) +} + +public typealias BatchLoadFunction = + @Sendable (_ keys: [Key]) async throws -> [DataLoaderValue] +private typealias LoaderQueue = [( + key: Key, + channel: Channel +)] + +/// DataLoader creates a public API for loading data from a particular +/// data back-end with unique keys such as the id column of a SQL table +/// or document name in a MongoDB database, given a batch loading function. +/// +/// Each DataLoader instance contains a unique memoized cache. Use caution +/// when used in long-lived applications or those which serve many users +/// with different access permissions and consider creating a new instance +/// per data request. +public actor DataLoader { + private let batchLoadFunction: BatchLoadFunction + private let options: DataLoaderOptions + + private var cache = [Key: Channel]() + private var queue = LoaderQueue() + + private var dispatchScheduled = false + + public init( + options: DataLoaderOptions = DataLoaderOptions(), + batchLoadFunction: @escaping BatchLoadFunction + ) { + self.options = options + self.batchLoadFunction = batchLoadFunction + } + + /// Loads a key, returning the value represented by that key. + public func load(key: Key) async throws -> Value { + let cacheKey = options.cacheKeyFunction?(key) ?? key + + if options.cachingEnabled, let cached = cache[cacheKey] { + return try await cached.value + } + + let channel = Channel() + + if options.batchingEnabled { + queue.append((key: key, channel: channel)) + + if let executionPeriod = options.executionPeriod, !dispatchScheduled { + Task.detached { + try await Task.sleep(nanoseconds: executionPeriod) + try await self.execute() + } + + dispatchScheduled = true + } + } else { + Task.detached { + do { + let results = try await self.batchLoadFunction([key]) + + if results.isEmpty { + await channel + .fail( + DataLoaderError + .noValueForKey("Did not return value for key: \(key)") + ) + } else { + let result = results[0] + + switch result { + case let .success(value): + await channel.fulfill(value) + case let .failure(error): + await channel.fail(error) + } + } + } catch { + await channel.fail(error) + } + } + } + + if options.cachingEnabled { + cache[cacheKey] = channel + } + + return try await channel.value + } + + /// Loads multiple keys, promising an array of values: + /// + /// ```swift + /// async let aAndB = try myLoader.loadMany(keys: [ "a", "b" ]) + /// ``` + /// + /// This is equivalent to the more verbose: + /// + /// ```swift + /// async let aAndB = [ + /// myLoader.load(key: "a"), + /// myLoader.load(key: "b") + /// ] + /// ``` + /// or + /// ```swift + /// async let a = myLoader.load(key: "a") + /// async let b = myLoader.load(key: "b") + /// ``` + public func loadMany(keys: [Key]) async throws -> [Value] { + guard !keys.isEmpty else { + return [] + } + + return try await keys.concurrentMap { try await self.load(key: $0) } + } + + /// Clears the value at `key` from the cache, if it exists. Returns itself for + /// method chaining. + @discardableResult + public func clear(key: Key) -> DataLoader { + let cacheKey = options.cacheKeyFunction?(key) ?? key + + cache.removeValue(forKey: cacheKey) + + return self + } + + /// Clears the entire cache. To be used when some event results in unknown + /// invalidations across this particular `DataLoader`. Returns itself for + /// method chaining. + @discardableResult + public func clearAll() -> DataLoader { + cache.removeAll() + + return self + } + + /// Adds the provied key and value to the cache. If the key already exists, no + /// change is made. Returns itself for method chaining. + @discardableResult + public func prime(key: Key, value: Value) async throws -> DataLoader { + let cacheKey = options.cacheKeyFunction?(key) ?? key + + if cache[cacheKey] == nil { + let channel = Channel() + + Task.detached { + await channel.fulfill(value) + } + + cache[cacheKey] = channel + } + + return self + } + + public func execute() async throws { + // Take the current loader queue, replacing it with an empty queue. + let batch = queue + + queue = [] + + if dispatchScheduled { + dispatchScheduled = false + } + + guard !batch.isEmpty else { + return () + } + + // If a maxBatchSize was provided and the queue is longer, then segment the + // queue into multiple batches, otherwise treat the queue as a single batch. + if let maxBatchSize = options.maxBatchSize, maxBatchSize > 0, maxBatchSize < batch.count { + try await batch.chunks(ofCount: maxBatchSize).asyncForEach { slicedBatch in + try await self.executeBatch(batch: Array(slicedBatch)) + } + } else { + try await executeBatch(batch: batch) + } + } + + private func executeBatch(batch: LoaderQueue) async throws { + let keys = batch.map { $0.key } + + if keys.isEmpty { + return + } + + // Step through the values, resolving or rejecting each Promise in the + // loaded queue. + do { + let values = try await batchLoadFunction(keys) + + if values.count != keys.count { + 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)" + ) + } + + for entry in batch.enumerated() { + let result = values[entry.offset] + + switch result { + case let .failure(error): + await entry.element.channel.fail(error) + case let .success(value): + await entry.element.channel.fulfill(value) + } + } + } catch { + await failedExecution(batch: batch, error: error) + } + } + + private func failedExecution(batch: LoaderQueue, error: Error) async { + for (key, channel) in batch { + _ = clear(key: key) + + await channel.fail(error) + } + } +} diff --git a/Sources/AsyncDataLoader/DataLoaderError.swift b/Sources/AsyncDataLoader/DataLoaderError.swift new file mode 100644 index 0000000..8366d04 --- /dev/null +++ b/Sources/AsyncDataLoader/DataLoaderError.swift @@ -0,0 +1,4 @@ +public enum DataLoaderError: Error { + case typeError(String) + case noValueForKey(String) +} diff --git a/Sources/AsyncDataLoader/DataLoaderOptions.swift b/Sources/AsyncDataLoader/DataLoaderOptions.swift new file mode 100644 index 0000000..bab2c56 --- /dev/null +++ b/Sources/AsyncDataLoader/DataLoaderOptions.swift @@ -0,0 +1,40 @@ +public struct DataLoaderOptions: Sendable { + /// Default `true`. Set to `false` to disable batching, invoking + /// `batchLoadFunction` with a single load key. This is + /// equivalent to setting `maxBatchSize` to `1`. + public let batchingEnabled: Bool + + /// Default `nil`. Limits the number of items that get passed in to the + /// `batchLoadFn`. May be set to `1` to disable batching. + public let maxBatchSize: Int? + + /// Default `true`. Set to `false` to disable memoization caching, creating a + /// new `EventLoopFuture` and new key in the `batchLoadFunction` + /// for every load of the same key. + public let cachingEnabled: Bool + + /// Default `2ms`. Defines the period of time that the DataLoader should + /// wait and collect its queue before executing. Faster times result + /// in smaller batches quicker resolution, slower times result in larger + /// batches but slower resolution. + /// This is irrelevant if batching is disabled. + public let executionPeriod: UInt64? + + /// Default `nil`. Produces cache key for a given load key. Useful + /// when objects are keys and two objects should be considered equivalent. + public let cacheKeyFunction: (@Sendable (Key) -> Key)? + + public init( + batchingEnabled: Bool = true, + cachingEnabled: Bool = true, + maxBatchSize: Int? = nil, + executionPeriod: UInt64? = 2_000_000, + cacheKeyFunction: (@Sendable (Key) -> Key)? = nil + ) { + self.batchingEnabled = batchingEnabled + self.cachingEnabled = cachingEnabled + self.executionPeriod = executionPeriod + self.maxBatchSize = maxBatchSize + self.cacheKeyFunction = cacheKeyFunction + } +} diff --git a/Sources/DataLoader/DataLoader.swift b/Sources/DataLoader/DataLoader.swift index dce6ce5..03374a2 100644 --- a/Sources/DataLoader/DataLoader.swift +++ b/Sources/DataLoader/DataLoader.swift @@ -26,7 +26,7 @@ public final class DataLoader { private var queue = LoaderQueue() private var dispatchScheduled = false - private let lock = Lock() + private let lock = NIOLock() public init( options: DataLoaderOptions = DataLoaderOptions(), diff --git a/Tests/AsyncDataLoaderTests/DataLoaderAbuseTests.swift b/Tests/AsyncDataLoaderTests/DataLoaderAbuseTests.swift new file mode 100644 index 0000000..d6b6b9b --- /dev/null +++ b/Tests/AsyncDataLoaderTests/DataLoaderAbuseTests.swift @@ -0,0 +1,114 @@ +import XCTest + +@testable import AsyncDataLoader + +/// Provides descriptive error messages for API abuse +class DataLoaderAbuseTests: XCTestCase { + func testFuntionWithNoValues() async throws { + let identityLoader = DataLoader( + options: DataLoaderOptions(batchingEnabled: false) + ) { _ in + [] + } + + async let value = identityLoader.load(key: 1) + + var didFailWithError: Error? + + do { + _ = try await value + } catch { + didFailWithError = error + } + + XCTAssertNotNil(didFailWithError) + } + + func testBatchFuntionMustPromiseAnArrayOfCorrectLength() async { + let identityLoader = DataLoader() { _ in + [] + } + + async let value = identityLoader.load(key: 1) + + var didFailWithError: Error? + + do { + _ = try await value + } catch { + didFailWithError = error + } + + XCTAssertNotNil(didFailWithError) + } + + func testBatchFuntionWithSomeValues() async throws { + let identityLoader = DataLoader() { keys in + var results = [DataLoaderValue]() + + for key in keys { + if key == 1 { + results.append(.success(key)) + } else { + results.append(.failure("Test error")) + } + } + + return results + } + + async let value1 = identityLoader.load(key: 1) + async let value2 = identityLoader.load(key: 2) + + var didFailWithError: Error? + + do { + _ = try await value2 + } catch { + didFailWithError = error + } + + XCTAssertNotNil(didFailWithError) + + let value = try await value1 + + XCTAssertTrue(value == 1) + } + + func testFuntionWithSomeValues() async throws { + let identityLoader = DataLoader( + options: DataLoaderOptions(batchingEnabled: false) + ) { keys in + var results = [DataLoaderValue]() + + for key in keys { + if key == 1 { + results.append(.success(key)) + } else { + results.append(.failure("Test error")) + } + } + + return results + } + + async let value1 = identityLoader.load(key: 1) + async let value2 = identityLoader.load(key: 2) + + var didFailWithError: Error? + + do { + _ = try await value2 + } catch { + didFailWithError = error + } + + XCTAssertNotNil(didFailWithError) + + let value = try await value1 + + XCTAssertTrue(value == 1) + } +} + +extension String: Swift.Error {} diff --git a/Tests/AsyncDataLoaderTests/DataLoaderTests.swift b/Tests/AsyncDataLoaderTests/DataLoaderTests.swift new file mode 100644 index 0000000..5987b14 --- /dev/null +++ b/Tests/AsyncDataLoaderTests/DataLoaderTests.swift @@ -0,0 +1,663 @@ +import XCTest + +@testable import AsyncDataLoader + +let sleepConstant = UInt64(2_000_000) + +actor Concurrent { + var wrappedValue: T + + func nonmutating(_ action: (T) throws -> Returned) async rethrows -> Returned { + try action(wrappedValue) + } + + func mutating(_ action: (inout T) throws -> Returned) async rethrows -> Returned { + try action(&wrappedValue) + } + + init(_ value: T) { + wrappedValue = value + } +} + +/// Primary API +/// The `try await Task.sleep(nanoseconds: 2_000_000)` introduces a small delay to simulate +/// asynchronous behavior and ensure that concurrent requests (`value1`, `value2`...) +/// are grouped into a single batch for processing, as intended by the batching settings. +final class DataLoaderTests: XCTestCase { + /// Builds a really really simple data loader' + func testReallyReallySimpleDataLoader() async throws { + let identityLoader = DataLoader( + options: DataLoaderOptions(batchingEnabled: false) + ) { keys in + keys.map { DataLoaderValue.success($0) } + } + + let value = try await identityLoader.load(key: 1) + + XCTAssertEqual(value, 1) + } + + /// Supports loading multiple keys in one call + func testLoadingMultipleKeys() async throws { + let identityLoader = DataLoader() { keys in + keys.map { DataLoaderValue.success($0) } + } + + let values = try await identityLoader.loadMany(keys: [1, 2]) + + XCTAssertEqual(values, [1, 2]) + + let empty = try await identityLoader.loadMany(keys: []) + + XCTAssertTrue(empty.isEmpty) + } + + // Batches multiple requests + func testMultipleRequests() async throws { + let loadCalls = Concurrent<[[Int]]>([]) + + let identityLoader = DataLoader( + options: DataLoaderOptions( + batchingEnabled: true, + executionPeriod: nil + ) + ) { keys in + await loadCalls.mutating { $0.append(keys) } + + return keys.map { DataLoaderValue.success($0) } + } + + async let value1 = identityLoader.load(key: 1) + async let value2 = identityLoader.load(key: 2) + + try await Task.sleep(nanoseconds: sleepConstant) + + var didFailWithError: Error? + + do { + _ = try await identityLoader.execute() + } catch { + didFailWithError = error + } + + XCTAssertNil(didFailWithError) + + let result1 = try await value1 + let result2 = try await value2 + + XCTAssertEqual(result1, 1) + XCTAssertEqual(result2, 2) + + let calls = await loadCalls.wrappedValue + + XCTAssertEqual(calls.map { $0.sorted() }, [[1, 2]]) + } + + /// Batches multiple requests with max batch sizes + func testMultipleRequestsWithMaxBatchSize() async throws { + let loadCalls = Concurrent<[[Int]]>([]) + + let identityLoader = DataLoader( + options: DataLoaderOptions( + batchingEnabled: true, + maxBatchSize: 2, + executionPeriod: nil + ) + ) { keys in + await loadCalls.mutating { $0.append(keys) } + + return keys.map { DataLoaderValue.success($0) } + } + + async let value1 = identityLoader.load(key: 1) + async let value2 = identityLoader.load(key: 2) + async let value3 = identityLoader.load(key: 3) + + try await Task.sleep(nanoseconds: sleepConstant) + + var didFailWithError: Error? + + do { + _ = try await identityLoader.execute() + } catch { + didFailWithError = error + } + + XCTAssertNil(didFailWithError) + + let result1 = try await value1 + let result2 = try await value2 + let result3 = try await value3 + + XCTAssertEqual(result1, 1) + XCTAssertEqual(result2, 2) + XCTAssertEqual(result3, 3) + + let calls = await loadCalls.wrappedValue + + XCTAssertEqual(calls.first?.count, 2) + XCTAssertEqual(calls.last?.count, 1) + } + + /// Coalesces identical requests + func testCoalescesIdenticalRequests() async throws { + let loadCalls = Concurrent<[[Int]]>([]) + + let identityLoader = DataLoader( + options: DataLoaderOptions(executionPeriod: nil) + ) { keys in + await loadCalls.mutating { $0.append(keys) } + + return keys.map { DataLoaderValue.success($0) } + } + + async let value1 = identityLoader.load(key: 1) + async let value2 = identityLoader.load(key: 1) + + try await Task.sleep(nanoseconds: sleepConstant) + + var didFailWithError: Error? + + do { + _ = try await identityLoader.execute() + } catch { + didFailWithError = error + } + + XCTAssertNil(didFailWithError) + + let result1 = try await value1 + let result2 = try await value2 + + XCTAssertTrue(result1 == 1) + XCTAssertTrue(result2 == 1) + + let calls = await loadCalls.wrappedValue + + XCTAssertTrue(calls.map { $0.sorted() } == [[1]]) + } + + // Caches repeated requests + func testCachesRepeatedRequests() async throws { + let loadCalls = Concurrent<[[String]]>([]) + + let identityLoader = DataLoader( + options: DataLoaderOptions(executionPeriod: nil) + ) { keys in + await loadCalls.mutating { $0.append(keys) } + + return keys.map { DataLoaderValue.success($0) } + } + + async let value1 = identityLoader.load(key: "A") + async let value2 = identityLoader.load(key: "B") + + try await Task.sleep(nanoseconds: sleepConstant) + + var didFailWithError: Error? + + do { + _ = try await identityLoader.execute() + } catch { + didFailWithError = error + } + + XCTAssertNil(didFailWithError) + + let result1 = try await value1 + let result2 = try await value2 + + XCTAssertTrue(result1 == "A") + XCTAssertTrue(result2 == "B") + + let calls = await loadCalls.wrappedValue + + XCTAssertTrue(calls.map { $0.sorted() } == [["A", "B"]]) + + async let value3 = identityLoader.load(key: "A") + async let value4 = identityLoader.load(key: "C") + + try await Task.sleep(nanoseconds: sleepConstant) + + var didFailWithError2: Error? + + do { + _ = try await identityLoader.execute() + } catch { + didFailWithError2 = error + } + + XCTAssertNil(didFailWithError2) + + let result3 = try await value3 + let result4 = try await value4 + + XCTAssertTrue(result3 == "A") + XCTAssertTrue(result4 == "C") + + let calls2 = await loadCalls.wrappedValue + + XCTAssertTrue(calls2.map { $0.sorted() } == [["A", "B"], ["C"]]) + + async let value5 = identityLoader.load(key: "A") + async let value6 = identityLoader.load(key: "B") + async let value7 = identityLoader.load(key: "C") + + try await Task.sleep(nanoseconds: sleepConstant) + + var didFailWithError3: Error? + + do { + _ = try await identityLoader.execute() + } catch { + didFailWithError3 = error + } + + XCTAssertNil(didFailWithError3) + + let result5 = try await value5 + let result6 = try await value6 + let result7 = try await value7 + + XCTAssertTrue(result5 == "A") + XCTAssertTrue(result6 == "B") + XCTAssertTrue(result7 == "C") + + let calls3 = await loadCalls.wrappedValue + + XCTAssertTrue(calls3.map { $0.sorted() } == [["A", "B"], ["C"]]) + } + + /// Clears single value in loader + func testClearSingleValueLoader() async throws { + let loadCalls = Concurrent<[[String]]>([]) + + let identityLoader = DataLoader( + options: DataLoaderOptions(executionPeriod: nil) + ) { keys in + await loadCalls.mutating { $0.append(keys) } + + return keys.map { DataLoaderValue.success($0) } + } + + async let value1 = identityLoader.load(key: "A") + async let value2 = identityLoader.load(key: "B") + + try await Task.sleep(nanoseconds: sleepConstant) + + var didFailWithError: Error? + + do { + _ = try await identityLoader.execute() + } catch { + didFailWithError = error + } + + XCTAssertNil(didFailWithError) + + let result1 = try await value1 + let result2 = try await value2 + + XCTAssertTrue(result1 == "A") + XCTAssertTrue(result2 == "B") + + let calls = await loadCalls.wrappedValue + + XCTAssertTrue(calls.map { $0.sorted() } == [["A", "B"]]) + + await identityLoader.clear(key: "A") + + async let value3 = identityLoader.load(key: "A") + async let value4 = identityLoader.load(key: "B") + + try await Task.sleep(nanoseconds: sleepConstant) + + var didFailWithError2: Error? + + do { + _ = try await identityLoader.execute() + } catch { + didFailWithError2 = error + } + + XCTAssertNil(didFailWithError2) + + let result3 = try await value3 + let result4 = try await value4 + + XCTAssertTrue(result3 == "A") + XCTAssertTrue(result4 == "B") + + let calls2 = await loadCalls.wrappedValue + + XCTAssertTrue(calls2.map { $0.sorted() } == [["A", "B"], ["A"]]) + } + + /// Clears all values in loader + func testClearsAllValuesInLoader() async throws { + let loadCalls = Concurrent<[[String]]>([]) + + let identityLoader = DataLoader( + options: DataLoaderOptions(executionPeriod: nil) + ) { keys in + await loadCalls.mutating { $0.append(keys) } + + return keys.map { DataLoaderValue.success($0) } + } + + async let value1 = identityLoader.load(key: "A") + async let value2 = identityLoader.load(key: "B") + + try await Task.sleep(nanoseconds: sleepConstant) + + var didFailWithError: Error? + + do { + _ = try await identityLoader.execute() + } catch { + didFailWithError = error + } + + XCTAssertNil(didFailWithError) + + let result1 = try await value1 + let result2 = try await value2 + + XCTAssertTrue(result1 == "A") + XCTAssertTrue(result2 == "B") + + let calls = await loadCalls.wrappedValue + + XCTAssertTrue(calls.map { $0.sorted() } == [["A", "B"]]) + + await identityLoader.clearAll() + + async let value3 = identityLoader.load(key: "A") + async let value4 = identityLoader.load(key: "B") + + try await Task.sleep(nanoseconds: sleepConstant) + + var didFailWithError2: Error? + + do { + _ = try await identityLoader.execute() + } catch { + didFailWithError2 = error + } + + XCTAssertNil(didFailWithError2) + + let result3 = try await value3 + let result4 = try await value4 + + XCTAssertTrue(result3 == "A") + XCTAssertTrue(result4 == "B") + + let calls2 = await loadCalls.wrappedValue + + XCTAssertTrue(calls2.map { $0.sorted() } == [["A", "B"], ["A", "B"]]) + } + + // Allows priming the cache + func testAllowsPrimingTheCache() async throws { + let loadCalls = Concurrent<[[String]]>([]) + + let identityLoader = DataLoader( + options: DataLoaderOptions(executionPeriod: nil) + ) { keys in + await loadCalls.mutating { $0.append(keys) } + + return keys.map { DataLoaderValue.success($0) } + } + + try await identityLoader.prime(key: "A", value: "A") + + async let value1 = identityLoader.load(key: "A") + async let value2 = identityLoader.load(key: "B") + + try await Task.sleep(nanoseconds: sleepConstant) + + var didFailWithError: Error? + + do { + _ = try await identityLoader.execute() + } catch { + didFailWithError = error + } + + XCTAssertNil(didFailWithError) + + let result1 = try await value1 + let result2 = try await value2 + + XCTAssertTrue(result1 == "A") + XCTAssertTrue(result2 == "B") + + let calls = await loadCalls.wrappedValue + + XCTAssertTrue(calls.map { $0.sorted() } == [["B"]]) + } + + /// Does not prime keys that already exist + func testDoesNotPrimeKeysThatAlreadyExist() async throws { + let loadCalls = Concurrent<[[String]]>([]) + + let identityLoader = DataLoader( + options: DataLoaderOptions(executionPeriod: nil) + ) { keys in + await loadCalls.mutating { $0.append(keys) } + + return keys.map { DataLoaderValue.success($0) } + } + + try await identityLoader.prime(key: "A", value: "X") + + async let value1 = identityLoader.load(key: "A") + async let value2 = identityLoader.load(key: "B") + + try await Task.sleep(nanoseconds: sleepConstant) + + var didFailWithError: Error? + + do { + _ = try await identityLoader.execute() + } catch { + didFailWithError = error + } + + XCTAssertNil(didFailWithError) + + let result1 = try await value1 + let result2 = try await value2 + + XCTAssertTrue(result1 == "X") + XCTAssertTrue(result2 == "B") + + try await identityLoader.prime(key: "A", value: "Y") + try await identityLoader.prime(key: "B", value: "Y") + + async let value3 = identityLoader.load(key: "A") + async let value4 = identityLoader.load(key: "B") + + try await Task.sleep(nanoseconds: sleepConstant) + + var didFailWithError2: Error? + + do { + _ = try await identityLoader.execute() + } catch { + didFailWithError2 = error + } + + XCTAssertNil(didFailWithError2) + + let result3 = try await value3 + let result4 = try await value4 + + XCTAssertTrue(result3 == "X") + XCTAssertTrue(result4 == "B") + + let calls = await loadCalls.wrappedValue + + XCTAssertTrue(calls.map { $0.sorted() } == [["B"]]) + } + + /// Allows forcefully priming the cache + func testAllowsForcefullyPrimingTheCache() async throws { + let loadCalls = Concurrent<[[String]]>([]) + + let identityLoader = DataLoader( + options: DataLoaderOptions(executionPeriod: nil) + ) { keys in + await loadCalls.mutating { $0.append(keys) } + + return keys.map { DataLoaderValue.success($0) } + } + + try await identityLoader.prime(key: "A", value: "X") + + async let value1 = identityLoader.load(key: "A") + async let value2 = identityLoader.load(key: "B") + + try await Task.sleep(nanoseconds: sleepConstant) + + var didFailWithError: Error? + + do { + _ = try await identityLoader.execute() + } catch { + didFailWithError = error + } + + XCTAssertNil(didFailWithError) + + let result1 = try await value1 + let result2 = try await value2 + + XCTAssertTrue(result1 == "X") + XCTAssertTrue(result2 == "B") + + try await identityLoader.clear(key: "A").prime(key: "A", value: "Y") + try await identityLoader.clear(key: "B").prime(key: "B", value: "Y") + + async let value3 = identityLoader.load(key: "A") + async let value4 = identityLoader.load(key: "B") + + try await Task.sleep(nanoseconds: sleepConstant) + + var didFailWithError2: Error? + + do { + _ = try await identityLoader.execute() + } catch { + didFailWithError2 = error + } + + XCTAssertNil(didFailWithError2) + + let result3 = try await value3 + let result4 = try await value4 + + XCTAssertTrue(result3 == "Y") + XCTAssertTrue(result4 == "Y") + + let calls = await loadCalls.wrappedValue + + XCTAssertTrue(calls.map { $0.sorted() } == [["B"]]) + } + + func testAutoExecute() async throws { + let identityLoader = DataLoader( + options: DataLoaderOptions(executionPeriod: sleepConstant) + ) { keys in + + keys.map { DataLoaderValue.success($0) } + } + + async let value = identityLoader.load(key: "A") + + // Don't manually call execute, but wait for more than 2ms + usleep(3000) + + let result = try await value + + XCTAssertNotNil(result) + } + + func testErrorResult() async throws { + let loaderErrorMessage = "TEST" + + // Test throwing loader without auto-executing + let throwLoader = DataLoader( + options: DataLoaderOptions(executionPeriod: nil) + ) { _ in + throw DataLoaderError.typeError(loaderErrorMessage) + } + + async let value = throwLoader.load(key: 1) + + try await Task.sleep(nanoseconds: sleepConstant) + + var didFailWithError: DataLoaderError? + + do { + _ = try await throwLoader.execute() + } catch { + didFailWithError = error as? DataLoaderError + } + + XCTAssertNil(didFailWithError) + + var didFailWithError2: DataLoaderError? + + do { + _ = try await value + } catch { + didFailWithError2 = error as? DataLoaderError + } + + var didFailWithErrorText2 = "" + + switch didFailWithError2 { + case let .typeError(text): + didFailWithErrorText2 = text + case .noValueForKey: + break + case .none: + break + } + + XCTAssertEqual(didFailWithErrorText2, loaderErrorMessage) + + // Test throwing loader with auto-executing + let throwLoaderAutoExecute = DataLoader( + options: DataLoaderOptions() + ) { _ in + throw DataLoaderError.typeError(loaderErrorMessage) + } + + async let valueAutoExecute = throwLoaderAutoExecute.load(key: 1) + + var didFailWithError3: DataLoaderError? + + do { + _ = try await valueAutoExecute + } catch { + didFailWithError3 = error as? DataLoaderError + } + + var didFailWithErrorText3 = "" + + switch didFailWithError3 { + case let .typeError(text): + didFailWithErrorText3 = text + case .noValueForKey: + break + case .none: + break + } + + XCTAssertEqual(didFailWithErrorText3, loaderErrorMessage) + } +}