codex: fix cancellation and sendability

Author: Thibault Wittemberg

Sources/AsyncChannels/AsyncBufferedChannel.swift

@@ -29,7 +29,7 @@ import OrderedCollections
 /// sut.send(3)
 /// sut.finish()
 /// ```
-public final class AsyncBufferedChannel<Element>: AsyncSequence, Sendable {
+public final class AsyncBufferedChannel<Element: Sendable>: AsyncSequence, Sendable {
   public typealias Element = Element
   public typealias AsyncIterator = Iterator
 

Sources/AsyncSubjects/AsyncPassthroughSubject.swift

@@ -30,7 +30,7 @@
 /// passthrough.send(2)
 /// passthrough.send(.finished)
 /// ```
-public final class AsyncPassthroughSubject<Element>: AsyncSubject {
+public final class AsyncPassthroughSubject<Element: Sendable>: AsyncSubject {
   public typealias Element = Element
   public typealias Failure = Never
   public typealias AsyncIterator = Iterator

Sources/Combiners/Merge/MergeStateMachine.swift

@@ -240,7 +240,7 @@ struct MergeStateMachine<Element>: Sendable {
         }
       }
 
-      if case .termination = regulatedElement, case .element(.failure) = regulatedElement {
+      if case .element(.failure) = regulatedElement {
         self.task.cancel()
       }
 

Sources/Operators/AsyncMulticastSequence.swift

@@ -105,37 +105,43 @@ where Base.Element == Subject.Element, Subject.Failure == Error, Base.AsyncItera
   }
 
   func next() async {
-    await Task {
-      let (canAccessBase, iterator) = self.state.withCriticalRegion { state -> (Bool, Base.AsyncIterator?) in
-        switch state {
-          case .available(let iterator):
-            state = .busy
-            return (true, iterator)
-          case .busy:
-            return (false, nil)
-        }
-      }
-
-      guard canAccessBase, var iterator = iterator else { return }
-
-      let toSend: Result<Element?, Error>
-      do {
-        let element = try await iterator.next()
-        toSend = .success(element)
-      } catch {
-        toSend = .failure(error)
+    guard !Task.isCancelled else { return }
+
+    let (canAccessBase, iterator) = self.state.withCriticalRegion { state -> (Bool, Base.AsyncIterator?) in
+      switch state {
+        case .available(let iterator):
+          state = .busy
+          return (true, iterator)
+        case .busy:
+          return (false, nil)
       }
+    }
 
+    guard canAccessBase, var iterator = iterator else { return }
+    defer {
       self.state.withCriticalRegion { state in
         state = .available(iterator)
       }
+    }
 
-      switch toSend {
-        case .success(.some(let element)): self.subject.send(element)
-        case .success(.none): self.subject.send(.finished)
-        case .failure(let error): self.subject.send(.failure(error))
-      }
-    }.value
+    guard !Task.isCancelled else { return }
+
+    let toSend: Result<Element?, Error>
+    do {
+      let element = try await iterator.next()
+      toSend = .success(element)
+    } catch {
+      guard !Task.isCancelled else { return }
+      toSend = .failure(error)
+    }
+
+    guard !Task.isCancelled else { return }
+
+    switch toSend {
+      case .success(.some(let element)): self.subject.send(element)
+      case .success(.none): self.subject.send(.finished)
+      case .failure(let error): self.subject.send(.failure(error))
+    }
   }
 
   public func makeAsyncIterator() -> AsyncIterator {

Tests/Combiners/Merge/AsyncMergeSequenceTests.swift

@@ -41,6 +41,34 @@ private struct TimedAsyncSequence<Element>: AsyncSequence, AsyncIteratorProtocol
   }
 }
 
+private struct CancellationAwareSequence<Element>: AsyncSequence, AsyncIteratorProtocol {
+  typealias Element = Element
+  typealias AsyncIterator = CancellationAwareSequence
+
+  let onStart: @Sendable () -> Void
+  let onCancel: @Sendable () -> Void
+  var hasStarted = false
+
+  mutating func next() async throws -> Element? {
+    if !hasStarted {
+      hasStarted = true
+      onStart()
+    }
+
+    do {
+      try await Task.sleep(nanoseconds: 5_000_000_000)
+      return nil
+    } catch {
+      onCancel()
+      return nil
+    }
+  }
+
+  func makeAsyncIterator() -> AsyncIterator {
+    self
+  }
+}
+
 final class AsyncMergeSequenceTests: XCTestCase {
   func testMerge_merges_sequences_according_to_the_timeline_using_asyncSequences() async throws {
     // -- 0 ------------------------------- 1000 ----------------------------- 2000 -
@@ -306,4 +334,34 @@ final class AsyncMergeSequenceTests: XCTestCase {
 
     task.cancel()
   }
+
+  func testMerge_cancels_other_bases_on_error() async {
+    let baseStartedExpectation = expectation(description: "The blocking base has started")
+    let baseCancelledExpectation = expectation(description: "The blocking base has been cancelled")
+
+    let blockingBase = CancellationAwareSequence<Int>(
+      onStart: { baseStartedExpectation.fulfill() },
+      onCancel: { baseCancelledExpectation.fulfill() }
+    )
+    let failingBase = TimedAsyncSequence(intervalInMills: [0, 0], sequence: [1, 2], indexOfError: 1)
+
+    let sut = merge(failingBase, blockingBase)
+    var iterator = sut.makeAsyncIterator()
+
+    do {
+      _ = try await iterator.next()
+    } catch {
+      XCTFail("The first element should not fail")
+    }
+    await fulfillment(of: [baseStartedExpectation], timeout: 1)
+
+    do {
+      _ = try await iterator.next()
+      XCTFail("The iteration should fail")
+    } catch {
+      XCTAssertEqual(error as? MockError, MockError(code: 1))
+    }
+
+    await fulfillment(of: [baseCancelledExpectation], timeout: 1)
+  }
 }

Tests/Operators/AsyncMulticastSequenceTests.swift

@@ -39,6 +39,39 @@ private class SpyAsyncSequenceForNumberOfIterators<Element>: AsyncSequence {
   }
 }
 
+private struct CancellationAwareSequence: AsyncSequence {
+  typealias Element = Int
+  typealias AsyncIterator = Iterator
+
+  let onStart: @Sendable () -> Void
+  let onCancel: @Sendable () -> Void
+
+  func makeAsyncIterator() -> AsyncIterator {
+    Iterator(onStart: self.onStart, onCancel: self.onCancel)
+  }
+
+  struct Iterator: AsyncIteratorProtocol {
+    let onStart: @Sendable () -> Void
+    let onCancel: @Sendable () -> Void
+    var hasStarted = false
+
+    mutating func next() async throws -> Int? {
+      if !hasStarted {
+        hasStarted = true
+        onStart()
+      }
+
+      do {
+        try await Task.sleep(nanoseconds: 5_000_000_000)
+        return nil
+      } catch {
+        onCancel()
+        return nil
+      }
+    }
+  }
+}
+
 final class AsyncMulticastSequenceTests: XCTestCase {
   func test_multiple_loops_receive_elements_from_single_baseIterator() {
     let taskHaveIterators = expectation(description: "All tasks have their iterator")
@@ -156,4 +189,26 @@ final class AsyncMulticastSequenceTests: XCTestCase {
       XCTAssertEqual(error as? MockError, expectedError)
     }
   }
+
+  func test_multicast_cancels_upstream_when_consumer_cancels() async {
+    let upstreamStartedExpectation = expectation(description: "Upstream started")
+    let upstreamCancelledExpectation = expectation(description: "Upstream cancelled")
+
+    let base = CancellationAwareSequence(
+      onStart: { upstreamStartedExpectation.fulfill() },
+      onCancel: { upstreamCancelledExpectation.fulfill() }
+    )
+    let stream = AsyncThrowingPassthroughSubject<Int, Error>()
+    let sut = base.multicast(stream).autoconnect()
+
+    let task = Task {
+      var iterator = sut.makeAsyncIterator()
+      _ = try? await iterator.next()
+    }
+
+    await fulfillment(of: [upstreamStartedExpectation], timeout: 1)
+    task.cancel()
+
+    await fulfillment(of: [upstreamCancelledExpectation], timeout: 1)
+  }
 }