Fix app auto-reconnect on background

Also moves the handling logic from WSClient to ChatClient.

Basically, when we disconnected, we automatically set a timer for reconnection, and reconnected while in background.

Sources/StreamChat/ChatClient.swift

@@ -196,6 +196,8 @@ public class _ChatClient<ExtraData: ExtraDataTypes> {
 
     private(set) lazy var internetConnection = environment.internetConnection()
     private(set) lazy var clientUpdater = environment.clientUpdaterBuilder(self)
+    /// Used for starting and ending background tasks. Hides platform specific logic.
+    private lazy var backgroundTaskScheduler = environment.backgroundTaskSchedulerBuilder()
 
     /// The environment object containing all dependencies of this `Client` instance.
     private let environment: Environment
@@ -299,6 +301,11 @@ public class _ChatClient<ExtraData: ExtraDataTypes> {
         currentUserId = fetchCurrentUserIdFromDatabase()
 
         clientUpdater.reloadUserIfNeeded(completion: completion)
+
+        backgroundTaskScheduler?.startListeningForAppStateUpdates(
+            onEnteringBackground: { [weak self] in self?.handleAppDidEnterBackground() },
+            onEnteringForeground: { [weak self] in self?.handleAppDidBecomeActive() }
+        )
     }
 
     deinit {
@@ -342,6 +349,42 @@ public class _ChatClient<ExtraData: ExtraDataTypes> {
             waiters.removeAll()
         }
     }
+
+    private func handleAppDidEnterBackground() {
+        guard connectionStatus == .connected else { return }
+        guard config.staysConnectedInBackground else {
+            // We immediately disconnect
+            clientUpdater.disconnect(source: .systemInitiated)
+            return
+        }
+        guard let scheduler = backgroundTaskScheduler else { return }
+
+        let succeed = scheduler.beginTask { [weak self] in
+            self?.clientUpdater.disconnect(source: .systemInitiated)
+            // We need to call `endBackgroundTask` else our app will be killed
+            self?.cancelBackgroundTaskIfNeeded()
+        }
+
+        if !succeed {
+            // Can't initiate a background task, close the connection
+            clientUpdater.disconnect(source: .systemInitiated)
+        }
+    }
+
+    private func handleAppDidBecomeActive() {
+        cancelBackgroundTaskIfNeeded()
+
+        guard config.shouldConnectAutomatically else { return }
+        guard connectionStatus != .connected && connectionStatus != .connected else {
+            // We are connected or connecting anyway
+            return
+        }
+        clientUpdater.connect()
+    }
+
+    private func cancelBackgroundTaskIfNeeded() {
+        backgroundTaskScheduler?.endTask()
+    }
 }
 
 extension _ChatClient {
@@ -396,6 +439,20 @@ extension _ChatClient {
         var internetConnection: () -> InternetConnection = { InternetConnection() }
 
         var clientUpdaterBuilder = ChatClientUpdater<ExtraData>.init
+
+        var backgroundTaskSchedulerBuilder: () -> BackgroundTaskScheduler? = {
+            if Bundle.main.isAppExtension {
+                // No background task scheduler exists for app extensions.
+                return nil
+            } else {
+                #if os(iOS)
+                return IOSBackgroundTaskScheduler()
+                #else
+                // No need for background schedulers on macOS, app continues running when inactive.
+                return nil
+                #endif
+            }
+        }
     }
 }
 

Sources/StreamChat/Config/ChatClientConfig.swift

@@ -66,6 +66,19 @@ public struct ChatClientConfig {
     /// Is `true` by default.
     public var shouldConnectAutomatically = true
 
+    /// If set to `true`, the `ChatClient` will try to stay connected while app is backgrounded.
+    /// If set to `false`, websocket disconnects immediately when app is backgrounded.
+    ///
+    /// This flag aims to reduce unnecessary reconnections while quick app switches,
+    /// like when a user just checks a notification or another app.
+    /// `ChatClient` tries to stay connected while in background up to 5 minutes.
+    /// Usually, disconnection occurs around 2-3 minutes.
+    /// If you're using manual connection flow (`shouldConnectAutomatically` set to `false`),
+    /// you should handle connection when opening app from background.
+    ///
+    /// Default value is `true`
+    public var staysConnectedInBackground = true
+
     /// Creates a new instance of `ChatClientConfig`.
     ///
     /// - Parameter apiKey: The API key of the chat app the `ChatClient` connects to.

Sources/StreamChat/WebSocketClient/Engine/URLSessionWebSocketEngine.swift

@@ -92,6 +92,11 @@ class URLSessionWebSocketEngine: NSObject, WebSocketEngine, URLSessionDataDelega
     }
 
     func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
+        // If we received this callback because we closed the WS connection
+        // intentionally, `error` param will be `nil`.
+        // Delegate is already informed with `didCloseWith` callback,
+        // so we don't need to call delegate again.
+        guard let error = error else { return }
         delegate?.webSocketDidDisconnect(error: WebSocketEngineError(error: error))
     }
 }

Sources/StreamChat/WebSocketClient/WebSocketClient.swift

