From 259c8415910c2bd0327dbf8239237eabc3ff7147 Mon Sep 17 00:00:00 2001
From: Danny Canter <danny_canter@apple.com>
Date: Thu, 7 Aug 2025 11:34:45 -0700
Subject: [PATCH] ContainerizationExtras: Timeout adjustments (#254)

Three things:

1. Timeout.run today didn't allow a throwing closure.
2. Add a new constructor for Timeout that takes in a Duration directly
so we can be more granular than seconds.
3. Add unit tests because we all love those :)
---
 Sources/ContainerizationExtras/Timeout.swift  | 27 ++++++++-
 .../TestTimeout.swift                         | 57 +++++++++++++++++++
 2 files changed, 82 insertions(+), 2 deletions(-)
 create mode 100644 Tests/ContainerizationExtrasTests/TestTimeout.swift

diff --git a/Sources/ContainerizationExtras/Timeout.swift b/Sources/ContainerizationExtras/Timeout.swift
index 42b6e55a..557922f9 100644
--- a/Sources/ContainerizationExtras/Timeout.swift
+++ b/Sources/ContainerizationExtras/Timeout.swift
@@ -23,11 +23,11 @@ public struct Timeout {
     /// doesn't finish in the provided `seconds` amount.
     public static func run<T: Sendable>(
         seconds: UInt32,
-        operation: @escaping @Sendable () async -> T
+        operation: @escaping @Sendable () async throws -> T
     ) async throws -> T {
         try await withThrowingTaskGroup(of: T.self) { group in
             group.addTask {
-                await operation()
+                try await operation()
             }
 
             group.addTask {
@@ -43,4 +43,27 @@ public struct Timeout {
             return result
         }
     }
+
+    public static func run<T: Sendable>(
+        for duration: Duration,
+        operation: @escaping @Sendable () async throws -> T
+    ) async throws -> T {
+        try await withThrowingTaskGroup(of: T.self) { group in
+            group.addTask {
+                try await operation()
+            }
+
+            group.addTask {
+                try await Task.sleep(for: duration)
+                throw CancellationError()
+            }
+
+            guard let result = try await group.next() else {
+                fatalError()
+            }
+
+            group.cancelAll()
+            return result
+        }
+    }
 }
diff --git a/Tests/ContainerizationExtrasTests/TestTimeout.swift b/Tests/ContainerizationExtrasTests/TestTimeout.swift
new file mode 100644
index 00000000..a3bfe36b
--- /dev/null
+++ b/Tests/ContainerizationExtrasTests/TestTimeout.swift
@@ -0,0 +1,57 @@
+//===----------------------------------------------------------------------===//
+// Copyright © 2025 Apple Inc. and the Containerization project authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//===----------------------------------------------------------------------===//
+
+import Foundation
+import Testing
+
+@testable import ContainerizationExtras
+
+final class TestTimeoutType {
+    @Test
+    func testNoCancellation() async throws {
+        await #expect(throws: Never.self) {
+            try await Timeout.run(
+                for: .seconds(5),
+                operation: {
+                    return
+                })
+        }
+    }
+
+    @Test
+    func testCancellationError() async throws {
+        await #expect(throws: CancellationError.self) {
+            try await Timeout.run(
+                for: .milliseconds(50),
+                operation: {
+                    try await Task.sleep(for: .seconds(2))
+                })
+        }
+    }
+
+    @Test
+    func testClosureError() async throws {
+        // Check that we get the closures error if we don't timeout, but
+        // the closure does throw before.
+        await #expect(throws: POSIXError.self) {
+            try await Timeout.run(
+                for: .seconds(10),
+                operation: {
+                    throw POSIXError(.E2BIG)
+                })
+        }
+    }
+}