Fix memory leaks

Author: carson-katri

Package.resolved

@@ -50,7 +50,7 @@ public class LiveSessionCoordinator<R: RootRegistry>: ObservableObject {
     var socket: LiveViewNativeCore.Socket?
 
     private var liveReloadChannel: LiveViewNativeCore.LiveChannel?
-    private var liveReloadListener: AsyncThrowingStream<LiveViewNativeCore.EventPayload, any Error>?
+    private var liveReloadListener: Channel.EventStream?
     private var liveReloadListenerLoop: Task<(), any Error>?
 
     private var cancellables = Set<AnyCancellable>()

Sources/LiveViewNative/Coordinators/LiveSessionCoordinator.swift

@@ -54,8 +54,9 @@ public class LiveViewCoordinator<R: RootRegistry>: ObservableObject {
     private(set) internal var eventSubject = PassthroughSubject<(String, Payload), Never>()
     private(set) internal var eventHandlers = Set<AnyCancellable>()
 
-    private var eventListener: AsyncThrowingStream<LiveViewNativeCore.EventPayload, any Error>?
+    private var eventListener: Channel.EventStream?
     private var eventListenerLoop: Task<(), any Error>?
+    private var statusListener: Channel.StatusStream?
     private var statusListenerLoop: Task<(), any Error>?
 
     private(set) internal var liveViewModel = LiveViewModel()
@@ -289,10 +290,11 @@ public class LiveViewCoordinator<R: RootRegistry>: ObservableObject {
         let channel = liveChannel.channel()
         self.channel = channel
 
+        let statusListener = channel.statusStream()
+        self.statusListener = statusListener
         statusListenerLoop = Task { @MainActor [weak self] in
-            for try await status in channel.statusStream() {
-                guard let self else { return }
-                self.internalState = switch status {
+            for try await status in statusListener {
+                self?.internalState = switch status {
                 case .joined:
                     .connected
                 case .joining, .waitingForSocketToConnect, .waitingToJoin:

Sources/LiveViewNative/Coordinators/LiveViewCoordinator.swift

@@ -57,7 +57,7 @@ struct CoordinatorEnvironment {
     private final class Storage {
         let pushEvent: @MainActor (String, String, Any, Int?) async throws -> [String:Any]?
         let elementChanged: @MainActor (NodeRef) -> ObservableObjectPublisher
-        let document: Document
+        weak var document: Document?
 
         init<R: CustomRegistry>(_ coordinator: LiveViewCoordinator<R>, document: Document) {
             self.pushEvent = coordinator.pushEvent
@@ -73,7 +73,7 @@ struct CoordinatorEnvironment {
     var elementChanged: @MainActor (NodeRef) -> ObservableObjectPublisher {
         storage.elementChanged
     }
-    var document: Document {
+    var document: Document? {
         storage.document
     }
 

Sources/LiveViewNative/Environment.swift

@@ -231,7 +231,7 @@ public struct LiveView<
         .environment(\.stylesheet, session.stylesheet ?? .init(content: [], classes: [:]))
         .environment(\.reconnectLiveView, .init(baseURL: session.url, action: session.reconnect))
         .environmentObject(session)
-        .task(priority: .userInitiated) {
+        .task {
             await session.connect()
         }
         .onChange(of: scenePhase) { newValue in

Sources/LiveViewNative/Live/LiveView.swift

@@ -145,30 +145,38 @@ public struct Event: @preconcurrency DynamicProperty, @preconcurrency Decodable
 
         init() {}
 
+        deinit {
+            self.handlerTask?.cancel()
+        }
+        
         func update(coordinator: CoordinatorEnvironment?, debounce: Double?, throttle: Double?) {
             guard handlerTask == nil || debounce != self.debounce || throttle != self.throttle
             else { return }
             handlerTask?.cancel()
             self.debounce = debounce
             self.throttle = throttle
+            let pushEvent = coordinator?.pushEvent
             if let debounce = debounce {
-                handlerTask = Task { [weak didSendSubject] in
+                handlerTask = Task { [weak channel, weak didSendSubject, pushEvent] in
+                    guard let channel else { return }
                     for await event in channel.debounce(for: .milliseconds(debounce)) {
-                        _ = try await coordinator?.pushEvent(event.type, event.event, event.payload, event.target)
+                        _ = try await pushEvent?(event.type, event.event, event.payload, event.target)
                         didSendSubject?.send()
                     }
                 }
             } else if let throttle = throttle {
-                handlerTask = Task { @MainActor [weak didSendSubject] in
-                    for await event in channel.throttle(for: .milliseconds(throttle)) {
-                        _ = try await coordinator?.pushEvent(event.type, event.event, event.payload, event.target)
+                handlerTask = Task { @MainActor [weak channel, weak didSendSubject, pushEvent] in
+                    guard let channel else { return }
+                    for await event in channel._throttle(for: .milliseconds(throttle)) {
+                        _ = try await pushEvent?(event.type, event.event, event.payload, event.target)
                         didSendSubject?.send()
                     }
                 }
             } else {
-                handlerTask = Task { @MainActor [weak didSendSubject] in
+                handlerTask = Task { @MainActor [weak channel, weak didSendSubject, pushEvent] in
+                    guard let channel else { return }
                     for await event in channel {
-                        _ = try await coordinator?.pushEvent(event.type, event.event, event.payload, event.target)
+                        _ = try await pushEvent?(event.type, event.event, event.payload, event.target)
                         didSendSubject?.send()
                     }
                 }

Sources/LiveViewNative/Property Wrappers/Event.swift

@@ -140,7 +140,7 @@ extension ObservedElement {
             observeChildren: Bool
         ) {
             guard cancellable == nil || (observeChildren && self.observedChildIDs != self.resolvedChildIDs) else { return }
-            self.resolvedElement = context.document[id].asElement()
+            self.resolvedElement = context.document?[id].asElement()
             self.resolvedChildren = Array(self.resolvedElement.children())
             self._resolvedChildIDs = nil
 
@@ -161,7 +161,7 @@ extension ObservedElement {
             cancellable = self.elementChangedPublisher
                 .sink { [weak self] _ in
                     guard let self else { return }
-                    self.resolvedElement = context.document[id].asElement()
+                    self.resolvedElement = context.document?[id].asElement()
                     self.resolvedChildren = Array(self.resolvedElement.children())
                     self._resolvedChildIDs = nil
                     self.objectWillChange.send()

Sources/LiveViewNative/Property Wrappers/ObservedElement.swift

@@ -108,14 +108,28 @@ struct ImageView<Root: RootRegistry>: View {
             }
         })
     }
+    
+    var fileUploadImage: SwiftUI.Image? {
+        guard let phxUploadRef
+        else { return nil }
+        #if os(macOS)
+        return liveViewModel
+            .fileUpload(id: phxUploadRef)
+            .flatMap({ NSImage(data: $0.data) })
+            .flatMap(Image.init(nsImage:))
+        #else
+        return liveViewModel
+            .fileUpload(id: phxUploadRef)
+            .flatMap({ UIImage(data: $0.data) })
+            .flatMap(Image.init(uiImage:))
+        #endif
+    }
 
     var image: SwiftUI.Image? {
         if let overrideImage {
             return overrideImage
-        } else if let phxUploadRef,
-                  let image = liveViewModel.fileUpload(id: phxUploadRef).flatMap({ UIImage(data: $0.data) })
-        {
-            return SwiftUI.Image(uiImage: image)
+        } else if let fileUploadImage {
+            return fileUploadImage
         } else if let systemName {
             return SwiftUI.Image(systemName: systemName, variableValue: variableValue)
         } else if let name {

Sources/LiveViewNative/Views/Images/ImageView.swift

@@ -59,18 +59,52 @@ extension Node: Identifiable {
 }
 
 extension Channel {
-    public func eventStream() -> AsyncThrowingStream<EventPayload, any Error> {
-        let events = self.events()
-        return AsyncThrowingStream(unfolding: {
-            return try await events.event()
-        })
+    struct EventStream: AsyncSequence {
+        let events: Events
+        
+        init(for channel: Channel) {
+            self.events = channel.events()
+        }
+        
+        func makeAsyncIterator() -> AsyncIterator {
+            .init(events: events)
+        }
+        
+        struct AsyncIterator: AsyncIteratorProtocol {
+            let events: Events
+            
+            func next() async throws -> EventPayload? {
+                try await events.event()
+            }
+        }
+    }
+    
+    func eventStream() -> EventStream {
+        return EventStream(for: self)
+    }
+    
+    final class StatusStream: AsyncSequence {
+        let statuses: ChannelStatuses
+        
+        init(for channel: Channel) {
+            self.statuses = channel.statuses()
+        }
+        
+        func makeAsyncIterator() -> AsyncIterator {
+            .init(statuses: statuses)
+        }
+        
+        struct AsyncIterator: AsyncIteratorProtocol {
+            let statuses: ChannelStatuses
+            
+            func next() async throws -> ChannelStatus? {
+                try await statuses.status()
+            }
+        }
     }
 
-    public func statusStream() -> AsyncThrowingStream<ChannelStatus, any Error> {
-        let statuses = self.statuses()
-        return AsyncThrowingStream(unfolding: {
-            return try await statuses.status()
-        })
+    func statusStream() -> StatusStream {
+        StatusStream(for: self)
     }
 }