Timeout Handler

Sources/Timeout.swift

@@ -0,0 +1,170 @@
+//
+//  Timeout.swift
+//  swift-timeout
+//
+//  Created by Simon Whitty on 02/06/2025.
+//  Copyright 2025 Simon Whitty
+//
+//  Distributed under the permissive MIT license
+//  Get the latest version from here:
+//
+//  https://github.com/swhitty/swift-timeout
+//
+//  Permission is hereby granted, free of charge, to any person obtaining a copy
+//  of this software and associated documentation files (the "Software"), to deal
+//  in the Software without restriction, including without limitation the rights
+//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+//  copies of the Software, and to permit persons to whom the Software is
+//  furnished to do so, subject to the following conditions:
+//
+//  The above copyright notice and this permission notice shall be included in all
+//  copies or substantial portions of the Software.
+//
+//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+//  SOFTWARE.
+//
+
+#if compiler(>=6.0)
+import Foundation
+
+public struct Timeout: Sendable {
+    fileprivate var canary: @Sendable () -> Void
+    fileprivate let shared: SharedState
+
+    @discardableResult
+    public func expire(seconds: TimeInterval) -> Bool {
+        enqueue {
+            try await Task.sleep(nanoseconds: UInt64(seconds * 1_000_000_000))
+            throw TimeoutError("Task timed out before completion. Timeout: \(seconds) seconds.")
+        }
+    }
+
+    @discardableResult
+    public func expireImmediatley() -> Bool {
+        enqueue(flagAsComplete: true) {
+            throw TimeoutError("Task timed out before completion. expireImmediatley()")
+        }
+    }
+
+    @discardableResult
+    public func cancelExpiration() -> Bool {
+        enqueue {
+            try await Task.sleepIndefinitely()
+        }
+    }
+
+    struct State {
+        var running: Task<Never, any Error>?
+        var pending: (@Sendable () async throws -> Never)?
+        var isComplete: Bool = false
+    }
+
+    final class SharedState: Sendable {
+        let state: Mutex<State>
+
+        init(pending: @escaping @Sendable () async throws -> Never) {
+            state = Mutex(.init(pending: pending))
+        }
+    }
+}
+
+@available(macOS 13.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *)
+public extension Timeout {
+
+    @discardableResult
+    func expire<C: Clock>(
+        after instant: C.Instant,
+        tolerance: C.Instant.Duration? = nil,
+        clock: C
+    ) -> Bool {
+        enqueue {
+            try await Task.sleep(until: instant, tolerance: tolerance, clock: clock)
+            throw TimeoutError("Task timed out before completion. Deadline: \(instant).")
+        }
+    }
+
+    @discardableResult
+    func expire(
+        after instant: ContinuousClock.Instant,
+        tolerance: ContinuousClock.Instant.Duration? = nil
+    ) -> Bool {
+        expire(after: instant, tolerance: tolerance, clock: ContinuousClock())
+    }
+}
+
+extension Timeout {
+
+    init(
+        canary: @escaping @Sendable () -> Void,
+        pending closure: @escaping @Sendable () async throws -> Never
+    ) {
+        self.canary = canary
+        self.shared = .init(pending: closure)
+    }
+
+    @discardableResult
+    func enqueue(flagAsComplete: Bool = false, closure: @escaping @Sendable () async throws -> Never) -> Bool {
+        shared.state.withLock { s in
+            guard !s.isComplete else { return false }
+            s.pending = closure
+            s.running?.cancel()
+            s.isComplete = flagAsComplete
+            return true
+        }
+    }
+
+    func startPendingTask() -> Task<Never, any Error>? {
+        return shared.state.withLock { s in
+            guard let pending = s.pending else {
+                s.isComplete = true
+                return nil
+            }
+            let task = Task { try await pending() }
+            s.pending = nil
+            s.running = task
+            return task
+        }
+    }
+
+    func waitForTimeout() async throws {
+        var lastError: (any Error)?
+        while let task = startPendingTask() {
+            do {
+                try await withTaskCancellationHandler {
+                    try await task.value
+                } onCancel: {
+                    task.cancel()
+                }
+            } catch is CancellationError {
+                lastError = nil
+            } catch {
+                lastError = error
+            }
+        }
+
+        if let lastError {
+            throw lastError
+        }
+    }
+}
+
+func withNonEscapingTimeout<T>(
+    _ timeout: @escaping @Sendable () async throws -> Never,
+    isolation: isolated (any Actor)? = #isolation,
+    body: (Timeout) async throws -> sending T
+) async throws -> sending T {
+    // canary ensuring Timeout does not escape at runtime.
+    // Swift 6.2 and later enforce at compile time with ~Escapable
+    try await withoutActuallyEscaping({ @Sendable in }) { escaping in
+        _ = isolation
+        let timeout = Timeout(canary: escaping, pending: timeout)
+        return try await Transferring(body(timeout))
+    }.value
+}
+
+#endif

Sources/withThrowingTimeout.swift

