Stop encoding when the task is cancelled

SJSAssetExportSession/SampleWriter.swift

@@ -44,6 +44,7 @@ actor SampleWriter {
     private var audioInput: AVAssetWriterInput?
     private var videoOutput: AVAssetReaderVideoCompositionOutput?
     private var videoInput: AVAssetWriterInput?
+    private var isCancelled = false
 
     init(
         asset: sending AVAsset,
@@ -90,15 +91,19 @@ actor SampleWriter {
     }
 
     func writeSamples() async throws {
+        try Task.checkCancellation()
+
         progressContinuation?.yield(0.0)
 
         writer.startWriting()
-        reader.startReading()
         writer.startSession(atSourceTime: timeRange.start)
+        reader.startReading()
+        try Task.checkCancellation()
 
         await encodeVideoTracks()
-        await encodeAudioTracks()
+        try Task.checkCancellation()
 
+        await encodeAudioTracks()
         try Task.checkCancellation()
 
         guard reader.status != .cancelled && writer.status != .cancelled else {
@@ -173,36 +178,78 @@ actor SampleWriter {
         self.videoInput = videoInput
     }
 
+    func cancel() async {
+        isCancelled = true
+    }
+
     // MARK: - Encoding
 
     private func encodeAudioTracks() async {
         // Don't do anything when we have no audio to encode.
         guard audioInput != nil, audioOutput != nil else { return }
 
-        return await withCheckedContinuation { continuation in
+        return await withTaskCancellationHandler {
+            await withCheckedContinuation { continuation in
                 self.audioInput!.requestMediaDataWhenReady(on: queue) {
-                let hasMoreSamples = self.assumeIsolated { _self in
+                    self.assumeIsolated { _self in
-                    _self.writeReadySamples(output: _self.audioOutput!, input: _self.audioInput!)
+                        guard !_self.isCancelled else {
+                            log.debug("Cancelled while encoding audio")
+                            _self.reader.cancelReading()
+                            _self.writer.cancelWriting()
+                            continuation.resume()
+                            return
                         }
+
+                        let hasMoreSamples = _self.writeReadySamples(
+                            output: _self.audioOutput!,
+                            input: _self.audioInput!
+                        )
                         if !hasMoreSamples {
+                            log.debug("Finished encoding audio")
                             continuation.resume()
                         }
                     }
                 }
             }
+        } onCancel: {
+            log.debug("Task cancelled while encoding audio")
+            Task {
+                await self.cancel()
+            }
+        }
+    }
 
     private func encodeVideoTracks() async {
-        return await withCheckedContinuation { continuation in
+        return await withTaskCancellationHandler {
+            await withCheckedContinuation { continuation in
                 self.videoInput!.requestMediaDataWhenReady(on: queue) {
-                let hasMoreSamples = self.assumeIsolated { _self in
+                    self.assumeIsolated { _self in
-                    _self.writeReadySamples(output: _self.videoOutput!, input: _self.videoInput!)
+                        guard !_self.isCancelled else {
+                            log.debug("Cancelled while encoding video")
+                            _self.reader.cancelReading()
+                            _self.writer.cancelWriting()
+                            continuation.resume()
+                            return
                         }
+
+                        let hasMoreSamples = _self.writeReadySamples(
+                            output: _self.videoOutput!,
+                            input: _self.videoInput!
+                        )
                         if !hasMoreSamples {
+                            log.debug("Finished encoding video")
                             continuation.resume()
                         }
                     }
                 }
             }
+        } onCancel: {
+            log.debug("Task cancelled while encoding video")
+            Task {
+                await self.cancel()
+            }
+        }
+    }
 
     private func writeReadySamples(output: AVAssetReaderOutput, input: AVAssetWriterInput) -> Bool {
         while input.isReadyForMoreMediaData {

SJSAssetExportSessionTests/SJSAssetExportSessionTests.swift

@@ -297,4 +297,28 @@ final class ExportSessionTests {
             )
         }
     }
+
+    @Test func test_export_cancellation() async throws {
+        let sourceURL = resourceURL(named: "test-720p-h264-24fps", withExtension: "mov")
+        let destinationURL💥 = makeTemporaryURL()
+        let task = Task {
+            let sourceAsset = AVURLAsset(url: sourceURL, options: [
+                AVURLAssetPreferPreciseDurationAndTimingKey: true,
+            ])
+            let subject = ExportSession()
+            try await subject.export(
+                asset: sourceAsset,
+                video: .codec(.h264, width: 1280, height: 720),
+                to: destinationURL💥.url,
+                as: .mov
+            )
+            Issue.record("Task should be cancelled long before we get here")
+        }
+        NSLog("Sleeping for 0.3s")
+        try await Task.sleep(for: .milliseconds(300))
+        NSLog("Cancelling task")
+        task.cancel()
+        try? await task.value // Wait for task to complete
+        NSLog("Task has finished executing")
+    }
 }