From fa1f39bb2b60b0c219f764e07b4fe7b2e31940fb Mon Sep 17 00:00:00 2001
From: Sami Samhuri <sami@samhuri.net>
Date: Wed, 10 Jul 2024 23:30:28 -0700
Subject: [PATCH] Stop encoding when the task is cancelled

---
 SJSAssetExportSession/SampleWriter.swift      | 79 +++++++++++++++----
 .../SJSAssetExportSessionTests.swift          | 24 ++++++
 2 files changed, 87 insertions(+), 16 deletions(-)

diff --git a/SJSAssetExportSession/SampleWriter.swift b/SJSAssetExportSession/SampleWriter.swift
index a5b65d2..eff3700 100644
--- a/SJSAssetExportSession/SampleWriter.swift
+++ b/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,34 +178,76 @@ 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
-            self.audioInput!.requestMediaDataWhenReady(on: queue) {
-                let hasMoreSamples = self.assumeIsolated { _self in
-                    _self.writeReadySamples(output: _self.audioOutput!, input: _self.audioInput!)
-                }
-                if !hasMoreSamples {
-                    continuation.resume()
+        return await withTaskCancellationHandler {
+            await withCheckedContinuation { continuation in
+                self.audioInput!.requestMediaDataWhenReady(on: queue) {
+                    self.assumeIsolated { _self in
+                        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
-            self.videoInput!.requestMediaDataWhenReady(on: queue) {
-                let hasMoreSamples = self.assumeIsolated { _self in
-                    _self.writeReadySamples(output: _self.videoOutput!, input: _self.videoInput!)
-                }
-                if !hasMoreSamples {
-                    continuation.resume()
+        return await withTaskCancellationHandler {
+            await withCheckedContinuation { continuation in
+                self.videoInput!.requestMediaDataWhenReady(on: queue) {
+                    self.assumeIsolated { _self in
+                        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()
+            }
         }
     }
 
diff --git a/SJSAssetExportSessionTests/SJSAssetExportSessionTests.swift b/SJSAssetExportSessionTests/SJSAssetExportSessionTests.swift
index 17ec16b..3d30e64 100644
--- a/SJSAssetExportSessionTests/SJSAssetExportSessionTests.swift
+++ b/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")
+    }
 }