feat: Resume after disconnect

Sources/DiscordKit/GatewayClient.swift

@@ -4,35 +4,52 @@ import FoundationNetworking
 #endif
 
 actor GatewayClient {
-    private let ws: URLSessionWebSocketTask
+    private var ws: URLSessionWebSocketTask
     private let token: String
     private(set) var open = false
-    var sequenceNum: Int? = nil
+    private var sequenceNum: Int? = nil
+    private var sessionID: String? = nil
+    private var resumeURL: URL? = nil
+    private let intents: Intents
+    private let gatewayURL: URL
+    private var hbTask: Task<(), any Error>? = nil
 
-    init(gatewayURL: URL, token: String) {
+
+    init(gatewayURL: URL, token: String, intents: Intents) {
+        self.intents = intents
+        self.token = token
+        self.gatewayURL = gatewayURL
         let queryItems = [URLQueryItem(name: "v", value: "10"), URLQueryItem(name: "encoding", value: "json")]
         ws = URLSession.shared.webSocketTask(with: gatewayURL.appending(queryItems: queryItems))
-        self.token = token
     }
 
-    func openConnection(intents: Intents) async throws {
+    func openConnection() async throws {
         ws.resume()
-        open = true
-        guard case .hello(let helloMessage) = try await getMessage().d else { throw GatewayError.mismatchedOpcode }
-        dump(helloMessage)
-
-        let heartbeatTask = Task() {
-            try await Task.sleep(for: .milliseconds(Int.random(in: 0...helloMessage.heartbeat_interval)))
-            try await sendHeartbeat()
-
-            while !Task.isCancelled {
-                try await Task.sleep(for: .milliseconds(helloMessage.heartbeat_interval))
-                try await sendHeartbeat()
-            }
-        }
+        try await setupHeartbeat()
 
         try await sendIdentify(intents: intents)
-        _ = await heartbeatTask.result
+        guard case .ready(let readyMessage) = try await getMessage().d else { throw GatewayError.connectionFailure }
+        open = true
+        sessionID = readyMessage.session_id
+        resumeURL = readyMessage.resume_gateway_url
+    }
+
+    func setupHeartbeat() async throws {
+        guard case .hello(let helloMessage) = try await getMessage().d else { throw GatewayError.connectionFailure }
+
+        self.hbTask = Task.detached { [self] in
+            do {
+                try await Task.sleep(for: .milliseconds(Int.random(in: 0...helloMessage.heartbeat_interval)))
+                try await sendHeartbeat()
+
+                while !Task.isCancelled {
+                    try await Task.sleep(for: .milliseconds(helloMessage.heartbeat_interval))
+                    try await sendHeartbeat()
+                }
+            } catch {
+                print("Heartbeat task canceled")
+            }
+        }
     }
 
     func sendIdentify(intents: Intents) async throws {
@@ -58,11 +75,15 @@ actor GatewayClient {
         var gwMessage: GatewayMessage? = nil
         let json = JSONDecoder()
         while gwMessage == nil {
-            let wsMessage = try await ws.receive()
-            guard case .string(let str) = wsMessage else { throw GatewayError.invalidMessage }
-            strBuffer.append(str)
             do {
+                let wsMessage = try await ws.receive()
+                //print(wsMessage)
+                guard case .string(let str) = wsMessage else { throw GatewayError.invalidMessage }
+                strBuffer.append(str)
                 gwMessage = try json.decode(GatewayMessage.self, from: Data(strBuffer.utf8))
+            } catch URLError.networkConnectionLost {
+                self.open = false
+                try await reconnect()
             } catch DecodingError.dataCorrupted {
                 continue
             }
@@ -72,26 +93,75 @@ actor GatewayClient {
         return gwMessage
     }
 
+    private func reconnect() async throws {
+        hbTask?.cancel()
+        do {
+            try await attemptResume()
+        } catch {
+            print(error)
+            ws.cancel()
+            let queryItems = [URLQueryItem(name: "v", value: "10"), URLQueryItem(name: "encoding", value: "json")]
+            ws = URLSession.shared.webSocketTask(with: gatewayURL.appending(queryItems: queryItems))
+            try await openConnection()
+        }
+    }
+
+    private func attemptResume() async throws {
+        guard
+            ws.closeCode.rawValue != 4004 && ws.closeCode.rawValue < 4010,
+            let resumeURL = resumeURL,
+            let sessionID = sessionID,
+            let sequenceNum = sequenceNum
+        else {
+            if let closeReason = ws.closeReason {
+                throw GatewayError.disconected(code: ws.closeCode.rawValue, message: String(data: closeReason, encoding: .utf8))
+            } else {
+                throw GatewayError.disconected(code: ws.closeCode.rawValue, message: "unknown")
+            }
+        }
+        ws.cancel()
+        let queryItems = [URLQueryItem(name: "v", value: "10"), URLQueryItem(name: "encoding", value: "json")]
+        ws = URLSession.shared.webSocketTask(with: resumeURL.appending(queryItems: queryItems))
+        ws.resume()
+        try await setupHeartbeat()
+
+        let payload = """
+        {
+            "op": 6,
+            "d": {
+                "token": "\(token)",
+                "session_id": "\(sessionID)",
+                "seq": \(sequenceNum)
+            }
+        }
+        """ // Im lazy
+        try await ws.send(.string(payload))
+        self.open = true
+    }
+
     private func sendHeartbeat() async throws {
+        guard self.open else { print("conn closed, skipping hb"); return }
         let hbMessage = "{\"op\":1,\"d\":\(sequenceNum == nil ? "null" : String(sequenceNum!))}"
         try await ws.send(.string(hbMessage))
     }
 
-    var events: AsyncStream<GatewayMessage> {
+    var events: AsyncStream<GatewayPayload> {
         AsyncStream { [self] in
-            var event: GatewayMessage? = nil
-            while event == nil {
+            var payload: GatewayPayload? = nil
+            while payload == nil {
                 do {
-                    while await !open {
+                    while await !self.open {
                         try await Task.sleep(for: .seconds(1))
                     }
-                    event = try await getMessage()
-                    if event!.op == 1 { try await sendHeartbeat() }
+                    let event = try await getMessage()
+                    dump(event)
+                    if event.op == 1 { try await sendHeartbeat() }
+                    payload = event.d
                 } catch {
                     print("Error listening to gateway: \(error)")
                 }
             }
-            return event!
+            return payload!
         }
     }
 
@@ -101,4 +171,6 @@ public enum GatewayError: Error {
     case invalidMessage
     case invalidOpcode
     case mismatchedOpcode
+    case disconected(code: Int, message: String?)
+    case connectionFailure
 }