Adding async await to DataLoader

.github/workflows/test.yml

@@ -6,15 +6,30 @@ on:
   pull_request:
     branches: [ master ]
 jobs:
-  build:
+  linux-build:
     name: Build and test on ${{ matrix.os }}
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        os: [macos-latest, ubuntu-latest]
+        os: [ubuntu-latest]
     steps:
       - uses: actions/checkout@v2
-      - uses: fwal/setup-swift@v1
+      - uses: swift-actions/setup-swift@v1
+      - name: Build
+        run: swift build
+      - name: Run tests
+        run: swift test
+  macos-build:
+    name: Build and test on ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [macos-latest]
+    steps:
+      - uses: actions/checkout@v2
+      - uses: maxim-lobanov/setup-xcode@v1
+        with:
+          xcode-version: latest-stable
       - name: Build
         run: swift build
       - name: Run tests

Sources/DataLoader/DataLoader.swift

@@ -217,3 +217,51 @@ final public class DataLoader<Key: Hashable, Value> {
         }
     }
 }
+
+#if compiler(>=5.5) && canImport(_Concurrency)
+
+/// Batch load function using async await 
+public typealias ConcurrentBatchLoadFunction<Key, Value> = @Sendable (_ keys: [Key]) async throws -> [DataLoaderFutureValue<Value>]
+
+public extension DataLoader {
+    @available(macOS 12, iOS 15, watchOS 8, tvOS 15, *)
+    convenience init(
+        on eventLoop: EventLoop,
+        options: DataLoaderOptions<Key, Value> = DataLoaderOptions(),
+        throwing asyncThrowingLoadFunction: @escaping ConcurrentBatchLoadFunction<Key, Value>
+    ) {
+        self.init(options: options, batchLoadFunction: { keys in
+            let promise = eventLoop.next().makePromise(of: [DataLoaderFutureValue<Value>].self)
+            promise.completeWithTask {
+                try await asyncThrowingLoadFunction(keys)
+            }
+            return promise.futureResult
+        })
+    }
+
+    /// Asynchronously loads a key, returning the value represented by that key.
+    @available(macOS 12, iOS 15, watchOS 8, tvOS 15, *)
+    func load(key: Key, on eventLoopGroup: EventLoopGroup) async throws -> Value {
+        try await load(key: key, on: eventLoopGroup).get()
+    }
+
+    /// Asynchronously loads multiple keys, promising an array of values:
+    ///
+    /// ```
+    /// let aAndB = try await myLoader.loadMany(keys: [ "a", "b" ], on: eventLoopGroup)
+    /// ```
+    ///
+    /// This is equivalent to the more verbose:
+    ///
+    /// ```
+    /// async let a = myLoader.load(key: "a", on: eventLoopGroup)
+    /// async let b = myLoader.load(key: "b", on: eventLoopGroup)
+    /// let aAndB = try await a + b
+    /// ```
+    @available(macOS 12, iOS 15, watchOS 8, tvOS 15, *)
+    func loadMany(keys: [Key], on eventLoopGroup: EventLoopGroup) async throws -> [Value] {
+        try await loadMany(keys: keys, on: eventLoopGroup).get()
+    }
+}
+
+#endif

Tests/DataLoaderTests/DataLoaderAsyncTests.swift

@@ -0,0 +1,117 @@
+import XCTest
+import NIO
+
+@testable import DataLoader
+
+#if compiler(>=5.5) && canImport(_Concurrency)
+
+@available(macOS 12, iOS 15, watchOS 8, tvOS 15, *)
+actor Concurrent<T> {
+    var wrappedValue: T
+
+    func nonmutating<Returned>(_ action: (T) throws -> Returned) async rethrows -> Returned {
+        try action(wrappedValue)
+    }
+
+    func mutating<Returned>(_ action: (inout T) throws -> Returned) async rethrows -> Returned {
+        try action(&wrappedValue)
+    }
+
+    init(_ value: T) {
+        self.wrappedValue = value
+    }
+}
+
+
+/// Primary API
+@available(macOS 12, iOS 15, watchOS 8, tvOS 15, *)
+final class DataLoaderAsyncTests: XCTestCase {
+
+    /// Builds a really really simple data loader with async await
+    func testReallyReallySimpleDataLoader() async throws {
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer {
+            XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully())
+        }
+
+        let identityLoader = DataLoader<Int, Int>(
+            on: eventLoopGroup.next(),
+            options: DataLoaderOptions(batchingEnabled: false)
+        ) { keys async in
+            let task = Task {
+                keys.map { DataLoaderFutureValue.success($0) }
+            }
+            return await task.value
+        }
+
+        let value = try await identityLoader.load(key: 1, on: eventLoopGroup)
+
+        XCTAssertEqual(value, 1)
+    }
+
+    /// Supports loading multiple keys in one call
+    func testLoadingMultipleKeys() async throws {
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer {
+            XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully())
+        }
+
+        let identityLoader = DataLoader<Int, Int>(on: eventLoopGroup.next()) { keys in
+            let task = Task {
+                keys.map { DataLoaderFutureValue.success($0) }
+            }
+            return await task.value
+        }
+
+        let values = try await identityLoader.loadMany(keys: [1, 2], on: eventLoopGroup)
+
+        XCTAssertEqual(values, [1,2])
+
+        let empty = try await identityLoader.loadMany(keys: [], on: eventLoopGroup)
+
+        XCTAssertTrue(empty.isEmpty)
+    }
+
+    // Batches multiple requests
+    func testMultipleRequests() async throws {
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer {
+            XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully())
+        }
+
+        let loadCalls = Concurrent<[[Int]]>([])
+
+        let identityLoader = DataLoader<Int, Int>(
+            on: eventLoopGroup.next(),
+            options: DataLoaderOptions(
+                batchingEnabled: true,
+                executionPeriod: nil
+            )
+        ) { keys in
+            await loadCalls.mutating { $0.append(keys) }
+            let task = Task {
+                keys.map { DataLoaderFutureValue.success($0) }
+            }
+            return await task.value
+        }
+
+        async let value1 = identityLoader.load(key: 1, on: eventLoopGroup)
+        async let value2 = identityLoader.load(key: 2, on: eventLoopGroup)
+
+        /// Have to wait for a split second because Tasks may not be executed before this statement
+        try await Task.sleep(nanoseconds: 500_000_000)
+
+        XCTAssertNoThrow(try identityLoader.execute())
+
+        let result1 = try await value1
+        XCTAssertEqual(result1, 1)
+        let result2 = try await value2
+        XCTAssertEqual(result2, 2)
+
+        let calls = await loadCalls.wrappedValue
+        XCTAssertEqual(calls.count, 1)
+        XCTAssertEqual(calls.map { $0.sorted() }, [[1, 2]])
+    }
+}
+
+#endif