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,11 +752,13 @@ extension Optional where Wrapped == String {
     }
 }
 
+/// Runs `body`, and then runs `onCleanup` if body throws an error, or if the parent task is cancelled. In the latter case, `onCleanup` may be run concurrently with `body`. `body` is guaranteed to run exactly once. `onCleanup` is guaranteed to run only once, or not at all.
 internal func withAsyncTaskCleanupHandler<Result>(
     _ body: () async throws -> Result,
     onCleanup handler: @Sendable @escaping () async -> Void,
     isolation: isolated (any Actor)? = #isolation
 ) async rethrows -> Result {
+    let (runCancellationHandlerStream, runCancellationHandlerContinuation) = AsyncThrowingStream.makeStream(of: Void.self)
     return try await withThrowingTaskGroup(
         of: Void.self,
         returning: Result.self
@@ -767,15 +769,38 @@ internal func withAsyncTaskCleanupHandler<Result>(
             // 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()
+            runCancellationHandlerContinuation.finish(throwing: CancellationError())
+        }
+
+        group.addTask {
+            // Enumerate the async stream until it completes or throws an error.
+            // Since we signal completion of the stream from cancellation or the
+            // parent task or the body throwing, this ensures that we run the
+            // cleanup handler exactly once in any failure scenario, and also do
+            // so _immediately_ if the failure scenario is due to parent task
+            // cancellation. We do so in a detached Task to prevent cancellation
+            // of the parent task from interrupting enumeration of the stream itself.
+            await Task.detached {
+                do {
+                    var iterator = runCancellationHandlerStream.makeAsyncIterator()
+                    while let _ = try await iterator.next() {
+                    }
+                } catch {
+                    await handler()
+                }
+            }.value
+        }
+
+        defer {
+            group.cancelAll()
         }
 
         do {
             let result = try await body()
-            group.cancelAll()
+            runCancellationHandlerContinuation.finish()
             return result
         } catch {
-            await handler()
+            runCancellationHandlerContinuation.finish(throwing: error)
             throw error
         }
     }