Zenn
🍁

Swift Concurrencyでタイムアウトする

2024/05/23に公開
Swift
concurrency
tech

時間のかかる処理の例

func calc() async throws -> Int {
    var sum: Int = 0
    for i in 0 ..< 10 {
        sum += i
        try? await Task.sleep(for: .seconds(0.1))
        try Task.checkCancellation()
    }
    return sum
}

パターン１

素朴にやる方法。Taskを２つ立てタイムレースさせてタイムアウトのTaskの方が速かったら時間のかかる処理のTaskをキャンセルする。

@main
struct TestCode {
    static func main() async {
        let task = Task {
            try await calc()
        }

        let timeout = Duration.seconds(0.3)
        let timeoutTask = Task {
            try await Task.sleep(for: timeout)
            task.cancel()
        }

        do {
            let value = try await task.value
            timeoutTask.cancel()
            print(value)
        } catch {
            print(error.localizedDescription)
        }
    }
}

パターン２

TaskGroupを使ってタイムレースさせる方法。

@main
struct TestCode {
    static func main() async {
        do {
            let value = try await withThrowingTaskGroup(of: Int.self) { group in
                group.addTask {
                    try await calc()
                }
                group.addTask {
                    let timeout = Duration.seconds(0.3)
                    try await Task.sleep(for: timeout)
                    throw CancellationError()
                }
                defer { group.cancelAll() }
                guard let result = try await group.next() else {
                    throw CancellationError()
                }
                return result
            }
            print(value)
        } catch {
            print(error.localizedDescription)
        }
    }
}

defer { group.cancelAll() }することが大事。

もしも時間のかかる処理がTaskのキャンセルに対応していない場合、当然ですが処理はキャンセルされずに処理にかかる時間だけawaitでvalueの取得を待つことになります。そのため例えば、処理が無限に終わらないような関数に対してこれらのタイムアウトパターンを適用しても、処理がリークして無限に待機することになってしまいます。よくある具体例としてはwithCheckedContinuation()やwithCheckedThrowingContinuationはスコープ内の処理に対してキャンセル対応されていないため、continuation.resume()をし忘れると無限に待機します。

Discussion