@@ -44,6 +44,17 @@ public func withThrowingTimeout<T>(
     isolation: isolated (any Actor)? = #isolation,
     seconds: TimeInterval,
     body: () async throws -> sending T
+) async throws -> sending T {
+    try await _withThrowingTimeout(isolation: isolation, body: { _ in try await body() }) {
+        try await Task.sleep(nanoseconds: UInt64(seconds * 1_000_000_000))
+        throw TimeoutError("Task timed out before completion. Timeout: \(seconds) seconds.")
+    }.value
+}
+
+public func withThrowingTimeout<T>(
+    isolation: isolated (any Actor)? = #isolation,
+    seconds: TimeInterval,
+    body: (Timeout) async throws -> sending T
 ) async throws -> sending T {
     try await _withThrowingTimeout(isolation: isolation, body: body) {
         try await Task.sleep(nanoseconds: UInt64(seconds * 1_000_000_000))
@@ -59,7 +70,7 @@ public func withThrowingTimeout<T, C: Clock>(
     clock: C,
     body: () async throws -> sending T
 ) async throws -> sending T {
-    try await _withThrowingTimeout(isolation: isolation, body: body) {
+    try await _withThrowingTimeout(isolation: isolation, body: { _ in try await body() }) {
         try await Task.sleep(until: instant, tolerance: tolerance, clock: clock)
         throw TimeoutError("Task timed out before completion. Deadline: \(instant).")
     }.value
@@ -72,39 +83,41 @@ public func withThrowingTimeout<T>(
     tolerance: ContinuousClock.Instant.Duration? = nil,
     body: () async throws -> sending T
 ) async throws -> sending T {
-    try await _withThrowingTimeout(isolation: isolation, body: body) {
+    try await _withThrowingTimeout(isolation: isolation, body: { _ in try await body() }) {
         try await Task.sleep(until: instant, tolerance: tolerance, clock: ContinuousClock())
         throw TimeoutError("Task timed out before completion. Deadline: \(instant).")
     }.value
 }
 
 private func _withThrowingTimeout<T>(
     isolation: isolated (any Actor)? = #isolation,
-    body: () async throws -> sending T,
-    timeout: @Sendable @escaping () async throws -> Never
+    body: (Timeout) async throws -> sending T,
+    timeout closure: @Sendable @escaping () async throws -> Never
 ) async throws -> Transferring<T> {
     try await withoutActuallyEscaping(body) { escapingBody in
-        let bodyTask = Task {
-            defer { _ = isolation }
-            return try await Transferring(escapingBody())
-        }
-        let timeoutTask = Task {
-            defer { bodyTask.cancel() }
-            try await timeout()
-        }
+        try await withNonEscapingTimeout(closure) { timeout in
+            let bodyTask = Task {
+                defer { _ = isolation }
+                return try await Transferring(escapingBody(timeout))
+            }
+            let timeoutTask = Task {
+                defer { bodyTask.cancel() }
+                try await timeout.waitForTimeout()
+            }
 
-        let bodyResult = await withTaskCancellationHandler {
-            await bodyTask.result
-        } onCancel: {
-            bodyTask.cancel()
-        }
-        timeoutTask.cancel()
+            let bodyResult = await withTaskCancellationHandler {
+                await bodyTask.result
+            } onCancel: {
+                bodyTask.cancel()
+            }
+            timeoutTask.cancel()
 
-        if case .failure(let timeoutError) = await timeoutTask.result,
-           timeoutError is TimeoutError {
-            throw timeoutError
-        } else {
-            return try bodyResult.get()
+            if case .failure(let timeoutError) = await timeoutTask.result,
+               timeoutError is TimeoutError {
+                throw timeoutError
+            } else {
+                return try bodyResult.get()
+            }
         }
     }
 }

Tests/withThrowingTimeoutTests.swift

@@ -30,7 +30,7 @@
 //
 
 #if canImport(Testing)
-import Timeout
+@testable import Timeout
 import Foundation
 import Testing
 
@@ -137,6 +137,46 @@ struct WithThrowingTimeoutTests {
             }
         }
     }
+
+    @Test
+    func timeout_ExpiresImmediatley() async throws {
+        await #expect(throws: TimeoutError.self) {
+            try await withThrowingTimeout(seconds: 1_000) { timeout in
+                timeout.expireImmediatley()
+            }
+        }
+    }
+
+    @Test
+    func timeout_ExpiresAfterSeconds() async throws {
+        await #expect(throws: TimeoutError.self) {
+            try await withThrowingTimeout(seconds: 1_000) { timeout in
+                timeout.expire(seconds: 0.1)
+                try await Task.sleepIndefinitely()
+            }
+        }
+    }
+
+    @Test
+    func timeout_ExpiresAfterDeadline() async throws {
+        await #expect(throws: TimeoutError.self) {
+            try await withThrowingTimeout(seconds: 1_000) { timeout in
+                timeout.expire(after: .now + .seconds(0.1))
+                try await Task.sleepIndefinitely()
+            }
+        }
+    }
+
+    @Test
+    func timeout_ExpirationCancels() async throws {
+        #expect(
+            try await withThrowingTimeout(seconds: 0.1) { timeout in
+                timeout.cancelExpiration()
+                try await Task.sleep(for: .seconds(0.3))
+                return "Fish"
+            } == "Fish"
+        )
+    }
 }
 
 public struct NonSendable<T> {