Cancel unstructured task when current task is cancelled.

sideeffect-io · Aug 2, 2022 · 27d1fbd · 27d1fbd
1 parent e31495e
commit 27d1fbd
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/Sources/Operators/AsyncSequence+SwitchToLatest.swift b/Sources/Operators/AsyncSequence+SwitchToLatest.swift
@@ -74,6 +74,11 @@ public struct AsyncSwitchToLatestSequence<UpstreamAsyncSequence: AsyncSequence>:
             self.currentTask = task
         }
 
+        @Sendable
+        func cancelCurrentTask() {
+            self.currentTask?.cancel()
+        }
+
         /// iterates over the upstream sequence and maintain the current async iterator while cancelling the current .next() task for each new element
         func startUpstreamIterator() async throws {
             guard !self.hasStarted else { return }
@@ -130,7 +135,7 @@ public struct AsyncSwitchToLatestSequence<UpstreamAsyncSequence: AsyncSequence>:
 
             // if a task is cancelled while waiting with the next element (a new element arrived in the root iterator)
             // we create a new task and wait for the elements from the new child iterator
-            while noValueHasBeenEmitted {
+            while noValueHasBeenEmitted, !Task.isCancelled {
                 currentTask = Task {
                     do {
                         return try await localUpstreamIteratorManager.nextOnCurrentChildIterator()
@@ -141,7 +146,9 @@ public struct AsyncSwitchToLatestSequence<UpstreamAsyncSequence: AsyncSequence>:
                     }
                 }
                 localUpstreamIteratorManager.setCurrentTask(task: currentTask)
-                emittedElement = try await currentTask.value
+                emittedElement = try await withTaskCancellationHandler(handler: localUpstreamIteratorManager.cancelCurrentTask) {
+                    try await currentTask.value
+                }
                 noValueHasBeenEmitted = (emittedElement == nil && currentTask.isCancelled)
             }