Cleanup

Author: pblazej

Sources/LiveKit/Broadcast/IPC/BroadcastUploader.swift

@@ -45,8 +45,7 @@ final class BroadcastUploader: Sendable, Loggable {
         self.channel = channel
 
         let messageLoopTask = Task.observing(channel.incomingMessages(BroadcastIPCHeader.self), by: self) { observer, message in
-            let (header, _) = message
-            observer.processMessage(header)
+            observer.processMessageHeader(message.0)
         } onFailure: { observer, error in
             observer.log("IPCChannel returned error: \(error)")
         }
@@ -64,9 +63,6 @@ final class BroadcastUploader: Sendable, Loggable {
 
     /// Close the connection to the receiver.
     func close() {
-        state.mutate {
-            $0.messageLoopTask = nil
-        }
         channel.close()
     }
 
@@ -87,10 +83,10 @@ final class BroadcastUploader: Sendable, Loggable {
             let rotation = VideoRotation(sampleBuffer.replayKitOrientation ?? .up)
             do {
                 let (metadata, imageData) = try imageCodec.encode(sampleBuffer)
-                Task { [weak self, channel] in
+                Task {
                     let header = BroadcastIPCHeader.image(metadata, rotation)
                     try await channel.send(header: header, payload: imageData)
-                    self?.state.mutate { $0.isUploadingImage = false }
+                    state.mutate { $0.isUploadingImage = false }
                 }
             } catch {
                 state.mutate { $0.isUploadingImage = false }
@@ -99,7 +95,7 @@ final class BroadcastUploader: Sendable, Loggable {
         case .audioApp:
             guard state.shouldUploadAudio else { return }
             let (metadata, audioData) = try audioCodec.encode(sampleBuffer)
-            Task { [channel] in
+            Task {
                 let header = BroadcastIPCHeader.audio(metadata)
                 try await channel.send(header: header, payload: audioData)
             }
@@ -108,7 +104,7 @@ final class BroadcastUploader: Sendable, Loggable {
         }
     }
 
-    private func processMessage(_ header: BroadcastIPCHeader) {
+    private func processMessageHeader(_ header: BroadcastIPCHeader) {
         switch header {
         case let .wantsAudio(wantsAudio):
             state.mutate { $0.shouldUploadAudio = wantsAudio }

Sources/LiveKit/Core/DataChannelPair.swift

@@ -58,6 +58,7 @@ class DataChannelPair: NSObject, @unchecked Sendable, Loggable {
     }
 
     private let _state: StateSync<State>
+
     private let eventContinuation: AsyncStream<ChannelEvent>.Continuation
     private var eventLoopTask: AnyTaskCancellable?
 
@@ -255,13 +256,13 @@ class DataChannelPair: NSObject, @unchecked Sendable, Loggable {
         _state = StateSync(State(lossy: lossyChannel,
                                  reliable: reliableChannel))
 
-        let (eventStream, continuation) = AsyncStream.makeStream(of: ChannelEvent.self)
-        eventContinuation = continuation
-
         if let delegate {
             delegates.add(delegate: delegate)
         }
 
+        let (eventStream, continuation) = AsyncStream.makeStream(of: ChannelEvent.self)
+        eventContinuation = continuation
+
         super.init()
 
         eventLoopTask = Task.observing(eventStream, by: self, state: Buffers()) { observer, event, buffers in

Sources/LiveKit/Core/SignalClient.swift

@@ -140,7 +140,7 @@ actor SignalClient: Loggable {
                                              connectOptions: connectOptions)
 
             let messageLoopTask = Task.observing(socket, by: self) { observer, message in
-                await observer._onWebSocketMessage(message: message)
+                await observer.onWebSocketMessage(message)
             } onFailure: { observer, error in
                 await observer.cleanUp(withError: error)
             }
@@ -241,7 +241,7 @@ private extension SignalClient {
         await _requestQueue.processIfResumed(request, elseEnqueue: request.canBeQueued())
     }
 
-    func _onWebSocketMessage(message: URLSessionWebSocketTask.Message) async {
+    func onWebSocketMessage(_ message: URLSessionWebSocketTask.Message) async {
         let response: Livekit_SignalResponse? = switch message {
         case let .data(data): try? Livekit_SignalResponse(serializedBytes: data)
         case let .string(string): try? Livekit_SignalResponse(jsonString: string)

Sources/LiveKit/DataStream/Incoming/IncomingStreamManager.swift

@@ -24,7 +24,7 @@ actor IncomingStreamManager: Loggable {
         let openTime: TimeInterval
         let continuation: StreamReaderSource.Continuation
         var readLength = 0
-        var task: AnyTaskCancellable?
+        var handlerTask: AnyTaskCancellable?
     }
 
     /// Mapping between stream ID and descriptor for open streams.
@@ -133,16 +133,13 @@ actor IncomingStreamManager: Loggable {
             continuation = $0
         }
 
-        var descriptor = Descriptor(
+        let descriptor = Descriptor(
             info: info,
             openTime: Date.timeIntervalSinceReferenceDate,
-            continuation: continuation
+            continuation: continuation,
+            handlerTask: Task.detached { try await handler(source, identity) }.cancellable()
         )
 
-        descriptor.task = Task.detached {
-            try await handler(source, identity)
-        }.cancellable()
-
         openStreams[info.id] = descriptor
     }
 

Tests/LiveKitCoreTests/TaskObserveTests.swift

@@ -68,7 +68,7 @@ class TaskObserveTests: LKTestCase {
 
     func testStreamBreaksWhenObserverDeallocates() async throws {
         var observer: TestObserver? = TestObserver(id: "dealloc-test")
-        weak let weakObserver = observer
+        weak var weakObserver = observer
 
         let (stream, continuation) = AsyncStream.makeStream(of: Int.self)
 
@@ -88,6 +88,7 @@ class TaskObserveTests: LKTestCase {
         try await Task.sleep(nanoseconds: 50_000_000)
 
         XCTAssertNil(weakObserver, "Observer should have been deallocated")
+        weakObserver = nil
 
         continuation.yield(3)
         continuation.yield(4)