TaskGroup Syntax Sugar

TaskGroupExtensionTest.swift

import XCTest

final class TaskGroupExtensionTest: XCTestCase {
    func testWhenAll() async throws {
        // 모든 Task 가 완료될 때 까지 기다리고 인자 순서대에 맞춰서 결괏값 목록 반환
        // Act
        let results = await TaskGroup.whenAll(
            { await self.fetchTestDataAsync(numberOfInvocation: 0, delay: .microseconds(100)) },
            { await self.fetchTestDataAsync(numberOfInvocation: 1, delay: .microseconds(100)) },
            { await self.fetchTestDataAsync(numberOfInvocation: 2, delay: .microseconds(100)) },
            { await self.fetchTestDataAsync(numberOfInvocation: 3, delay: .microseconds(100)) },
            { await self.fetchTestDataAsync(numberOfInvocation: 4, delay: .microseconds(100)) },
            { await self.fetchTestDataAsync(numberOfInvocation: 5, delay: .microseconds(100)) },
            { await self.fetchTestDataAsync(numberOfInvocation: 6, delay: .microseconds(100)) },
            { await self.fetchTestDataAsync(numberOfInvocation: 7, delay: .microseconds(100)) },
            { await self.fetchTestDataAsync(numberOfInvocation: 8, delay: .microseconds(100)) },
            { await self.fetchTestDataAsync(numberOfInvocation: 9, delay: .microseconds(100)) }
        )
        
        // Assert
        for (i, result) in results.enumerated() {
            assert(try! result.get() == "비동기 데이터 \(i)")
        }
    }

    func testWhenAny() async throws {
        // 가장 빨리 완료된 Task 의 결괏값을 반환
        // Act
        let result = await TaskGroup.whenAny(
            { await self.fetchTestDataAsync(numberOfInvocation: 0, delay: .seconds(10)) },
            { await self.fetchTestDataAsync(numberOfInvocation: 1, delay: .seconds(10)) },
            { await self.fetchTestDataAsync(numberOfInvocation: 2, delay: .seconds(10)) },
            { await self.fetchTestDataAsync(numberOfInvocation: 3, delay: .seconds(10)) },
            { await self.fetchTestDataAsync(numberOfInvocation: 4, delay: .microseconds(100)) },
            { await self.fetchTestDataAsync(numberOfInvocation: 5, delay: .seconds(10)) },
            { await self.fetchTestDataAsync(numberOfInvocation: 6, delay: .seconds(10)) },
            { await self.fetchTestDataAsync(numberOfInvocation: 7, delay: .seconds(10)) },
            { await self.fetchTestDataAsync(numberOfInvocation: 8, delay: .seconds(10)) },
            { await self.fetchTestDataAsync(numberOfInvocation: 9, delay: .seconds(10)) }
        )
        
        // Assert
        assert(try! result.get() == "비동기 데이터 4")
    }
    
    private func fetchTestDataAsync(
        numberOfInvocation: Int,
        delay: Duration
    ) async -> String {
        try? await Task.sleep(for: delay)
        return "비동기 데이터 \(numberOfInvocation)"
    }
}

extension TaskGroup {
    static func whenAll(
        _ tasks: () async throws -> ChildTaskResult...
    ) async -> [Result<ChildTaskResult, Error>] {
        return await withTaskGroup(of: Result<ChildTaskResult, Error>.self) { group in
            for task in tasks {
                group.addTask {
                    do {
                        return .success(try await task())
                    } catch {
                        return .failure(error)
                    }
                }
            }
            
            await group.waitForAll()
            
            var results: [Result<ChildTaskResult, Error>] = []
            results.reserveCapacity(tasks.count)
            for await task in group {
                results.append(task)
            }
            return results
        }
    }
    
    static func whenAny(
        _ tasks: () async throws -> ChildTaskResult...
    ) async -> Result<ChildTaskResult, Error> {
        if (tasks.isEmpty) {
            fatalError("tasks should not be empty")
        }
        
        return await withTaskGroup(of: Result<ChildTaskResult, Error>.self) { group in
            for task in tasks {
                group.addTask {
                    do {
                        return .success(try await task())
                    } catch {
                        return .failure(error)
                    }
                }
            }

            let result = await group.first { _ in true }
            group.cancelAll()
            return result!
        }
    }
}