timeout-cont-demo.swift

import Foundation

struct TimeoutError: Error {}

extension Result where Failure == any Error {
    init(catching: () async throws -> Success) async {
        do {
            self = .success(try await catching())
        } catch {
            self = .failure(error)
        }
    }
}

func myAsyncProc() async {
    // continuationで実装する自由な処理
    await withCheckedContinuation { (c) in
        // 中身は5秒かかる非同期処理を模した無意味な例
        Task {
            try await Task.sleep(for: .seconds(5))
            c.resume()
        }
    }
}

final class SingleSink<T>: @unchecked Sendable {
    init(con: CheckedContinuation<T, any Error>) {
        self.con = con
    }

    let lock = NSLock()
    let con: CheckedContinuation<T, any Error>
    var isCompleted = false

    func send(result: Result<T, any Error>) {
        lock.withLock {
            if isCompleted { return }
            isCompleted = true
            con.resume(with: result)
        }
    }
}

func runWithTimeout2<T>(
    timeout: Duration,
    _ operation: @escaping @Sendable () async throws -> T
) async throws -> T {
    return try await withCheckedThrowingContinuation { (con) in
        let sink = SingleSink(con: con)

        Task {
            let result = await Result {
                try await operation()
            }
            sink.send(result: result)
        }

        Task {
            try await Task.sleep(for: timeout)
            sink.send(result: .failure(TimeoutError()))
        }
    }
}


func main() async throws {
    // 自由な処理に、タイムアウトを設定して呼び出す
    do {
        try await runWithTimeout2(timeout: .seconds(1)) {
            await myAsyncProc()
        }
        print("proc completed")
    } catch {
        if error is TimeoutError {
            // 1秒後にこれが呼ばれて嬉しい
            print("timeout")
        } else {
            fatalError("error: \(error)")
        }
    }
}

try await main()