Merge pull request #1655 from shogo4405/feature/fix-memory-leak

fix memory-leak a MediaMixer instance.
shogo4405 · Jan 1, 2025 · b3ceb9e · b3ceb9e
2 parents 203e873 + 96d6ff7
commit b3ceb9e
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 14 deletions.
diff --git a/HaishinKit/Sources/Mixer/AudioCaptureUnit.swift b/HaishinKit/Sources/Mixer/AudioCaptureUnit.swift
@@ -24,9 +24,9 @@ final class AudioCaptureUnit: CaptureUnit {
         return audioMixer.inputFormats
     }
     var output: AsyncStream<(AVAudioPCMBuffer, AVAudioTime)> {
-        let (stream, continutation) = AsyncStream<(AVAudioPCMBuffer, AVAudioTime)>.makeStream()
-        self.continutation = continutation
-        return stream
+        AsyncStream<(AVAudioPCMBuffer, AVAudioTime)> { continutation in
+            self.continutation = continutation
+        }
     }
     private lazy var audioMixer: any AudioMixer = {
         if isMultiTrackAudioMixingEnabled {
@@ -106,6 +106,10 @@ final class AudioCaptureUnit: CaptureUnit {
             break
         }
     }
+
+    func finish() {
+        continutation?.finish()
+    }
 }
 
 extension AudioCaptureUnit: AudioMixerDelegate {

diff --git a/HaishinKit/Sources/Mixer/MediaMixer.swift b/HaishinKit/Sources/Mixer/MediaMixer.swift
@@ -346,6 +346,9 @@ public final actor MediaMixer {
     }
 
     func setVideoRenderingMode(_ mode: VideoMixerSettings.Mode) {
+        guard isRunning else {
+            return
+        }
         switch mode {
         case .passthrough:
             Task { @ScreenActor in
@@ -355,7 +358,7 @@ public final actor MediaMixer {
             Task { @ScreenActor in
                 displayLink.preferredFramesPerSecond = await Int(frameRate)
                 displayLink.startRunning()
-                for await updateFrame in displayLink.updateFrames where displayLink.isRunning {
+                for await updateFrame in displayLink.updateFrames {
                     guard let buffer = screen.makeSampleBuffer(updateFrame) else {
                         continue
                     }
@@ -390,7 +393,7 @@ extension MediaMixer: AsyncRunner {
         }
         isRunning = true
         Task {
-            for await inputs in videoIO.inputs where isRunning {
+            for await inputs in videoIO.inputs {
                 Task { @ScreenActor in
                     var sampleBuffer = inputs.1
                     screen.append(inputs.0, buffer: sampleBuffer)
@@ -405,14 +408,14 @@ extension MediaMixer: AsyncRunner {
             }
         }
         Task {
-            for await video in videoIO.output where isRunning {
+            for await video in videoIO.output {
                 for output in outputs where await output.videoTrackId == UInt8.max {
                     output.mixer(self, didOutput: video)
                 }
             }
         }
         Task {
-            for await audio in audioIO.output where isRunning {
+            for await audio in audioIO.output {
                 for output in outputs where await output.audioTrackId == UInt8.max {
                     output.mixer(self, didOutput: audio.0, when: audio.1)
                 }
@@ -450,6 +453,8 @@ extension MediaMixer: AsyncRunner {
         if useManualCapture {
             session.stopRunning()
         }
+        audioIO.finish()
+        videoIO.finish()
         cancellables.forEach { $0.cancel() }
         cancellables.removeAll()
         Task { @ScreenActor in

diff --git a/HaishinKit/Sources/Mixer/VideoCaptureUnit.swift b/HaishinKit/Sources/Mixer/VideoCaptureUnit.swift
@@ -62,15 +62,15 @@ final class VideoCaptureUnit: CaptureUnit {
     #endif
 
     var inputs: AsyncStream<(UInt8, CMSampleBuffer)> {
-        let (stream, continutation) = AsyncStream<(UInt8, CMSampleBuffer)>.makeStream()
-        self.inputsContinuation = continutation
-        return stream
+        AsyncStream<(UInt8, CMSampleBuffer)> { continutation in
+            self.inputsContinuation = continutation
+        }
     }
 
     var output: AsyncStream<CMSampleBuffer> {
-        let (stream, continutation) = AsyncStream<CMSampleBuffer>.makeStream()
-        self.outputContinuation = continutation
-        return stream
+        AsyncStream<CMSampleBuffer> { continutation in
+            self.outputContinuation = continutation
+        }
     }
 
     private lazy var videoMixer = {
@@ -153,6 +153,11 @@ final class VideoCaptureUnit: CaptureUnit {
         return .init(track: track, videoMixer: videoMixer)
     }
 
+    func finish() {
+        inputsContinuation?.finish()
+        outputContinuation?.finish()
+    }
+
     @available(tvOS 17.0, *)
     private func device(for track: UInt8) -> VideoDeviceUnit? {
         #if os(tvOS)

diff --git a/HaishinKit/Sources/Screen/DisplayLinkChoreographer.swift b/HaishinKit/Sources/Screen/DisplayLinkChoreographer.swift
@@ -141,7 +141,8 @@ extension DisplayLinkChoreographer: Runner {
         guard isRunning else {
             return
         }
-        displayLink = nil
         isRunning = false
+        displayLink = nil
+        continutation?.finish()
     }
 }