@@ -5,13 +5,6 @@
 import Foundation
 
 class WebSocketClient {
-    /// Additional options for configuring web socket behavior.
-    struct Options: OptionSet {
-        let rawValue: Int
-        /// When the app enters background, `WebSocketClient` starts a long term background task and stays connected.
-        static let staysConnectedInBackground = Options(rawValue: 1 << 0)
-    }
-
     /// The notification center `WebSocketClient` uses to send notifications about incoming events.
     let eventNotificationCenter: EventNotificationCenter
 
@@ -27,11 +20,6 @@ class WebSocketClient {
 
             pingController.connectionStateDidChange(connectionState)
 
-            if case .disconnected = connectionState {
-                // No reconnection attempts are scheduled
-                cancelBackgroundTaskIfNeeded()
-            }
-
             // Publish Connection event with the new state
             let event = ConnectionStatusUpdated(webSocketConnectionState: connectionState)
             eventNotificationCenter.process(event)
@@ -40,9 +28,6 @@ class WebSocketClient {
 
     weak var connectionStateDelegate: ConnectionStateDelegate?
 
-    /// Web socket connection options
-    var options: Options = [.staysConnectedInBackground]
-
     /// The endpoint used for creating a web socket connection.
     ///
     /// Changing this value doesn't automatically update the existing connection. You need to manually call `disconnect`
@@ -68,9 +53,6 @@ class WebSocketClient {
 
     /// An object describing reconnection behavior after the web socket is disconnected.
     private var reconnectionStrategy: WebSocketClientReconnectionStrategy
-
-    /// Used for starting and ending background tasks. Hides platform specific logic
-    private lazy var backgroundTaskScheduler: BackgroundTaskScheduler? = environment.backgroundTaskScheduler
 
     /// The internet connection observer we use for recovering when the connection was offline for some time.
     private let internetConnection: InternetConnection
@@ -118,11 +100,6 @@ class WebSocketClient {
         self.internetConnection = internetConnection
 
         self.eventNotificationCenter = eventNotificationCenter
-
-        backgroundTaskScheduler?.startListeningForAppStateUpdates(
-            onEnteringBackground: { [weak self] in self?.handleAppDidEnterBackground() },
-            onEnteringForeground: { [weak self] in self?.handleAppDidBecomeActive() }
-        )
     }
 
     /// Connects the web connect.
@@ -163,32 +140,6 @@ class WebSocketClient {
             engine?.disconnect()
         }
     }
-
-    private func handleAppDidEnterBackground() {
-        guard options.contains(.staysConnectedInBackground),
-              connectionState.isActive,
-              let scheduler = backgroundTaskScheduler
-        else { return }
-
-        let succeed = scheduler.beginTask { [weak self] in
-            self?.disconnect(source: .systemInitiated)
-            // We need to call `endBackgroundTask` else our app will be killed
-            self?.cancelBackgroundTaskIfNeeded()
-        }
-
-        if !succeed {
-            // Can't initiate a background task, close the connection
-            disconnect(source: .systemInitiated)
-        }
-    }
-
-    private func handleAppDidBecomeActive() {
-        cancelBackgroundTaskIfNeeded()
-    }
-
-    private func cancelBackgroundTaskIfNeeded() {
-        backgroundTaskScheduler?.endTask()
-    }
 }
 
 protocol ConnectionStateDelegate: AnyObject {
@@ -217,20 +168,6 @@ extension WebSocketClient {
                 return StarscreamWebSocketProvider(request: $0, sessionConfiguration: $1, callbackQueue: $2)
             }
         }
-
-        var backgroundTaskScheduler: BackgroundTaskScheduler? = {
-            if Bundle.main.isAppExtension {
-                // No background task scheduler exists for app extensions.
-                return nil
-            } else {
-                #if os(iOS)
-                return IOSBackgroundTaskScheduler()
-                #else
-                // No need for background schedulers on macOS, app continues running when inactive.
-                return nil
-                #endif
-            }
-        }()
     }
 }
 
@@ -274,7 +211,10 @@ extension WebSocketClient: WebSocketEngineDelegate {
 
     func webSocketDidDisconnect(error engineError: WebSocketEngineError?) {
         // Reconnection shouldn't happen for manually initiated disconnect
+        // or system initated disconnect (app enters background)
         let shouldReconnect = connectionState != .disconnecting(source: .userInitiated)
+            && connectionState != .disconnecting(source: .systemInitiated)
+        print("### \(connectionState) \(engineError)")
 
         let disconnectionError: Error?
         if case let .disconnecting(.serverInitiated(webSocketError)) = connectionState {

Sources/StreamChat/Workers/ChatClientUpdater.swift

@@ -137,7 +137,7 @@ class ChatClientUpdater<ExtraData: ExtraDataTypes> {
 
     /// Disconnects the chat client the controller represents from the chat servers. No further updates from the servers
     /// are received.
-    func disconnect() {
+    func disconnect(source: WebSocketConnectionState.DisconnectionSource = .userInitiated) {
         // Disconnecting is not possible in connectionless mode (duh)
         guard client.config.isClientInActiveMode else {
             log.error(ClientError.ClientIsNotInActiveMode().localizedDescription)
@@ -150,7 +150,7 @@ class ChatClientUpdater<ExtraData: ExtraDataTypes> {
         }
 
         // Disconnect the web socket
-        client.webSocketClient?.disconnect(source: .userInitiated)
+        client.webSocketClient?.disconnect(source: source)
 
         // Reset `connectionId`. This would happen asynchronously by the callback from WebSocketClient anyway, but it's
         // safer to do it here synchronously to immediately stop all API calls.