Resolve a race condition that caused repeated testing of channel cancellation to sometimes hang (apple#80)

* Resolve a race condition that caused repeated testing of channel cancellation to sometimes hang

* Better name for removal for cancellation in inner Async(Throwing)Channel

Author: phausler

Sources/AsyncAlgorithms/AsyncChannel.swift

@@ -41,15 +41,24 @@ public final class AsyncChannel<Element: Sendable>: AsyncSequence, Sendable {
   struct Awaiting: Hashable {
     var generation: Int
     var continuation: UnsafeContinuation<Element?, Never>?
+    let cancelled: Bool
 
     init(generation: Int, continuation: UnsafeContinuation<Element?, Never>) {
       self.generation = generation
       self.continuation = continuation
+      cancelled = false
     }
 
     init(placeholder generation: Int) {
       self.generation = generation
       self.continuation = nil
+      cancelled = false
+    }
+    
+    init(cancelled generation: Int) {
+      self.generation = generation
+      self.continuation = nil
+      cancelled = true
     }
 
     func hash(into hasher: inout Hasher) {
@@ -66,12 +75,15 @@ public final class AsyncChannel<Element: Sendable>: AsyncSequence, Sendable {
     case pending([UnsafeContinuation<UnsafeContinuation<Element?, Never>?, Never>])
     case awaiting(Set<Awaiting>)
 
-    mutating func remove(_ generation: Int) -> UnsafeContinuation<Element?, Never>? {
+    mutating func cancel(_ generation: Int) -> UnsafeContinuation<Element?, Never>? {
       switch self {
       case .awaiting(var awaiting):
         let continuation = awaiting.remove(Awaiting(placeholder: generation))?.continuation
         self = .awaiting(awaiting)
         return continuation
+      case .idle:
+        self = .awaiting([Awaiting(cancelled: generation)])
+        return nil
       default:
         return nil
       }
@@ -97,12 +109,13 @@ public final class AsyncChannel<Element: Sendable>: AsyncSequence, Sendable {
 
   func cancel(_ generation: Int) {
     state.withCriticalRegion { state in
-      state.emission.remove(generation)
+      state.emission.cancel(generation)
     }?.resume(returning: nil)
   }
 
   func next(_ generation: Int) async -> Element? {
     return await withUnsafeContinuation { continuation in
+      var cancelled = false
       state.withCriticalRegion { state -> UnsafeResumption<UnsafeContinuation<Element?, Never>?, Never>? in
         switch state.emission {
         case .idle:
@@ -117,11 +130,17 @@ public final class AsyncChannel<Element: Sendable>: AsyncSequence, Sendable {
           }
           return UnsafeResumption(continuation: send, success: continuation)
         case .awaiting(var nexts):
-          nexts.insert(Awaiting(generation: generation, continuation: continuation))
+          if nexts.update(with: Awaiting(generation: generation, continuation: continuation)) != nil {
+            nexts.remove(Awaiting(placeholder: generation))
+            cancelled = true
+          }
           state.emission = .awaiting(nexts)
           return nil
         }
       }?.resume()
+      if cancelled {
+        continuation.resume(returning: nil)
+      }
     }
   }
 

Sources/AsyncAlgorithms/AsyncThrowingChannel.swift

@@ -46,15 +46,24 @@ public final class AsyncThrowingChannel<Element: Sendable, Failure: Error>: Asyn
   struct Awaiting: Hashable {
     var generation: Int
     var continuation: UnsafeContinuation<Element?, Error>?
+    let cancelled: Bool
 
     init(generation: Int, continuation: UnsafeContinuation<Element?, Error>) {
       self.generation = generation
       self.continuation = continuation
+      cancelled = false
     }
 
     init(placeholder generation: Int) {
       self.generation = generation
       self.continuation = nil
+      cancelled = false
+    }
+    
+    init(cancelled generation: Int) {
+      self.generation = generation
+      self.continuation = nil
+      cancelled = true
     }
 
     func hash(into hasher: inout Hasher) {
@@ -71,12 +80,15 @@ public final class AsyncThrowingChannel<Element: Sendable, Failure: Error>: Asyn
     case pending([UnsafeContinuation<UnsafeContinuation<Element?, Error>?, Never>])
     case awaiting(Set<Awaiting>)
 
-    mutating func remove(_ generation: Int) -> UnsafeContinuation<Element?, Error>? {
+    mutating func cancel(_ generation: Int) -> UnsafeContinuation<Element?, Error>? {
       switch self {
       case .awaiting(var awaiting):
         let continuation = awaiting.remove(Awaiting(placeholder: generation))?.continuation
         self = .awaiting(awaiting)
         return continuation
+      case .idle:
+        self = .awaiting([Awaiting(cancelled: generation)])
+        return nil
       default:
         return nil
       }
@@ -102,12 +114,13 @@ public final class AsyncThrowingChannel<Element: Sendable, Failure: Error>: Asyn
 
   func cancel(_ generation: Int) {
     state.withCriticalRegion { state in
-      state.emission.remove(generation)
+      state.emission.cancel(generation)
     }?.resume(returning: nil)
   }
 
   func next(_ generation: Int) async throws -> Element? {
     return try await withUnsafeThrowingContinuation { continuation in
+      var cancelled = false
       state.withCriticalRegion { state -> UnsafeResumption<UnsafeContinuation<Element?, Error>?, Never>? in
         switch state.emission {
         case .idle:
@@ -122,11 +135,17 @@ public final class AsyncThrowingChannel<Element: Sendable, Failure: Error>: Asyn
           }
           return UnsafeResumption(continuation: send, success: continuation)
         case .awaiting(var nexts):
-          nexts.insert(Awaiting(generation: generation, continuation: continuation))
+          if nexts.update(with: Awaiting(generation: generation, continuation: continuation)) != nil {
+            nexts.remove(Awaiting(placeholder: generation))
+            cancelled = true
+          }
           state.emission = .awaiting(nexts)
           return nil
         }
       }?.resume()
+      if cancelled {
+        continuation.resume(returning: nil)
+      }
     }
   }