Add Clock and Tolerance to AsyncTimeoutSequence

Sources/AsyncTimeoutSequence.swift

@@ -40,28 +40,42 @@ 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<Self> {
-        AsyncTimeoutSequence(base: self, duration: duration)
+    func timeout<C: Clock>(
+        duration: C.Duration,
+        tolerance: C.Instant.Duration? = nil,
+        clock: C = ContinuousClock()
+    ) -> AsyncTimeoutSequence<Self> {
+        AsyncTimeoutSequence(
+            base: self,
+            duration: duration,
+            tolerance: tolerance,
+            clock: clock
+        )
     }
 }
 
 public struct AsyncTimeoutSequence<Base: AsyncSequence>: AsyncSequence where Base.Element: Sendable {
     public typealias Element = Base.Element
 
     private let base: Base
-    private let interval: TimeoutInterval
+    private let interval: TimeoutInterval<Base.Element?>
 
     public init(base: Base, seconds: TimeInterval) {
         self.base = base
         self.interval = .timeInterval(seconds)
     }
 
     @available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *)
-    public init(base: Base, duration: Duration) {
+    public init<C: Clock>(
+        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<Base: AsyncSequence>: AsyncSequence where Bas
 
     public struct AsyncIterator: AsyncIteratorProtocol {
         private var iterator: Base.AsyncIterator
-        private let interval: TimeoutInterval
+        private let interval: TimeoutInterval<Base.Element?>
 
-        init(iterator: Base.AsyncIterator, interval: TimeoutInterval) {
+        fileprivate init(iterator: Base.AsyncIterator, interval: TimeoutInterval<Base.Element?>) {
             self.iterator = iterator
             self.interval = interval
         }
@@ -91,29 +105,47 @@ public struct AsyncTimeoutSequence<Base: AsyncSequence>: 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()
                 }
             }
         }
     }
 }
 
-enum TimeoutInterval {
+private enum TimeoutInterval<T: Sendable> {
     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<C: Clock>(
+            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()
+            }
         }
     }
 }

Tests/AsyncTimeoutSequenceTests.swift

@@ -71,4 +71,29 @@ struct AsyncTimeoutSequenceTests {
             try await iterator.next()
         }
     }
+
+    @Test
+    func timeoutDurationWithSuspendingClock() async throws {
+        let (stream, continuation) = AsyncStream<Int>.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()
+        }
+    }
 }