Improve Home Assistant connection reliability

XIONControlPanel/Servers/ServerMultiplex.swift

@@ -8,11 +8,12 @@
 
 import Foundation
 
-protocol ServerMultiplexDelegate
+protocol ServerMultiplexDelegate: AnyObject
 {
     func serverMultiplex(_ multiplex: ServerMultiplex, didAddDevices devices: [AnyDevice])
     func serverMultiplex(_ multiplex: ServerMultiplex, devicesStateChanged devices: [AnyDevice])
     func serverMultiplex(_ multiplex: ServerMultiplex, didReceiveAcknowledgementFromServer server: Server)
+    func serverMultiplexConnectionStatusDidChange(_ multiplex: ServerMultiplex)
     func serverMultiplex(_ multiplex: ServerMultiplex, didEncounterError error: Error)
 }
 
@@ -23,10 +24,11 @@ enum ServerMultiplexError : Error
 
 class ServerMultiplex
 {
-    public var delegate: ServerMultiplexDelegate?
+    public weak var delegate: ServerMultiplexDelegate?
     public private(set) var devices = Set<AnyDevice>()
     
     private var servers: [Server] = []
+    private var devicesByServer: [ObjectIdentifier: Set<AnyDevice>] = [:]
     public var numServers: Int { return servers.count }
     
     public func addServer(_ server: Server)
@@ -73,31 +75,43 @@ class ServerMultiplex
     
     public func refreshDevices()
     {
-        self.servers.forEach { server in
+        servers.forEach { server in
             server.connect { error in
                 if let error {
                     print("Multiplexer: error connecting server \(server): \(error)")
+                    self.handleError(forServer: server, error: error)
                 }
             }
-        }
-        
-        self.servers.forEach { (server: Server) in
+
             server.fetchDevices { (result: Result<[AnyDevice], Error>) in
                 self.handleServerFetchResult(forServer: server, result: result)
             }
         }
     }
+
+    public func disconnect()
+    {
+        servers.forEach { server in
+            server.disconnect { error in
+                if let error {
+                    self.handleError(forServer: server, error: error)
+                }
+            }
+        }
+    }
 }
 
 extension ServerMultiplex
 {
     private func handleServerFetchResult(forServer server: Server, result: Result<[AnyDevice], Error>)
     {
-        switch result {
-        case .success(let devices):
-            handleDevicesChanged(forServer: server, devicesChanged: devices)
-        case .failure(let error):
-            handleError(forServer: server, error: error)
+        DispatchQueue.main.async {
+            switch result {
+            case .success(let devices):
+                self.handleDevicesChanged(forServer: server, devicesChanged: devices)
+            case .failure(let error):
+                self.handleError(forServer: server, error: error)
+            }
         }
     }
     
@@ -107,28 +121,29 @@ extension ServerMultiplex
         self.delegate?.serverMultiplex(self, didReceiveAcknowledgementFromServer: server)
         
         // Then, optionally notify about new devices or device state changes
+        let serverID = ObjectIdentifier(server)
+        let previousDevices = devicesByServer[serverID] ?? []
         let newDevicesSet = Set<AnyDevice>(devicesChanged)
-        let additions = newDevicesSet.subtracting(self.devices)
+        let additions = newDevicesSet.subtracting(previousDevices)
+        let removals = previousDevices.subtracting(newDevicesSet)
         let changed = newDevicesSet.filter { (device: AnyDevice) in
-            if let existing = (devices.first { $0.hashValue == device.hashValue }) {
+            if let existing = previousDevices.first(where: { $0.serial == device.serial }) {
                 return existing.state != device.state
             }
             
             return false
         }
-        
-        self.devices = self.devices.union(newDevicesSet)
-        
-        if additions.count > 0 {
-            DispatchQueue.main.async {
-                self.delegate?.serverMultiplex(self, didAddDevices: Array(additions))
-            }
+
+        devicesByServer[serverID] = newDevicesSet
+        devices.subtract(previousDevices)
+        devices.formUnion(newDevicesSet)
+
+        if !additions.isEmpty || !removals.isEmpty {
+            delegate?.serverMultiplex(self, didAddDevices: Array(additions))
         }
         
-        if changed.count > 0 {
-            DispatchQueue.main.async {
-                self.delegate?.serverMultiplex(self, devicesStateChanged: Array(changed))
-            }
+        if !changed.isEmpty {
+            delegate?.serverMultiplex(self, devicesStateChanged: Array(changed))
         }
     }
     
@@ -143,9 +158,21 @@ extension ServerMultiplex
 extension ServerMultiplex: ServerDelegate
 {
     func server(_ server: any Server, deviceChangedState subjectDevice: AnyDevice) {
-        guard let device = devices.first(where: { $0.serial == subjectDevice.serial }) else { return }
-        device.state = subjectDevice.state
-        
-        delegate?.serverMultiplex(self, devicesStateChanged: [device])
+        DispatchQueue.main.async {
+            guard let device = self.devices.first(where: {
+                $0.serial == subjectDevice.serial
+            }) else {
+                return
+            }
+
+            device.state = subjectDevice.state
+            self.delegate?.serverMultiplex(self, devicesStateChanged: [device])
+        }
+    }
+
+    func server(_ server: any Server, connectionStatusChanged status: ConnectionStatus) {
+        DispatchQueue.main.async {
+            self.delegate?.serverMultiplexConnectionStatusDidChange(self)
+        }
     }
 }