Ping/Pong heartbeat logic

Sources/LiveKit/Core/SignalClient.swift

@@ -56,6 +56,9 @@ internal class SignalClient: MulticastDelegate<SignalClientDelegate> {
     private var webSocket: WebSocket?
     private var latestJoinResponse: Livekit_JoinResponse?
 
+    private var pingIntervalTimer: DispatchQueueTimer?
+    private var pingTimeoutTimer: DispatchQueueTimer?
+
     init() {
         super.init()
 
@@ -66,7 +69,8 @@ internal class SignalClient: MulticastDelegate<SignalClientDelegate> {
 
             guard let self = self else { return }
 
-            if oldState.connectionState != state.connectionState {
+            // connectionState did update
+            if state.connectionState != oldState.connectionState {
                 self.log("\(oldState.connectionState) -> \(state.connectionState)")
             }
 
@@ -106,7 +110,6 @@ internal class SignalClient: MulticastDelegate<SignalClientDelegate> {
                 return WebSocket.connect(url: url,
                                          onMessage: self.onWebSocketMessage,
                                          onDisconnect: { reason in
-                                            self.webSocket = nil
                                             self.cleanUp(reason: reason)
                                          })
             }.then(on: queue) { (webSocket: WebSocket) -> Void in
@@ -145,8 +148,11 @@ internal class SignalClient: MulticastDelegate<SignalClientDelegate> {
 
         _state.mutate { $0.connectionState = .disconnected(reason: reason) }
 
+        pingIntervalTimer = nil
+        pingTimeoutTimer = nil
+
         if let socket = webSocket {
-            socket.cleanUp(reason: reason)
+            socket.cleanUp(reason: reason, notify: false)
             socket.onMessage = nil
             socket.onDisconnect = nil
             self.webSocket = nil
@@ -284,6 +290,7 @@ private extension SignalClient {
         case .join(let joinResponse):
             responseQueueState = .suspended
             latestJoinResponse = joinResponse
+            restartPingTimer()
             notify { $0.signalClient(self, didReceive: joinResponse) }
             _state.mutate { $0.joinResponseCompleter.set(value: joinResponse) }
 
@@ -343,7 +350,7 @@ private extension SignalClient {
         case .refreshToken(let token):
             notify { $0.signalClient(self, didUpdate: token) }
         case .pong(let r):
-            log("pong: \(r)")
+            onReceivedPong(r)
         }
     }
 }
@@ -584,6 +591,7 @@ internal extension SignalClient {
         return sendRequest(r)
     }
 
+    @discardableResult
     func sendLeave() -> Promise<Void> {
         log()
 
@@ -594,6 +602,7 @@ internal extension SignalClient {
         return sendRequest(r)
     }
 
+    @discardableResult
     func sendSimulate(scenario: SimulateScenario) -> Promise<Void> {
         log()
 
@@ -608,6 +617,70 @@ internal extension SignalClient {
 
         return sendRequest(r)
     }
+
+    @discardableResult
+    private func sendPing() -> Promise<Void> {
+        log("ping/pong: sending...")
+
+        let r = Livekit_SignalRequest.with {
+            $0.ping = Int64(Date().timeIntervalSince1970)
+        }
+
+        return sendRequest(r)
+    }
+}
+
+// MARK: - Server ping/pong logic
+
+private extension SignalClient {
+
+    func onPingIntervalTimer() {
+
+        guard let jr = latestJoinResponse else { return }
+
+        pingIntervalTimer?.suspend()
+        sendPing()
+
+        pingTimeoutTimer = {
+            let timer = DispatchQueueTimer(timeInterval: TimeInterval(jr.pingTimeout), queue: queue)
+            timer.handler = { [weak self] in
+                guard let self = self else { return }
+                self.log("ping/pong timed out", .error)
+                self.cleanUp(reason: .networkError(SignalClientError.serverPingTimedOut()))
+            }
+            timer.resume()
+            return timer
+        }()
+    }
+
+    func onReceivedPong(_ r: Int64) {
+
+        log("ping/pong received pong from server")
+        // clear timeout timer
+        pingTimeoutTimer = nil
+        // carry on...
+        pingIntervalTimer?.resume()
+    }
+
+    func restartPingTimer() {
+        // always suspend first
+        pingIntervalTimer = nil
+        pingTimeoutTimer = nil
+        // check received joinResponse already
+        guard let jr = latestJoinResponse,
+              // check server supports ping/pong
+              jr.pingTimeout > 0,
+              jr.pingInterval > 0 else { return }
+
+        log("ping/pong starting with interval: \(jr.pingInterval)")
+
+        pingIntervalTimer = {
+            let timer = DispatchQueueTimer(timeInterval: TimeInterval(jr.pingInterval), queue: queue)
+            timer.handler = { [weak self] in self?.onPingIntervalTimer() }
+            timer.resume()
+            return timer
+        }()
+    }
 }
 
 internal extension Livekit_SignalRequest {

Sources/LiveKit/Errors.swift

@@ -102,6 +102,7 @@ public enum SignalClientError: LiveKitError {
     case close(message: String? = nil)
     case connect(message: String? = nil)
     case timedOut(message: String? = nil)
+    case serverPingTimedOut(message: String? = nil)
 
     public var description: String {
         switch self {
@@ -110,6 +111,7 @@ public enum SignalClientError: LiveKitError {
         case .close(let message): return buildDescription("close", message)
         case .connect(let message): return buildDescription("connect", message)
         case .timedOut(let message): return buildDescription("timedOut", message)
+        case .serverPingTimedOut(let message): return buildDescription("serverPingTimedOut", message)
         }
     }
 }

Sources/LiveKit/Support/WebSocket.swift

@@ -81,7 +81,7 @@ internal class WebSocket: NSObject, URLSessionWebSocketDelegate, Loggable {
         return connectPromise!
     }
 
-    internal func cleanUp(reason: DisconnectReason?) {
+    internal func cleanUp(reason: DisconnectReason?, notify: Bool = true) {
 
         log("reason: \(String(describing: reason))")
 
@@ -102,7 +102,9 @@ internal class WebSocket: NSObject, URLSessionWebSocketDelegate, Loggable {
             connectPromise = nil
         }
 
-        onDisconnect?(reason)
+        if notify {
+            onDisconnect?(reason)
+        }
     }
 
     public func send(data: Data) -> Promise<Void> {