diff --git a/Sources/AsyncTimeoutSequence.swift b/Sources/AsyncTimeoutSequence.swift index 748425c..5975838 100644 --- a/Sources/AsyncTimeoutSequence.swift +++ b/Sources/AsyncTimeoutSequence.swift @@ -40,10 +40,19 @@ public extension AsyncSequence where Element: Sendable { } /// Creates an asynchronous sequence that throws error if any iteration - /// takes longer than provided `Duration`. + /// takes longer than provided `Duration` using the supplied `Clock`. @available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) - func timeout(duration: Duration) -> AsyncTimeoutSequence { - AsyncTimeoutSequence(base: self, duration: duration) + func timeout( + duration: C.Duration, + tolerance: C.Instant.Duration? = nil, + clock: C = ContinuousClock() + ) -> AsyncTimeoutSequence { + AsyncTimeoutSequence( + base: self, + duration: duration, + tolerance: tolerance, + clock: clock + ) } } @@ -51,7 +60,7 @@ public struct AsyncTimeoutSequence: AsyncSequence where Bas public typealias Element = Base.Element private let base: Base - private let interval: TimeoutInterval + private let interval: TimeoutInterval public init(base: Base, seconds: TimeInterval) { self.base = base @@ -59,9 +68,14 @@ public struct AsyncTimeoutSequence: AsyncSequence where Bas } @available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) - public init(base: Base, duration: Duration) { + public init( + base: Base, + duration: C.Duration, + tolerance: C.Instant.Duration? = nil, + clock: C = ContinuousClock() + ) { self.base = base - self.interval = .duration(.init(duration)) + self.interval = .duration(.init(duration: duration, tolerance: tolerance, clock: clock)) } public func makeAsyncIterator() -> AsyncIterator { @@ -73,9 +87,9 @@ public struct AsyncTimeoutSequence: AsyncSequence where Bas public struct AsyncIterator: AsyncIteratorProtocol { private var iterator: Base.AsyncIterator - private let interval: TimeoutInterval + private let interval: TimeoutInterval - init(iterator: Base.AsyncIterator, interval: TimeoutInterval) { + fileprivate init(iterator: Base.AsyncIterator, interval: TimeoutInterval) { self.iterator = iterator self.interval = interval } @@ -91,7 +105,7 @@ public struct AsyncTimeoutSequence: AsyncSequence where Bas guard #available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) else { fatalError("cannot occur") } - return try await withThrowingTimeout(after: .now + durationBox.value) { + return try await durationBox.withThrowingTimeout { try await self.iterator.next() } } @@ -99,21 +113,39 @@ public struct AsyncTimeoutSequence: AsyncSequence where Bas } } -enum TimeoutInterval { +private enum TimeoutInterval { case timeInterval(TimeInterval) case duration(DurationBox) struct DurationBox { - private let storage: Any + private typealias TimeoutClosure = (() async throws -> sending T) async throws -> sending T + + private let storage: TimeoutClosure @available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) - var value: Duration { - storage as! Duration + init( + duration: C.Duration, + tolerance: C.Instant.Duration? = nil, + clock: C + ) { + self.storage = { closure in + try await Timeout.withThrowingTimeout( + after: clock.now.advanced(by: duration), + tolerance: tolerance, + clock: clock + ) { + try await closure() + } + } } @available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *) - init(_ duration: Duration) { - self.storage = duration + func withThrowingTimeout( + _ closure: () async throws -> sending T + ) async throws -> T { + try await storage { + try await closure() + } } } } diff --git a/Tests/AsyncTimeoutSequenceTests.swift b/Tests/AsyncTimeoutSequenceTests.swift index 8eedab1..3853fbc 100644 --- a/Tests/AsyncTimeoutSequenceTests.swift +++ b/Tests/AsyncTimeoutSequenceTests.swift @@ -71,4 +71,29 @@ struct AsyncTimeoutSequenceTests { try await iterator.next() } } + + @Test + func timeoutDurationWithSuspendingClock() async throws { + let (stream, continuation) = AsyncStream.makeStream() + let t = Task { + continuation.yield(1) + try await Task.sleep(nanoseconds: 1_000) + continuation.yield(2) + try await Task.sleepIndefinitely() + } + defer { t.cancel() } + var iterator = stream + .timeout( + duration: .milliseconds(100), + tolerance: .zero, + clock: SuspendingClock() + ) + .makeAsyncIterator() + + #expect(try await iterator.next() == 1) + #expect(try await iterator.next() == 2) + await #expect(throws: TimeoutError.self) { + try await iterator.next() + } + } }