TimerWebsocketCleanup.swift

// ChatServer is a very simple chat server

import Foundation
import KituraWebSocket
import Foundation

class ChatService: WebSocketService {
    
    private var connections = [String: ConnectionWrapper]()
    
    func removeConnection(id: String) {
        connections.removeValue(forKey: id)
    }
    
    public func connected(connection: WebSocketConnection) {
        connections[connection.id] = ConnectionWrapper(connection: connection, whenDone: removeConnection)
    }
    
    public func disconnected(connection: WebSocketConnection, reason: WebSocketCloseReasonCode) {
        connections.removeValue(forKey: connection.id)
    }
    
    public func received(message: Data, from: WebSocketConnection) {
        from.close(reason: .invalidDataType, description: "Chat-Server only accepts text messages")
        connections.removeValue(forKey: from.id)
    }
    
    public func received(message: String, from: WebSocketConnection) {
        connections[from.id]?.refreshHeartbeat()
        for (connectionId, connectionWrapper) in connections {
            if connectionId != from.id {
                connectionWrapper.connection.send(message: message)
            }
        }
    }
}

class ConnectionWrapper {
    static let timeInterval: Int = 5
    let connection: WebSocketConnection
    let deregisterConnection: (String) -> Void
    var heartbeat: Date
    var timer: DispatchSourceTimer = DispatchSource.makeTimerSource()
    let dispatchTimeInterval: DispatchTimeInterval = DispatchTimeInterval.seconds((timeInterval))

    init(connection: WebSocketConnection, whenDone: @escaping (String) -> Void) {
        self.connection = connection
        self.deregisterConnection = whenDone
        self.heartbeat = Date()
        self.timerStart()
    }

    func checkAlive() {
        if abs(heartbeat.timeIntervalSinceNow) > Double(ConnectionWrapper.timeInterval) {
            if abs(heartbeat.timeIntervalSinceNow) > Double(ConnectionWrapper.timeInterval * 2) {
                connection.close(reason: .closedAbnormally, description: nil)
                deregisterConnection(self.connection.id)
                timer.cancel()
            }
            connection.ping()
        }
    }

    func refreshHeartbeat() {
        heartbeat = Date()
    }
    
    func timerStart() {
        timer.schedule(deadline: .now(), repeating: dispatchTimeInterval, leeway: DispatchTimeInterval.milliseconds(ConnectionWrapper.timeInterval))
        timer.setEventHandler(handler: self.checkAlive)
        timer.resume()
    }
}