simple websocket store

simplewebsocketstore.swift

import Foundation
import Vapor

@available(macOS 12.0, *)
actor WebSocketStore {
    private var webSockets: [WebSocket] = []
    private var heartbeats: [WebSocket: TimeInterval] = [:]
    private var channels: [String: [WebSocket]] = [:]

    struct CustomMessage: Content {
        let text: String
        let sender: String
    }

    func addWebSocket(_ ws: WebSocket) {
        webSockets.append(ws)
        // Start the heartbeat timer when a client connects
        startHeartbeatTimer(for: ws)

        // Add WebSocket message handling
        handleWebSocketMessages(ws)
    }

    func removeWebSocket(_ ws: WebSocket) {
        if let index = webSockets.firstIndex(where: { $0 === ws }) {
            webSockets.remove(at: index)
            // Stop the heartbeat timer when a client disconnects
            stopHeartbeatTimer(for: ws)
            // Remove the client from all channels
            removeClientFromChannels(ws)
        }
    }

    func broadcast(message: CustomMessage, toChannel channel: String) {
        guard let channelSubscribers = channels[channel] else {
            return
        }

        for ws in channelSubscribers {
            Task {
                // Use a JSON encoder to encode the custom message struct
                let jsonEncoder = JSONEncoder()
                if let jsonData = try? jsonEncoder.encode(message),
                   let jsonString = String(data: jsonData, encoding: .utf8) {
                    await ws.send(jsonString)
                }
            }
        }
    }

    func joinChannel(_ ws: WebSocket, channel: String) {
        if channels[channel] == nil {
            channels[channel] = []
        }
        channels[channel]?.append(ws)
    }

    func leaveChannel(_ ws: WebSocket, channel: String) {
        channels[channel]?.removeAll { $0 === ws }
    }

    private func handleWebSocketMessages(_ ws: WebSocket) {
        ws.onText { ws, text in
            // Handle text messages received
            print("Received text: \(text)")
        }

        ws.onBinary { ws, byteBuffer in
            // Handle binary messages received (ByteBuffer)
            print("Received binary data")
        }

        ws.onData { ws, data in
            // Handle raw data messages received
            print("Received data")
        }

        ws.onPing { ws in
            // Handle incoming ping messages
            print("Received ping")
        }

        ws.onPong { ws in
            // Handle incoming pong messages
            print("Received pong")
        }

        ws.onClose.whenComplete { _ in
            // Handle WebSocket close
            print("WebSocket connection closed")
            removeWebSocket(ws)
        }

        ws.onOpen.whenComplete { _ in
            // Handle WebSocket open
            print("WebSocket connection opened")
        }
    }

    private func startHeartbeatTimer(for ws: WebSocket) {
        // (Previous implementation remains the same)
    }

    private func stopHeartbeatTimer(for ws: WebSocket) {
        // (Previous implementation remains the same)
    }

    private func removeClientFromChannels(_ ws: WebSocket) {
        for (channel, channelSubscribers) in channels {
            channels[channel] = channelSubscribers.filter { $0 !== ws }
        }
    }
}