From 2dac7d58dc0fd492ddbd81f69293d4535f63a0f3 Mon Sep 17 00:00:00 2001
From: Sami Samhuri <sami@samhuri.net>
Date: Sun, 19 Jan 2025 15:01:41 -0800
Subject: [PATCH] Try to fix a possible retain cycle

---
 .../SJSAssetExportSession/SampleWriter.swift  | 80 ++++++++++++-------
 1 file changed, 49 insertions(+), 31 deletions(-)

diff --git a/Sources/SJSAssetExportSession/SampleWriter.swift b/Sources/SJSAssetExportSession/SampleWriter.swift
index dcdf51c..5e114b2 100644
--- a/Sources/SJSAssetExportSession/SampleWriter.swift
+++ b/Sources/SJSAssetExportSession/SampleWriter.swift
@@ -41,8 +41,8 @@ actor SampleWriter {
 
     // MARK: Internal state
 
-    private let reader: AVAssetReader
-    private let writer: AVAssetWriter
+    private var reader: AVAssetReader?
+    private var writer: AVAssetWriter?
     private var audioOutput: AVAssetReaderAudioMixOutput?
     private var audioInput: AVAssetWriterInput?
     private var videoOutput: AVAssetReaderVideoCompositionOutput?
@@ -111,15 +111,15 @@ actor SampleWriter {
 
         progressContinuation.yield(0.0)
 
-        writer.startWriting()
-        writer.startSession(atSourceTime: timeRange.start)
-        reader.startReading()
+        writer?.startWriting()
+        writer?.startSession(atSourceTime: timeRange.start)
+        reader?.startReading()
         try Task.checkCancellation()
 
         startEncodingAudioTracks()
         startEncodingVideoTracks()
 
-        while reader.status == .reading, writer.status == .writing {
+        while reader?.status == .reading, writer?.status == .writing {
             guard !Task.isCancelled else {
                 // Flag so that we stop writing samples
                 isCancelled = true
@@ -129,24 +129,29 @@ actor SampleWriter {
             try await Task.sleep(for: .milliseconds(10))
         }
 
-        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 {
-            reader.cancelReading()
-            throw Error.writeFailure(writer.error)
-        }
-        guard reader.status != .failed else {
-            writer.cancelWriting()
-            throw Error.readFailure(reader.error)
+        defer {
+            cleanUp()
         }
 
-        await withCheckedContinuation { continuation in
-            writer.finishWriting {
-                continuation.resume(returning: ())
+        guard !isCancelled, reader?.status != .cancelled, writer?.status != .cancelled else {
+            log.debug("Cancelled before writing samples")
+            cancelReaderAndWriter()
+            throw CancellationError()
+        }
+        guard writer?.status != .failed else {
+            reader?.cancelReading()
+            throw Error.writeFailure(writer?.error)
+        }
+        guard reader?.status != .failed else {
+            writer?.cancelWriting()
+            throw Error.readFailure(reader?.error)
+        }
+
+        if let writer {
+            await withCheckedContinuation { continuation in
+                writer.finishWriting {
+                    continuation.resume(returning: ())
+                }
             }
         }
 
@@ -169,7 +174,7 @@ actor SampleWriter {
         let audioOutput = AVAssetReaderAudioMixOutput(audioTracks: audioTracks, audioSettings: nil)
         audioOutput.alwaysCopiesSampleData = false
         audioOutput.audioMix = audioMix
-        guard reader.canAdd(audioOutput) else {
+        guard let reader, reader.canAdd(audioOutput) else {
             throw Error.setupFailure(.cannotAddAudioOutput)
         }
         reader.add(audioOutput)
@@ -177,7 +182,7 @@ actor SampleWriter {
 
         let audioInput = AVAssetWriterInput(mediaType: .audio, outputSettings: audioOutputSettings)
         audioInput.expectsMediaDataInRealTime = false
-        guard writer.canAdd(audioInput) else {
+        guard let writer, writer.canAdd(audioInput) else {
             throw Error.setupFailure(.cannotAddAudioInput)
         }
         writer.add(audioInput)
@@ -193,7 +198,7 @@ actor SampleWriter {
         )
         videoOutput.alwaysCopiesSampleData = false
         videoOutput.videoComposition = videoComposition
-        guard reader.canAdd(videoOutput) else {
+        guard let reader, reader.canAdd(videoOutput) else {
             throw Error.setupFailure(.cannotAddVideoOutput)
         }
         reader.add(videoOutput)
@@ -201,7 +206,7 @@ actor SampleWriter {
 
         let videoInput = AVAssetWriterInput(mediaType: .video, outputSettings: videoOutputSettings)
         videoInput.expectsMediaDataInRealTime = false
-        guard writer.canAdd(videoInput) else {
+        guard let writer, writer.canAdd(videoInput) else {
             throw Error.setupFailure(.cannotAddVideoInput)
         }
         writer.add(videoInput)
@@ -233,11 +238,25 @@ actor SampleWriter {
         }
     }
 
+    private func cancelReaderAndWriter() {
+        reader?.cancelReading()
+        writer?.cancelWriting()
+        cleanUp()
+    }
+
+    private func cleanUp() {
+        reader = nil
+        writer = nil
+        audioInput = nil
+        audioOutput = nil
+        videoInput = nil
+        videoOutput = nil
+    }
+
     private func writeAllReadySamples() {
         guard !isCancelled else {
             log.debug("Cancelled while writing samples")
-            reader.cancelReading()
-            writer.cancelWriting()
+            cancelReaderAndWriter()
             return
         }
 
@@ -254,11 +273,10 @@ actor SampleWriter {
         while input.isReadyForMoreMediaData {
             guard !isCancelled else {
                 log.debug("Cancelled while writing samples")
-                reader.cancelReading()
-                writer.cancelWriting()
+                cancelReaderAndWriter()
                 return false
             }
-            guard reader.status == .reading && writer.status == .writing,
+            guard reader?.status == .reading && writer?.status == .writing,
                   let sampleBuffer = output.copyNextSampleBuffer() else {
                 input.markAsFinished()
                 return false