From acde14a2e8a5e74f5cbf757547ff96d55ed44dd9 Mon Sep 17 00:00:00 2001
From: Ben Cherry <ben@livekit.io>
Date: Thu, 5 Dec 2024 00:34:23 -0800
Subject: [PATCH] Actually stop the track when the broadcast extension socket
 is closed (#520)

This closes https://github.com/livekit/client-sdk-swift/issues/446

Two fixes:

~1. It seems that stopping the broadcast via tapping on the system
indicator and choosing "stop broadcast" calls the `broadcastPaused`
method instead of the `broadcastFinished` method on the SampleHandler.
The only way I could trigger `broadcastFinished` in testing was by
locking my device. I could not find a way to trigger `broadcastResumed`
so I just went ahead and treated "pause" as final, and close the socket
(I think maybe this is vestigial from older iterations of ReplayKit? not
sure, it's not an area I'm super familiar with)~
2. When the socket closes the frame reader stops but never notified the
BroadcastScreenCapturer that owns it, so the track would stay published
even as it stopped being updated. This fix adds a new callback method to
pass the closure up the stack and stop the track.

Not sure what happened in original testing but actually
broadcastFinished _is_ called as expected, and broadcastPaused is called
when you open the "stop capture" dialog, not when you finish it. I also
fixed it to properly unpublish the track, not sure why initial testing
showed simply ending capture to be enough.

---------

Co-authored-by: hiroshihorie <548776+hiroshihorie@users.noreply.github.com>
---
 .../Broadcast/BroadcastScreenCapturer.swift      |  6 ++++++
 .../Broadcast/SocketConnectionFrameReader.swift  |  2 ++
 .../LocalTrackPublication.swift                  | 16 ++++++++++++++++
 3 files changed, 24 insertions(+)

diff --git a/Sources/LiveKit/Broadcast/BroadcastScreenCapturer.swift b/Sources/LiveKit/Broadcast/BroadcastScreenCapturer.swift
index b7bb38048..3c77a7efe 100644
--- a/Sources/LiveKit/Broadcast/BroadcastScreenCapturer.swift
+++ b/Sources/LiveKit/Broadcast/BroadcastScreenCapturer.swift
@@ -62,6 +62,12 @@ class BroadcastScreenCapturer: BufferCapturer {
         frameReader.didCapture = { pixelBuffer, rotation in
             self.capture(pixelBuffer, rotation: rotation.toLKType())
         }
+        frameReader.didEnd = { [weak self] in
+            guard let self else { return }
+            Task {
+                try await self.stopCapture()
+            }
+        }
         frameReader.startCapture(with: socketConnection)
         self.frameReader = frameReader
 
diff --git a/Sources/LiveKit/Broadcast/SocketConnectionFrameReader.swift b/Sources/LiveKit/Broadcast/SocketConnectionFrameReader.swift
index a44c645b6..9ce40a485 100644
--- a/Sources/LiveKit/Broadcast/SocketConnectionFrameReader.swift
+++ b/Sources/LiveKit/Broadcast/SocketConnectionFrameReader.swift
@@ -137,6 +137,7 @@ class SocketConnectionFrameReader: NSObject {
 
     private var message: Message?
     var didCapture: ((CVPixelBuffer, RTCVideoRotation) -> Void)?
+    var didEnd: (() -> Void)?
 
     override init() {}
 
@@ -227,6 +228,7 @@ extension SocketConnectionFrameReader: StreamDelegate {
         case .endEncountered:
             logger.log(level: .debug, "server stream end encountered")
             stopCapture()
+            didEnd?()
         case .errorOccurred:
             logger.log(level: .debug, "server stream error encountered: \(aStream.streamError?.localizedDescription ?? "")")
         default:
diff --git a/Sources/LiveKit/TrackPublications/LocalTrackPublication.swift b/Sources/LiveKit/TrackPublications/LocalTrackPublication.swift
index 8270ea4a3..bde84b213 100644
--- a/Sources/LiveKit/TrackPublications/LocalTrackPublication.swift
+++ b/Sources/LiveKit/TrackPublications/LocalTrackPublication.swift
@@ -105,6 +105,22 @@ extension LocalTrackPublication: VideoCapturerDelegate {
             }
         }
     }
+
+    public func capturer(_ capturer: VideoCapturer, didUpdate state: VideoCapturer.CapturerState) {
+        // Broadcasts can always be stopped from system UI that bypasses our normal disable & unpublish methods.
+        // This check ensures that when this happens the track gets unpublished as well.
+        #if os(iOS)
+        if state == .stopped, capturer is BroadcastScreenCapturer {
+            Task {
+                guard let participant = try await self.requireParticipant() as? LocalParticipant else {
+                    return
+                }
+
+                try await participant.unpublish(publication: self)
+            }
+        }
+        #endif
+    }
 }
 
 extension LocalTrackPublication {