子タスクがキャンセルハンドルしなくても元スコープを抜け出せるtimeout

test.swift

import Foundation

func brokenTask() async {
  await withCheckedContinuation { (c) in
    // キャンセルシグナルが握りつぶされていて正しく中断しない処理
    Task {
      try await Task.sleep(for: .seconds(5))
      c.resume()
    }
  }
}

struct TimeoutError: Error {}

class OnceToken {
  private var once: Bool = true
  private let lock = NSLock()

  func run(_ op: () -> Void) {
    let ok = lock.withLock {
      defer { once = false }
      return once
    }
    if ok {
      op()
    }
  }
}

struct OnceContinuation<T, E: Error> {
  var c: CheckedContinuation<T, E>
  let once = OnceToken()

  public func resume(returning value: T) {
    once.run {
      c.resume(returning: value)
    }
  }

  public func resume(throwing error: E) {
    once.run {
      c.resume(throwing: error)
    }
  }
}

func runWithTimeout<T>(
  timeout: Duration,
  _ operation: @escaping @Sendable () async throws -> T
) async throws -> T {
  try await withCheckedThrowingContinuation { c in
    let c = OnceContinuation(c: c)
    let task = Task {
      do {
        c.resume(returning: try await operation())
      } catch {
        c.resume(throwing: error)
      }
    }
    Task {
      try await Task.sleep(for: timeout)
      task.cancel()
      c.resume(throwing: TimeoutError())
    }
  }
}

try await runWithTimeout(timeout: .milliseconds(500)) {
  print("Task begin")

  await brokenTask()
  // try await Task.sleep(for: .seconds(3))
  print("Task end")
}