withAsyncTaskCleanupHandler calls its cleanup handler twice if body throws

If the body closure given to withAsyncTaskCleanupHandler throws, its cleanup handler is called twice: once in the catch block where body is called, and then again in the task group task once the Task.sleep throws due to cancellation, which swallows the error and then continues to call the handler as well. That results in the teardown sequence being invoked twice, as well as the teardown sequence being invoked for non-error cases.

This patch ensures the cleanup handler is invoked in failure cases only, and only once.

Closes #80

Author: jakepetroules

Sources/Subprocess/Configuration.swift

@@ -752,31 +752,156 @@ extension Optional where Wrapped == String {
     }
 }
 
-internal func withAsyncTaskCleanupHandler<Result>(
-    _ body: () async throws -> Result,
+internal func withAsyncTaskCleanupHandler<Success>(
+    _ body: () async throws -> Success,
     onCleanup handler: @Sendable @escaping () async -> Void,
     isolation: isolated (any Actor)? = #isolation
-) async rethrows -> Result {
+) async rethrows -> Success {
+    let runCancellationHandlerPromise = Promise<Bool, Never>()
     return try await withThrowingTaskGroup(
         of: Void.self,
-        returning: Result.self
+        returning: Success.self
     ) { group in
         group.addTask {
             // Keep this task sleep indefinitely until the parent task is cancelled.
             // `Task.sleep` throws `CancellationError` when the task is canceled
             // before the time ends. We then run the cancel handler.
             do { while true { try await Task.sleep(nanoseconds: 1_000_000_000) } } catch {}
             // Run task cancel handler
-            await handler()
+            runCancellationHandlerPromise.fulfill(with: true)
+        }
+
+        group.addTask {
+            if await runCancellationHandlerPromise.value {
+                await handler()
+            }
+        }
+
+        defer {
+            group.cancelAll()
         }
 
         do {
             let result = try await body()
-            group.cancelAll()
+            runCancellationHandlerPromise.fulfill(with: false)
             return result
         } catch {
-            await handler()
+            runCancellationHandlerPromise.fulfill(with: true)
             throw error
         }
     }
 }
+
+#if canImport(Darwin)
+import os
+#else
+import Synchronization
+#endif
+
+/// Represents a placeholder for a value which will be provided by synchronous code running in another execution context.
+internal final class Promise<Success: Sendable, Failure: Swift.Error>: Sendable {
+    private struct State {
+        var waiters: [CheckedContinuation<Success, Failure>] = []
+        var value: Result<Success, Failure>?
+    }
+
+    #if canImport(Darwin)
+    private let state: OSAllocatedUnfairLock<State>
+    #else
+    private let state: Mutex<State>
+    #endif
+
+    /// Creates a new promise in the unfulfilled state.
+    public init() {
+        #if canImport(Darwin)
+        state = OSAllocatedUnfairLock(initialState: State(value: nil))
+        #else
+        state = Mutex(State(value: nil))
+        #endif
+    }
+
+    deinit {
+        state.withLock { state in
+            precondition(state.waiters.isEmpty, "Deallocated with remaining waiters")
+        }
+    }
+
+    /// Fulfills the promise with the specified result.
+    ///
+    /// - returns: Whether the promise was already fulfilled.
+    @discardableResult
+    public func fulfill(with result: Result<Success, Failure>) -> Bool {
+        let (waiters, alreadyFulfilled): ([CheckedContinuation<Success, Failure>], Bool) = state.withLock { state in
+            if state.value != nil {
+                return ([], true)
+            }
+            state.value = result
+            let waiters = state.waiters
+            state.waiters.removeAll()
+            return (waiters, false)
+        }
+
+        // Resume the continuations outside the lock to avoid potential deadlock if invoked in a cancellation handler.
+        for waiter in waiters {
+            waiter.resume(with: result)
+        }
+
+        return alreadyFulfilled
+    }
+
+    /// Fulfills the promise with the specified value.
+    ///
+    /// - returns: Whether the promise was already fulfilled.
+    @discardableResult
+    public func fulfill(with value: Success) -> Bool {
+        fulfill(with: .success(value))
+    }
+
+    /// Fulfills the promise with the specified error.
+    ///
+    /// - returns: Whether the promise was already fulfilled.
+    @discardableResult
+    public func fail(throwing error: Failure) -> Bool {
+        fulfill(with: .failure(error))
+    }
+}
+
+extension Promise where Success == Void {
+    /// Fulfills the promise.
+    ///
+    /// - returns: Whether the promise was already fulfilled.
+    @discardableResult
+    internal func fulfill() -> Bool {
+        fulfill(with: ())
+    }
+}
+
+extension Promise where Failure == Never {
+    /// Suspends if the promise is not yet fulfilled, and returns the value once it is.
+    internal var value: Success {
+        get async {
+            await withCheckedContinuation { continuation in
+                let value: Result<Success, Never>? = state.withLock { state in
+                    if let value = state.value {
+                        return value
+                    } else {
+                        state.waiters.append(continuation)
+                        return nil
+                    }
+                }
+
+                // Resume the continuations outside the lock to avoid potential deadlock if invoked in a cancellation handler.
+                if let value {
+                    continuation.resume(with: value)
+                }
+            }
+        }
+    }
+
+    /// Suspends if the promise is not yet fulfilled, and returns the result once it is.
+    internal var result: Result<Success, Failure> {
+        get async {
+            await .success(value)
+        }
+    }
+}