Use DispatchSourceTimer for delayed background tasks

Author: tagavari

AirMessage/Connection/ClientConnection.swift

@@ -60,49 +60,59 @@ class ClientConnection {
 		case pingExpiry
 	}
 
-	struct RunningTimer {
-		let timer: Timer
-		let callback: (ClientConnection) -> Void
-	}
-	
-	private var timerDict: [TimerType: RunningTimer] = [:]
+	private var timerDict: [TimerType: DispatchSourceTimer] = [:]
 
 	/**
 	 Cancels all pending expiry timers
 	 */
 	func cancelAllTimers() {
 		for (_, runningTimer) in timerDict {
-			runningTimer.timer.invalidate()
+			runningTimer.cancel()
 		}
 		timerDict.removeAll()
 	}
 
 	/**
 	 Cancels the expiry timer of the specified type
 	 */
-	func cancelTimer(ofType type: TimerType) {
-		timerDict[type]?.timer.invalidate()
+	func cancelTimer(ofType type: TimerType, queue: DispatchQueue) {
+		//Make sure we're on the provided dispatch queue
+		assertDispatchQueue(queue)
+		
+		timerDict[type]?.cancel()
 		timerDict[type] = nil
 	}
 
 	/**
 	 Schedules a timer of the specified type after interval to run the callback
 	 If a timer was previously scheduled of this type, that timer is cancelled and replaced with this one
 	 */
-	func startTimer(ofType type: TimerType, interval: TimeInterval, callback: @escaping (ClientConnection) -> Void) {
+	func startTimer(ofType type: TimerType, interval: TimeInterval, queue: DispatchQueue, callback: @escaping (ClientConnection) -> Void) {
+		//Make sure we're on the provided dispatch queue
+		assertDispatchQueue(queue)
+		
 		//Cancel existing timers
-		cancelTimer(ofType: type)
+		cancelTimer(ofType: type, queue: queue)
 
 		//Create and start the timer
-		let timer = Timer(timeInterval: interval, target: self, selector: #selector(onTimerExpire), userInfo: type, repeats: false)
-		RunLoop.main.add(timer, forMode: .common)
-		timerDict[type] = RunningTimer(timer: timer, callback: callback)
-	}
-	
-	@objc private func onTimerExpire(timer: Timer) {
-		//Invoke the callback of the specified type
-		let type = timer.userInfo as! TimerType
-		timerDict[type]?.callback(self)
-		timerDict[type] = nil
+		let timer = DispatchSource.makeTimerSource(queue: queue)
+		timer.schedule(deadline: .now() + interval, repeating: .never)
+		timer.setEventHandler { [weak self] in
+			//Make sure we're still on the same queue
+			assertDispatchQueue(queue)
+			
+			//Check our reference to self
+			guard let self = self else {
+				return
+			}
+			
+			//Invoke the callback
+			callback(self)
+			
+			//Remove this timer
+			self.timerDict[type] = nil
+		}
+		timer.resume()
+		timerDict[type] = timer
 	}
 }

AirMessage/Connection/CommConst.swift

@@ -13,7 +13,7 @@ class CommConst {
 	//Timeouts
 	static let handshakeTimeout: TimeInterval = 10 //10 seconds
 	static let pingTimeout: TimeInterval = 30 //30 seconds
-	static let keepAliveMillis: TimeInterval = 30 * 60 //30 minutes
+	static let keepAliveInterval: TimeInterval = 30 * 60 //30 minutes
 
 	static let maxPacketAllocation = 50 * 1024 * 1024 //50 MB
 

AirMessage/Connection/Connect/DataProxyConnect.swift

@@ -23,10 +23,11 @@ class DataProxyConnect: DataProxy {
 	private var idToken: String?
 
 	private var webSocket: WebSocket?
-	private var handshakeTimer: Timer?
-	private var pingTimer: Timer?
+	private var handshakeTimer: DispatchSourceTimer?
+	private var pingTimer: DispatchSourceTimer?
+	private static let pingFrequency: TimeInterval = 5 * 60
 
-	private var connectionRecoveryTimer: Timer?
+	private var connectionRecoveryTimer: DispatchSourceTimer?
 	private var connectionRecoveryCount = 0
 	//The max num of attempts before capping the delay time - not before giving up
 	private static let connectionRecoveryCountMax = 8
@@ -40,12 +41,24 @@ class DataProxyConnect: DataProxy {
 		self.idToken = idToken
 	}
 
-	@objc func startServer() {
+	deinit {
+		//Ensure the server proxy isn't running when we go out of scope
+		assert(!isActive, "DataProxyConnect was deinitialized while active")
+		
+		//Ensure timers are cleaned up
+		assert(pingTimer == nil, "DataProxyConnect was deinitialized with an active ping timer")
+		assert(handshakeTimer == nil, "DataProxyConnect was deinitialized with an active handshake timer")
+		assert(connectionRecoveryTimer == nil, "DataProxyConnect was deinitialized with an active connection recovery timer")
+	}
+	
+	func startServer() {
 		//Ignore if we're already connecting or connected
 		guard !isActive else { return }
 
 		//Cancel any active connection recover timers (in case the user initiated a reconnect)
-		stopConnectionRecoveryTimer()
+		processingQueue.sync {
+			stopConnectionRecoveryTimer()
+		}
 
 		//Build the URL
 		var queryItems = [
@@ -76,29 +89,24 @@ class DataProxyConnect: DataProxy {
 	}
 
 	func stopServer() {
+		//Ignore if we're not running
+		guard isActive else { return }
+		
 		processingQueue.sync {
-			//Ignore if we're not connected
-			guard isActive, let socket = webSocket else { return }
-			
-			//Cancel the handshake timer
-			stopHandshakeTimer()
+			let socket = webSocket!
 
 			//Clear connected clients
 			connectionsLock.withWriteLock {
 				connectionsMap.removeAll()
 			}
 			NotificationNames.postUpdateConnectionCount(0)
 
-			//Disconect
+			//Disconnect
 			socket.disconnect()
-			isActive = false
 			delegate?.dataProxy(self, didStopWithState: .stopped, isRecoverable: false)
 		}
-	}
-	
-	deinit {
-		//Make sure we stop the server when we go out of scope
-		stopServer()
+		
+		isActive = false
 	}
 
 	func send(message data: Data, to client: ClientConnection?, encrypt: Bool, onSent: (() -> ())?) {
@@ -193,23 +201,28 @@ class DataProxyConnect: DataProxy {
 		assertDispatchQueue(processingQueue)
 
 		//Cancel the old timer
-		pingTimer?.invalidate()
+		pingTimer?.cancel()
 
 		//Create the new timer
-		let timer = Timer(timeInterval: 5 * 60, target: self, selector: #selector(onPingTimer), userInfo: nil, repeats: true)
-		RunLoop.main.add(timer, forMode: .common)
+		let timer = DispatchSource.makeTimerSource(queue: processingQueue)
+		timer.schedule(deadline: .now() + DataProxyConnect.pingFrequency, repeating: DataProxyConnect.pingFrequency)
+		timer.setEventHandler(handler: onPingTimer)
+		timer.resume()
 		pingTimer = timer
 	}
 
 	private func stopPingTimer() {
 		//Make sure we're on the processing queue
 		assertDispatchQueue(processingQueue)
 
-		pingTimer?.invalidate()
+		pingTimer?.cancel()
 		pingTimer = nil
 	}
 
-	@objc private func onPingTimer() {
+	private func onPingTimer() {
+		//Make sure we're on the processing queue
+		assertDispatchQueue(processingQueue)
+		
 		//Ping
 		webSocket?.write(ping: Data())
 	}
@@ -221,40 +234,54 @@ class DataProxyConnect: DataProxy {
 		assertDispatchQueue(processingQueue)
 
 		//Cancel the old timer
-		handshakeTimer?.invalidate()
+		handshakeTimer?.cancel()
 
 		//Create the new timer
-		let timer = Timer(timeInterval: ConnectConstants.handshakeTimeout, target: self, selector: #selector(onHandshakeTimer), userInfo: nil, repeats: false)
-		RunLoop.main.add(timer, forMode: .common)
+		let timer = DispatchSource.makeTimerSource(queue: processingQueue)
+		timer.schedule(deadline: .now() + ConnectConstants.handshakeTimeout, repeating: .never)
+		timer.setEventHandler(handler: onHandshakeTimer)
+		timer.resume()
 		handshakeTimer = timer
 	}
 
 	private func stopHandshakeTimer() {
 		//Make sure we're on the processing queue
 		assertDispatchQueue(processingQueue)
 
-		handshakeTimer?.invalidate()
+		handshakeTimer?.cancel()
 		handshakeTimer = nil
 	}
 
-	@objc private func onHandshakeTimer() {
+	private func onHandshakeTimer() {
+		//Make sure we're on the processing queue
+		assertDispatchQueue(processingQueue)
+		
 		//Disconnect
 		webSocket?.disconnect()
+		
+		//Clean up
+		stopHandshakeTimer()
 	}
 
 	//MARK: Connection recovery timer
 
 	private func startConnectionRecoveryTimer() {
+		//Make sure we're on the processing queue
+		assertDispatchQueue(processingQueue)
+		
 		//Cancel the old timer
-		connectionRecoveryTimer?.invalidate()
+		connectionRecoveryTimer?.cancel()
 
 		//Wait an exponentially increasing wait period + a random delay
 		let randomOffset: TimeInterval = Double(arc4random()) / Double(UInt32.max)
 		let delay: TimeInterval = pow(2, Double(connectionRecoveryCount)) + randomOffset
 
 		//Create the new timer
-		let timer = Timer(timeInterval: delay, target: self, selector: #selector(startServer), userInfo: nil, repeats: false)
-		RunLoop.main.add(timer, forMode: .common)
+		//startServer() must be invoked from the main thread
+		let timer = DispatchSource.makeTimerSource(queue: DispatchQueue.main)
+		timer.schedule(deadline: .now() + delay, repeating: .never)
+		timer.setEventHandler(handler: startServer)
+		timer.resume()
 		connectionRecoveryTimer = timer
 
 		//Add to the attempt counter
@@ -264,7 +291,10 @@ class DataProxyConnect: DataProxy {
 	}
 
 	private func stopConnectionRecoveryTimer() {
-		connectionRecoveryTimer?.invalidate()
+		//Make sure we're on the processing queue
+		assertDispatchQueue(processingQueue)
+		
+		connectionRecoveryTimer?.cancel()
 		connectionRecoveryTimer = nil
 	}
 

AirMessage/Connection/ConnectionManager.swift

@@ -12,7 +12,7 @@ class ConnectionManager {
 	private let timerQueue = DispatchQueue(label: Bundle.main.bundleIdentifier! + ".connection.timer", qos: .utility)
 
 	private var dataProxy: DataProxy?
-	private var keepaliveTimer: Timer?
+	private var keepaliveTimer: DispatchSourceTimer?
 	private let fileDownloadRequestMapLock = ReadWriteLock()
 	private var fileDownloadRequestMap: [Int16: FileDownloadRequest] = [:]
 
@@ -77,19 +77,20 @@ class ConnectionManager {
 
 	//MARK: - Timers
 
-	@objc private func runKeepalive() {
+	private func runKeepalive() {
 		guard let dataProxy = dataProxy else { return }
 
+		//Make sure we're running on the timer queue
+		assertDispatchQueue(timerQueue)
+		
 		//Send a ping to all clients
 		ConnectionManager.sendMessageHeaderOnly(dataProxy, to: nil, ofType: NHT.ping, encrypt: false)
 
 		//Start ping response timers
 		dataProxy.connectionsLock.withReadLock {
 			for connection in dataProxy.connections {
-				timerQueue.sync {
-					connection.startTimer(ofType: .pingExpiry, interval: CommConst.pingTimeout) { [weak self] client in
-						self?.dataProxy?.disconnect(client: client)
-					}
+				connection.startTimer(ofType: .pingExpiry, interval: CommConst.pingTimeout, queue: timerQueue) { [weak self] client in
+					self?.dataProxy?.disconnect(client: client)
 				}
 			}
 		}
@@ -288,7 +289,7 @@ class ConnectionManager {
 	private func handleMessageAuthentication(dataProxy: DataProxy, packer messagePacker: inout AirPacker, from client: C) throws {
 		//Cancel the handshake expiry timer
 		timerQueue.sync {
-			client.cancelTimer(ofType: .handshakeExpiry)
+			client.cancelTimer(ofType: .handshakeExpiry, queue: timerQueue)
 		}
 
 		//Sends an authorization rejected message and closes the connection
@@ -1544,9 +1545,13 @@ extension ConnectionManager: DataProxyDelegate {
 
 		//Start the keepalive timer
 		if dataProxy.requiresPersistence {
-			let timer = Timer(timeInterval: CommConst.keepAliveMillis, target: self, selector: #selector(runKeepalive), userInfo: nil, repeats: true)
-			RunLoop.main.add(timer, forMode: .common)
-			keepaliveTimer = timer
+			timerQueue.sync {
+				let timer = DispatchSource.makeTimerSource(queue: timerQueue)
+				timer.schedule(deadline: .now() + CommConst.keepAliveInterval, repeating: CommConst.keepAliveInterval)
+				timer.setEventHandler(handler: runKeepalive)
+				timer.resume()
+				keepaliveTimer = timer
+			}
 		}
 
 		LogManager.log("Server started", level: .info)
@@ -1558,8 +1563,10 @@ extension ConnectionManager: DataProxyDelegate {
 		NotificationNames.postUpdateConnectionCount(0)
 
 		//Stop the keepalive timer
-		keepaliveTimer?.invalidate()
-		keepaliveTimer = nil
+		timerQueue.sync {
+			keepaliveTimer?.cancel()
+			keepaliveTimer = nil
+		}
 
 		if isRecoverable {
 			LogManager.log("Server paused", level: .info)
@@ -1612,7 +1619,7 @@ extension ConnectionManager: DataProxyDelegate {
 
 		//Start the expiry timer
 		timerQueue.sync {
-			client.startTimer(ofType: .handshakeExpiry, interval: CommConst.handshakeTimeout) { [weak self] client in
+			client.startTimer(ofType: .handshakeExpiry, interval: CommConst.handshakeTimeout, queue: timerQueue) { [weak self] client in
 				LogManager.log("Handshake response for client \(client.readableID) timed out, disconnecting", level: .debug)
 				self?.dataProxy?.disconnect(client: client)
 			}

AirMessage/Connection/Direct/DataProxyTCP.swift

@@ -190,8 +190,8 @@ class DataProxyTCP: DataProxy {
 	}
 
 	deinit {
-		//Make sure we stop the server when we go out of scope
-		stopServer()
+		//Ensure the server proxy isn't running when we go out of scope
+		assert(!serverRunning, "DataProxyTCP was deinitialized while active")
 	}
 
 	func send(message data: Data, to client: ClientConnection?, encrypt: Bool, onSent: (() -> ())?) {