Refactor Socket Reconnect (davidstump#145)

* Handling abnormal close

* Added specs for socket reconnect changes

Author: dsrees

Sources/client/Socket.swift

@@ -51,6 +51,7 @@ public class Socket {
 
 
 
+  
   //----------------------------------------------------------------------
   // MARK: - Public Attributes
   //----------------------------------------------------------------------
@@ -83,8 +84,11 @@ public class Socket {
   /// Interval between sending a heartbeat
   public var heartbeatInterval: TimeInterval = Defaults.heartbeatInterval
 
-  /// Internval between socket reconnect attempts
-  public var reconnectAfter: (Int) -> TimeInterval = Defaults.steppedBackOff
+  /// Interval between socket reconnect attempts, in seconds
+  public var reconnectAfter: (Int) -> TimeInterval = Defaults.reconnectSteppedBackOff
+  
+  /// Interval between channel rejoin attempts, in seconds
+  public var rejoinAfter: (Int) -> TimeInterval = Defaults.rejoinSteppedBackOff
 
   /// The optional function to receive logs
   public var logger: ((String) -> Void)?
@@ -133,6 +137,9 @@ public class Socket {
   /// Timer to use when attempting to reconnect
   var reconnectTimer: TimeoutTimer
 
+  /// True if the Socket closed cleaned. False if not (connection timeout, heartbeat, etc)
+  var closeWasClean: Bool = false
+  
   /// Websocket connection to the server
   var connection: WebSocketClient?
 
@@ -224,6 +231,9 @@ public class Socket {
     // Do not attempt to reconnect if the socket is currently connected
     guard !isConnected else { return }
 
+    // Reset the clean close flag when attempting to connect
+    self.closeWasClean = false
+    
     self.connection = self.transport(endPointUrl)
     self.connection?.delegate = self
     self.connection?.disableSSLCertValidation = disableSSLCertValidation
@@ -243,6 +253,10 @@ public class Socket {
   /// - paramter callback: Optional. Called when disconnected
   public func disconnect(code: CloseCode = CloseCode.normal,
                          callback: (() -> Void)? = nil) {
+      // The socket was closed cleanly by the User
+      self.closeWasClean = true
+    
+    // Reset any reconnects and teardown the socket connection
     self.reconnectTimer.reset()
     self.teardown(code: code, callback: callback)
   }
@@ -520,6 +534,9 @@ public class Socket {
   internal func onConnectionOpen() {
     self.logItems("transport", "Connected to \(endPoint)")
 
+    // Reset the closeWasClean flag now that the socket has been connected
+    self.closeWasClean = false
+    
     // Send any messages that were waiting for a connection
     self.flushSendBuffer()
 
@@ -541,13 +558,12 @@ public class Socket {
     self.heartbeatTimer?.invalidate()
     self.heartbeatTimer = nil
 
-    self.stateChangeCallbacks.close.forEach({ $0.call() })
+    // Only attempt to reconnect if the socket did not close normally
+    if (!self.closeWasClean) {
+      self.reconnectTimer.scheduleTimeout()
+    }
 
-    // If there was a non-normal event when the connection closed, attempt
-    // to schedule a reconnect attempt
-    let closeCode = CloseCode.init(rawValue: UInt16(code ?? 0))
-    guard closeCode != CloseCode.normal else { return }
-    self.reconnectTimer.scheduleTimeout()
+    self.stateChangeCallbacks.close.forEach({ $0.call() })
   }
 
   internal func onConnectionError(_ error: Error) {
@@ -585,7 +601,12 @@ public class Socket {
 
   /// Triggers an error event to all of the connected Channels
   internal func triggerChannelError() {
-    self.channels.forEach( { $0.trigger(event: ChannelEvent.error) } )
+    self.channels.forEach { (channel) in
+      // Only trigger a channel error if it is in an "opened" state
+      if !(channel.isErrored || channel.isLeaving || channel.isClosed) {
+        channel.trigger(event: ChannelEvent.error)
+      }
+    }
   }
 
   /// Send all messages that were buffered before the socket opened
@@ -637,6 +658,9 @@ public class Socket {
       self.logItems("transport",
                     "heartbeat timeout. Attempting to re-establish connection")
 
+      // Close the socket, flagging the closure as abnormal
+      self.abnormalClose("heartbeat timeout")
+      
       // Disconnect the socket manually. Do not use `teardown` or
       // `disconnect` as they will nil out the websocket delegate
       self.connection?.disconnect(forceTimeout: nil,
@@ -651,6 +675,19 @@ public class Socket {
               payload: [:],
               ref: self.pendingHeartbeatRef)
   }
+  
+  internal func abnormalClose(_ reason: String) {
+    self.closeWasClean = false
+    
+    /*
+     We use NORMAL here since the client is the one determining to close the
+     connection. However, we keep a flag `closeWasClean` set to false so that
+     the client knows that it should attempt to reconnect.
+     */
+    self.connection?.disconnect(forceTimeout: nil,
+                                closeCode: CloseCode.normal.rawValue)
+    
+  }
 }
 
 

Sources/client/Utilities/Defaults.swift

@@ -28,11 +28,17 @@ public class Defaults {
   /// Default interval to send heartbeats on
   public static let heartbeatInterval: TimeInterval = 30.0
 
-  /// Default reconnect function
-  public static let steppedBackOff: (Int) -> TimeInterval = { tries in
-    return tries > 4 ? 10 : [1, 2, 5, 10][tries - 1]
+  /// Default reconnect algorithm for the socket
+  public static let reconnectSteppedBackOff: (Int) -> TimeInterval = { tries in
+    return tries > 9 ? 5.0 : [0.01, 0.05, 0.1, 0.15, 0.2, 0.25, 0.5, 1.0, 2.0][tries - 1]
   }
 
+  /** Default rejoin algorithm for individual channels */
+  public static let rejoinSteppedBackOff: (Int) -> TimeInterval = { tries in
+    return tries > 3 ? 10 : [1, 2, 5][tries - 1]
+  }
+  
+  
   /// Default encode function, utilizing JSONSerialization.data
   public static let encode: ([String: Any]) -> Data = { json in
     return try! JSONSerialization

Tests/client/SocketSpec.swift

@@ -32,11 +32,17 @@ class SocketSpec: QuickSpec {
         expect(socket.timeout).to(equal(Defaults.timeoutInterval))
         expect(socket.heartbeatInterval).to(equal(Defaults.heartbeatInterval))
         expect(socket.logger).to(beNil())
-        expect(socket.reconnectAfter(1)).to(equal(1))
-        expect(socket.reconnectAfter(2)).to(equal(2))
-        expect(socket.reconnectAfter(3)).to(equal(5))
-        expect(socket.reconnectAfter(4)).to(equal(10))
-        expect(socket.reconnectAfter(5)).to(equal(10))
+        expect(socket.reconnectAfter(1)).to(equal(0.010)) // 10ms
+        expect(socket.reconnectAfter(2)).to(equal(0.050)) // 50ms
+        expect(socket.reconnectAfter(3)).to(equal(0.100)) // 100ms
+        expect(socket.reconnectAfter(4)).to(equal(0.150)) // 150ms
+        expect(socket.reconnectAfter(5)).to(equal(0.200)) // 200ms
+        expect(socket.reconnectAfter(6)).to(equal(0.250)) // 250ms
+        expect(socket.reconnectAfter(7)).to(equal(0.500)) // 500ms
+        expect(socket.reconnectAfter(8)).to(equal(1.000)) // 1_000ms (1s)
+        expect(socket.reconnectAfter(9)).to(equal(2.000)) // 2_000ms (2s)
+        expect(socket.reconnectAfter(10)).to(equal(5.00)) // 5_000ms (5s)
+        expect(socket.reconnectAfter(11)).to(equal(5.00)) // 5_000ms (5s)
       })
 
       it("overrides some defaults", closure: {
@@ -110,6 +116,7 @@ class SocketSpec: QuickSpec {
       })
     }
 
+
     describe("websocketProtocol") {
       it("returns wss when protocol is https", closure: {
         let socket = Socket("https://example.com/")
@@ -268,6 +275,13 @@ class SocketSpec: QuickSpec {
           .to(equal(CloseCode.normal.rawValue))
       })
 
+      it("flags the socket as closed cleanly", closure: {
+        expect(socket.closeWasClean).to(beFalse())
+        
+        socket.disconnect()
+        expect(socket.closeWasClean).to(beTrue())
+      })
+      
       it("calls callback", closure: {
         var callCount = 0
         socket.connect()
@@ -599,6 +613,13 @@ class SocketSpec: QuickSpec {
         expect(socket.pendingHeartbeatRef).to(beNil())
       })
 
+      it("does not schedule heartbeat if skipHeartbeat == true", closure: {
+        socket.skipHeartbeat = true
+        socket.resetHeartbeat()
+        
+        expect(socket.heartbeatTimer).to(beNil())
+      })
+      
       it("creates a timer and sends a heartbeat", closure: {
         mockWebSocket.isConnected = true
         socket.connect()
@@ -647,13 +668,19 @@ class SocketSpec: QuickSpec {
         mockWebSocket = WebSocketClientMock()
         mockTimeoutTimer = TimeoutTimerMock()
         socket = Socket(endPoint: "/socket", transport: mockWebSocketTransport)
-        socket.reconnectAfter = { _ in return 10 }
+//        socket.reconnectAfter = { _ in return 10 }
         socket.reconnectTimer = mockTimeoutTimer
-        socket.skipHeartbeat = true
+//        socket.skipHeartbeat = true
       }
 
-      it("does not schedule reconnectTimer timeout if normal close", closure: {
+      it("schedules reconnectTimer timeout if normal close", closure: {
         socket.onConnectionClosed(code: Int(CloseCode.normal.rawValue))
+        expect(mockTimeoutTimer.scheduleTimeoutCalled).to(beTrue())
+      })
+      
+      
+      it("does not schedule reconnectTimer timeout if normal close after explicit disconnect", closure: {
+        socket.disconnect()
         expect(mockTimeoutTimer.scheduleTimeoutCalled).to(beFalse())
       })
 
@@ -662,6 +689,58 @@ class SocketSpec: QuickSpec {
         expect(mockTimeoutTimer.scheduleTimeoutCalled).to(beTrue())
       })
 
+      it("schedules reconnectTimer timeout if connection cannot be made after a previous clean disconnect", closure: {
+        socket.disconnect()
+        socket.connect()
+        
+        socket.onConnectionClosed(code: 1001)
+        expect(mockTimeoutTimer.scheduleTimeoutCalled).to(beTrue())
+      })
+      
+      it("triggers channel error if joining", closure: {
+        let channel = socket.channel("topic")
+        var errorCalled = false
+        channel.on(ChannelEvent.error, callback: { _ in
+          errorCalled = true
+        })
+        
+        channel.join()
+        expect(channel.state).to(equal(.joining))
+
+        socket.onConnectionClosed(code: 1001)
+        expect(errorCalled).to(beTrue())
+      })
+      
+      it("triggers channel error if joined", closure: {
+        let channel = socket.channel("topic")
+        var errorCalled = false
+        channel.on(ChannelEvent.error, callback: { _ in
+          errorCalled = true
+        })
+        
+        channel.join().trigger("ok", payload: [:])
+        expect(channel.state).to(equal(.joined))
+        
+        socket.onConnectionClosed(code: 1001)
+        expect(errorCalled).to(beTrue())
+      })
+      
+      it("does not trigger channel error after leave", closure: {
+        let channel = socket.channel("topic")
+        var errorCalled = false
+        channel.on(ChannelEvent.error, callback: { _ in
+          errorCalled = true
+        })
+        
+        channel.join().trigger("ok", payload: [:])
+        channel.leave()
+        expect(channel.state).to(equal(.closed))
+        
+        socket.onConnectionClosed(code: 1001)
+        expect(errorCalled).to(beFalse())
+      })
+      
+      
       it("triggers onClose callbacks", closure: {
         var oneCalled = 0
         socket.onClose { oneCalled += 1 }
@@ -675,24 +754,6 @@ class SocketSpec: QuickSpec {
         expect(twoCalled).to(equal(1))
         expect(threeCalled).to(equal(0))
       })
-      
-      it("triggers channel error", closure: {
-        let channel = socket.channel("topic")
-        
-        var errorCalled = false
-        var errorEvent: String?
-        var closeCalled = false
-        channel.on(ChannelEvent.error, callback: { (msg) in
-          errorEvent = msg.event
-          errorCalled = true
-        })
-        channel.on(ChannelEvent.close, callback: { (_) in closeCalled = true })
-        
-        socket.onConnectionClosed(code: 1000)
-        expect(errorCalled).to(beTrue())
-        expect(errorEvent).to(equal(ChannelEvent.error))
-        expect(closeCalled).to(beFalse())
-      })
     }
 
     describe("onConnectionError") {
@@ -724,22 +785,48 @@ class SocketSpec: QuickSpec {
         expect(lastError).toNot(beNil())
       })
 
-      it("triggers channel error", closure: {
+      
+      it("triggers channel error if joining", closure: {
         let channel = socket.channel("topic")
+        var errorCalled = false
+        channel.on(ChannelEvent.error, callback: { _ in
+          errorCalled = true
+        })
+        
+        channel.join()
+        expect(channel.state).to(equal(.joining))
 
+        socket.onConnectionError(TestError.stub)
+        expect(errorCalled).to(beTrue())
+      })
+      
+      it("triggers channel error if joined", closure: {
+        let channel = socket.channel("topic")
         var errorCalled = false
-        var errorEvent: String?
-        var closeCalled = false
-        channel.on(ChannelEvent.error, callback: { (msg) in
-          errorEvent = msg.event
+        channel.on(ChannelEvent.error, callback: { _ in
           errorCalled = true
         })
-        channel.on(ChannelEvent.close, callback: { (_) in closeCalled = true })
+        
+        channel.join().trigger("ok", payload: [:])
+        expect(channel.state).to(equal(.joined))
 
         socket.onConnectionError(TestError.stub)
         expect(errorCalled).to(beTrue())
-        expect(errorEvent).to(equal(ChannelEvent.error))
-        expect(closeCalled).to(beFalse())
+      })
+      
+      it("does not trigger channel error after leave", closure: {
+        let channel = socket.channel("topic")
+        var errorCalled = false
+        channel.on(ChannelEvent.error, callback: { _ in
+          errorCalled = true
+        })
+        
+        channel.join().trigger("ok", payload: [:])
+        channel.leave()
+        expect(channel.state).to(equal(.closed))
+        
+        socket.onConnectionError(TestError.stub)
+        expect(errorCalled).to(beFalse())
       })
     }