Fix for value leaks its continuation (#2427)

Author: kjaklinovic

RxSwift/Traits/PrimitiveSequence/PrimitiveSequence+Concurrency.swift

@@ -29,10 +29,21 @@ public extension PrimitiveSequenceType where Trait == SingleTrait {
             return try await withTaskCancellationHandler(
                 operation: {
                     try await withCheckedThrowingContinuation { continuation in
+                        var didResume = false
                         disposable.setDisposable(
                             self.subscribe(
-                                onSuccess: { continuation.resume(returning: $0) },
-                                onFailure: { continuation.resume(throwing: $0) }
+                                onSuccess: {
+                                    didResume = true
+                                    continuation.resume(returning: $0)
+                                },
+                                onFailure: {
+                                    didResume = true
+                                    continuation.resume(throwing: $0)
+                                },
+                                onDisposed: {
+                                    guard !didResume else { return }
+                                    continuation.resume(throwing: CancellationError())
+                                }
                             )
                         )
                     }
@@ -69,16 +80,26 @@ public extension PrimitiveSequenceType where Trait == MaybeTrait {
                 operation: {
                     try await withCheckedThrowingContinuation { continuation in
                         var didEmit = false
+                        var didResume = false
                         disposable.setDisposable(
                             self.subscribe(
                                 onSuccess: { value in
                                     didEmit = true
+                                    didResume = true
                                     continuation.resume(returning: value)
                                 },
-                                onError: { error in continuation.resume(throwing: error) },
+                                onError: { error in
+                                    didResume = true
+                                    continuation.resume(throwing: error)
+                                },
                                 onCompleted: {
                                     guard !didEmit else { return }
+                                    didResume = true
                                     continuation.resume(returning: nil)
+                                },
+                                onDisposed: {
+                                    guard !didResume else { return }
+                                    continuation.resume(throwing: CancellationError())
                                 }
                             )
                         )
@@ -114,10 +135,21 @@ public extension PrimitiveSequenceType where Trait == CompletableTrait, Element
             return try await withTaskCancellationHandler(
                 operation: {
                     try await withCheckedThrowingContinuation { continuation in
+                        var didResume = false
                         disposable.setDisposable(
                             self.subscribe(
-                                onCompleted: { continuation.resume() },
-                                onError: { error in continuation.resume(throwing: error) }
+                                onCompleted: {
+                                    didResume = true
+                                    continuation.resume()
+                                },
+                                onError: { error in
+                                    didResume = true
+                                    continuation.resume(throwing: error)
+                                },
+                                onDisposed: {
+                                    guard !didResume else { return }
+                                    continuation.resume(throwing: CancellationError())
+                                }
                             )
                         )
                     }

Tests/RxSwiftTests/PrimitiveSequence+ConcurrencyTests.swift

@@ -41,6 +41,37 @@ extension PrimitiveSequenceConcurrencyTests {
             XCTAssertTrue(true)
         }
     }
+
+    func testSingleThrowsCancellationWithoutEvents() async throws {
+        let single = Single<Void>.never()
+
+        Task {
+            do {
+                try await single.value
+                XCTFail("Should not proceed beyond try")
+            } catch {
+                XCTAssertTrue(Task.isCancelled)
+                XCTAssertTrue(error is CancellationError)
+            }
+        }.cancel()
+    }
+
+    func testSingleNotThrowingCancellation() async throws {
+        let single = Single.just(())
+
+        let task = Task {
+            do {
+                try await single.value
+                XCTAssertTrue(true)
+            } catch {
+                XCTFail()
+            }
+        }
+
+        try await Task.sleep(nanoseconds: 1_000_000)
+        task.cancel()
+    }
+
 }
 
 // MARK: - Maybe
@@ -79,6 +110,49 @@ extension PrimitiveSequenceConcurrencyTests {
             XCTAssertTrue(true)
         }
     }
+
+    func testMaybeThrowsCancellationWithoutEvents() async throws {
+        let maybe = Maybe<Void>.never()
+
+        Task {
+            do {
+                try await maybe.value
+                XCTFail("Should not proceed beyond try")
+            } catch {
+                XCTAssertTrue(Task.isCancelled)
+                XCTAssertTrue(error is CancellationError)
+            }
+        }.cancel()
+    }
+
+    func testMaybeNotThrowingCancellationWhenCompleted() async throws {
+        let maybe = Maybe<Int>.empty()
+
+        Task {
+            do {
+                let value = try await maybe.value
+                XCTAssertNil(value)
+            } catch {
+                XCTFail("Should not throw an error")
+            }
+        }.cancel()
+    }
+
+    func testMaybeNotThrowingCancellation() async throws {
+        let maybe = Maybe.just(())
+
+        let task = Task {
+            do {
+                try await maybe.value
+                XCTAssertTrue(true)
+            } catch {
+                XCTFail("Should not throw an error")
+            }
+        }
+
+        try await Task.sleep(nanoseconds: 1_000_000)
+        task.cancel()
+    }
 }
 
 // MARK: - Completable
@@ -105,6 +179,33 @@ extension PrimitiveSequenceConcurrencyTests {
             XCTAssertTrue(true)
         }
     }
+
+    func testCompletableThrowsCancellationWithoutEvents() async throws {
+        let completable = Completable.never()
+
+        Task {
+            do {
+                try await completable.value
+                XCTFail()
+            } catch {
+                XCTAssertTrue(Task.isCancelled)
+                XCTAssertTrue(error is CancellationError)
+            }
+        }.cancel()
+    }
+
+    func testCompletableNotThrowingCancellation() async throws {
+        let completable = Completable.empty()
+
+        Task {
+            do {
+                try await completable.value
+                XCTAssertTrue(true)
+            } catch {
+                XCTFail("Should not throw an error")
+            }
+        }.cancel()
+    }
 }
 #endif