From 33152e4e4439519546eeda5c50a15cb5d4e55ce2 Mon Sep 17 00:00:00 2001
From: Sami Samhuri <sami@samhuri.net>
Date: Sun, 19 Jan 2025 14:05:46 -0800
Subject: [PATCH] Try to improve cancellation response, memory might be leaking

---
 Sources/SJSAssetExportSession/SampleWriter.swift | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/Sources/SJSAssetExportSession/SampleWriter.swift b/Sources/SJSAssetExportSession/SampleWriter.swift
index 7ae4dad..d6c8dd4 100644
--- a/Sources/SJSAssetExportSession/SampleWriter.swift
+++ b/Sources/SJSAssetExportSession/SampleWriter.swift
@@ -129,7 +129,10 @@ actor SampleWriter {
             try await Task.sleep(for: .milliseconds(10))
         }
 
-        guard reader.status != .cancelled && writer.status != .cancelled else {
+        guard !isCancelled, reader.status != .cancelled, writer.status != .cancelled else {
+            log.debug("Cancelled before writing samples")
+            reader.cancelReading()
+            writer.cancelWriting()
             throw CancellationError()
         }
         guard writer.status != .failed else {
@@ -253,6 +256,12 @@ actor SampleWriter {
 
     private func writeReadySamples(output: AVAssetReaderOutput, input: AVAssetWriterInput) -> Bool {
         while input.isReadyForMoreMediaData {
+            guard !isCancelled else {
+                log.debug("Cancelled while writing samples")
+                reader.cancelReading()
+                writer.cancelWriting()
+                return false
+            }
             guard reader.status == .reading && writer.status == .writing,
                   let sampleBuffer = output.copyNextSampleBuffer() else {
                 input.markAsFinished()