fix: userSession memory leak - WPB-18909

WireDomain/Sources/WireDomain/Components/UserSessionComponent.swift

@@ -126,4 +126,14 @@ public final class UserSessionComponent {
         )
     }
 
+    // MARK: - Teardown
+
+    /// Invalidate and cancel all network URLSessions to stop in-flight requests.
+    /// Call this before closing Core Data to prevent responses from being processed.
+    public func invalidateNetworkServices() {
+        restNetworkService.invalidateAndCancel()
+        websocketNetworkService.invalidateAndCancel()
+        blacklistNetworkService.invalidateAndCancel()
+    }
+
 }

WireDomain/Sources/WireDomain/Synchronization/OneOnOneResolver.swift

@@ -56,11 +56,18 @@ public struct OneOnOneResolver: OneOnOneResolverProtocol {
     public func resolveAllOneOnOneConversations() async throws {
         let usersIDs = try await userLocalStore.fetchAllUserIDsWithOneOnOneConversation()
 
-        await withTaskGroup(of: Void.self) { group in
+        try await withThrowingTaskGroup(of: Void.self) { group in
             for userID in usersIDs {
+                // Check if task is cancelled before adding more work
+                try Task.checkCancellation()
+
                 group.addTask {
                     do {
+                        try Task.checkCancellation()
                         try await resolveOneOnOneConversation(with: userID)
+                    } catch is CancellationError {
+                        // Task was cancelled, exit gracefully
+                        WireLogger.conversation.info("1-1 conversation resolution cancelled for userID \(userID)")
                     } catch {
                         /// skip conversation migration for this user
                         WireLogger.conversation.error(
@@ -70,6 +77,9 @@ public struct OneOnOneResolver: OneOnOneResolverProtocol {
                     }
                 }
             }
+
+            // Wait for all tasks, but this will throw CancellationError if parent is cancelled
+            try await group.waitForAll()
         }
     }
 

WireNetwork/Sources/WireNetwork/Network/NetworkService/NetworkService.swift

@@ -66,6 +66,13 @@ public final class NetworkService: NSObject, NetworkServiceProtocol {
         self.urlSession = urlSession
     }
 
+    /// Explicitly invalidate and cancel all tasks in the URLSession.
+    /// Call this during teardown to ensure in-flight requests are cancelled.
+    public func invalidateAndCancel() {
+        urlSession?.invalidateAndCancel()
+        urlSession = nil
+    }
+
     public func executeRequest(_ request: URLRequest) async throws -> (Data, HTTPURLResponse) {
         guard let urlSession else {
             throw NetworkServiceError.serviceNotConfigured

WireUI/Sources/WireLocators/Locators.swift

@@ -339,6 +339,12 @@ public enum Locators {
 
     }
 
+    public enum DeveloperToolsPage: String {
+
+        case userSessionMemoryLeakLabel
+
+    }
+
     public enum WireDrive {
 
         public enum FilesFilterPage: String {

wire-ios-sync-engine/Source/SessionManager/SessionManager.swift

@@ -859,7 +859,12 @@ public final class SessionManager: NSObject, SessionManagerType {
                 }
             }
 
-            self?.activeUserSession = nil
+            // Notify delegate that teardown is complete - session should deallocate after this
+            self?.delegate?.sessionManagerWillLogout(
+                environment: environment,
+                error: error,
+                userSessionCanBeTornDown: nil
+            )
         }
     }
 

wire-ios-sync-engine/Source/Synchronization/SyncAgent.swift

@@ -139,6 +139,18 @@ final class SyncAgent: NSObject, SyncAgentProtocol {
         }
     }
 
+    func terminate() async {
+        // Nil delegate first to prevent callbacks during/after suspension
+        delegate = nil
+        cancellables.removeAll()
+        await suspend()
+    }
+
+    func restart() async {
+        setupBindings()
+        resume()
+    }
+
     /// Suspend any ongoing sync tasks.
 
     func suspend() async {

wire-ios-sync-engine/Source/UserSession/ZMUserSession/ZMUserSession+Actions.swift

@@ -113,9 +113,9 @@ public extension ZMUserSession {
 
         let conversation = userInfo.conversation(in: managedObjectContext)
 
-        managedObjectContext.perform {
+        managedObjectContext.perform { [weak self] in
             conversation?.mutedMessageTypes = .all
-            self.managedObjectContext.saveOrRollback()
+            self?.managedObjectContext.saveOrRollback()
             BackgroundActivityFactory.shared.endBackgroundActivity(activity)
             completionHandler()
         }
@@ -136,7 +136,12 @@ public extension ZMUserSession {
             guard let self else { return }
 
             messageReplyObserver = nil
-            syncManagedObjectContext.performGroupedBlock {
+            syncManagedObjectContext.performGroupedBlock { [weak self] in
+                guard let self else {
+                    BackgroundActivityFactory.shared.endBackgroundActivity(activity)
+                    completionHandler()
+                    return
+                }
 
                 let conversationOnSyncContext = userInfo.conversation(in: self.syncManagedObjectContext)
                 if result == .failed {

wire-ios-sync-engine/Source/UserSession/ZMUserSession/ZMUserSession+Authentication.swift

@@ -17,6 +17,7 @@
 //
 
 import Foundation
+import WireLogging
 
 extension ZMUserSession {
 
@@ -70,17 +71,17 @@ extension ZMUserSession {
 
     func close(deleteCookie: Bool, completion: @escaping () -> Void) {
         // Clear all notifications associated with the account from the notification center
-        syncManagedObjectContext.performGroupedBlock {
-            self.localNotificationDispatcher?.cancelAllNotifications()
+        syncManagedObjectContext.performGroupedBlock { [weak self] in
+            self?.localNotificationDispatcher?.cancelAllNotifications()
         }
 
         if deleteCookie {
             deleteUserKeychainItems()
         }
 
-        syncManagedObjectContext.perform {
-            self.tearDown()
-        }
+        // Call tearDown directly (not in perform block)
+        // Network services are explicitly invalidated in tearDown before closing Core Data
+        tearDown()
 
         completion()
     }
@@ -105,31 +106,53 @@ extension ZMUserSession {
             return
         }
 
-        let payload: [String: Any] = if let password = credentials.password, !password.isEmpty {
-            ["password": password]
-        } else {
-            [:]
-        }
-
-        let request = ZMTransportRequest(
-            path: "/clients/\(selfClientIdentifier)",
-            method: .delete,
-            payload: payload as ZMTransportData,
-            apiVersion: apiVersion.rawValue
-        )
 
-        request.add(ZMCompletionHandler(on: managedObjectContext, block: { [weak self] response in
+        // Step 1 & 2: Stop all sync and work processing before sending DELETE request
+        Task { [weak self] in
             guard let self else { return }
 
-            if response.httpStatus == 200 {
-                delegate?.userDidLogout(accountId: accountID)
-                completion(.success(()))
+            await self.syncAgent?.terminate()
+
+            // Nil out syncAgent to prevent resume() calls
+            self.syncAgent = nil
+
+            if let workAgent = self.clientSessionComponent?.workAgent {
+                await workAgent.clearSchedulerQueue()
+                await workAgent.stop()
+            }
+
+            // Step 3: Send DELETE request
+            let payload: [String: Any] = if let password = credentials.password, !password.isEmpty {
+                ["password": password]
             } else {
-                completion(.failure(errorFromFailedDeleteResponse(response)))
+                [:]
             }
-        }))
 
-        transportSession.enqueueOneTime(request)
+            let request = ZMTransportRequest(
+                path: "/clients/\(selfClientIdentifier)",
+                method: .delete,
+                payload: payload as ZMTransportData,
+                apiVersion: apiVersion.rawValue
+            )
+
+            request.add(ZMCompletionHandler(on: self.managedObjectContext, block: { [weak self] response in
+                guard let self else { return }
+
+                // Step 4: Stop operationLoop and transportSession (both success and failure)
+                self.operationLoop?.tearDown()
+                self.operationLoop = nil
+                self.transportSession.tearDown()
+
+                if response.httpStatus == 200 {
+                    self.delegate?.userDidLogout(accountId: accountID)
+                    completion(.success(()))
+                } else {
+                    completion(.failure(self.errorFromFailedDeleteResponse(response)))
+                }
+            }))
+
+            self.transportSession.enqueueOneTime(request)
+        }
     }
 
     func errorFromFailedDeleteResponse(_ response: ZMTransportResponse!) -> NSError {
@@ -141,7 +164,7 @@ extension ZMUserSession {
                 .clientDeletedRemotely
             case "invalid-credentials",
                  "missing-auth",
-                 "bad-request": // in case the password not matching password format requirement
+                 "bad-request": // in case the password does not match password format requirement
                 .invalidCredentials
             default:
                 .unknownError

wire-ios-sync-engine/Source/UserSession/ZMUserSession/ZMUserSession+CertificateRevocationLists.swift

@@ -57,8 +57,11 @@ extension ZMUserSession {
             return
         }
 
-        Task {
-            let isE2EIFeatureEnabled = await managedObjectContext.perform { self.e2eiFeature.isEnabled }
+        Task { [weak self] in
+            guard let self else { return }
+            let isE2EIFeatureEnabled = await managedObjectContext.perform { [weak self] in
+                self?.e2eiFeature.isEnabled ?? false
+            }
             if isE2EIFeatureEnabled {
                 await cRLsChecker.checkExpiredCRLs()
             }

wire-ios-sync-engine/Source/UserSession/ZMUserSession/ZMUserSession+Clients.swift

@@ -42,7 +42,8 @@ public extension ZMUserSession {
 
     @objc
     func fetchAllClients() {
-        syncManagedObjectContext.performGroupedBlock {
+        syncManagedObjectContext.performGroupedBlock { [weak self] in
+            guard let self else { return }
             self.applicationStatusDirectory.clientUpdateStatus.needsToFetchClients(andVerifySelfClient: false)
             RequestAvailableNotification.notifyNewRequestsAvailable(self)
         }
@@ -55,7 +56,8 @@ public extension ZMUserSession {
         client.markForDeletion()
         client.managedObjectContext?.saveOrRollback()
 
-        syncManagedObjectContext.performGroupedBlock {
+        syncManagedObjectContext.performGroupedBlock { [weak self] in
+            guard let self else { return }
             self.applicationStatusDirectory.clientUpdateStatus.deleteClients(withCredentials: credentials)
             RequestAvailableNotification.notifyNewRequestsAvailable(self)
         }

wire-ios-sync-engine/Source/UserSession/ZMUserSession/ZMUserSession+LegalHold.swift

@@ -58,7 +58,13 @@ public extension ZMUserSession {
             return completionHandler(.invalidSelfUser)
         }
 
-        Task {
+        Task { [weak self] in
+            guard let self else {
+                return await MainActor.run {
+                    completionHandler(.invalidState)
+                }
+            }
+
             let selfUser = await syncManagedObjectContext.perform {
                 ZMUser.selfUser(in: self.syncManagedObjectContext)
             }

wire-ios-sync-engine/Source/UserSession/ZMUserSession/ZMUserSession+LifeCycle.swift

@@ -73,7 +73,8 @@ public extension ZMUserSession {
         // to ensure that they have timers running or are deleted/obfuscated if
         // needed. Note: ZMMessageTimer will only create a new timer for a message
         // if one does not already exist.
-        syncManagedObjectContext.performGroupedBlock {
+        syncManagedObjectContext.performGroupedBlock { [weak self] in
+            guard let self else { return }
             ZMMessage.deleteOldEphemeralMessages(self.syncManagedObjectContext)
         }
     }
@@ -87,7 +88,8 @@ public extension ZMUserSession {
         for changes in storedNotifications {
             NSManagedObjectContext.mergeChanges(fromRemoteContextSave: changes, into: [managedObjectContext])
 
-            syncManagedObjectContext.performGroupedBlock {
+            syncManagedObjectContext.performGroupedBlock { [weak self] in
+                guard let self else { return }
                 NSManagedObjectContext.mergeChanges(
                     fromRemoteContextSave: changes,
                     into: [self.syncManagedObjectContext]
@@ -97,16 +99,17 @@ public extension ZMUserSession {
 
         // we only process pending changes on sync context bc changes on the
         // ui context will be processed when we do the save.
-        syncManagedObjectContext.performGroupedBlock {
-            self.syncManagedObjectContext.processPendingChanges()
+        syncManagedObjectContext.performGroupedBlock { [weak self] in
+            self?.syncManagedObjectContext.processPendingChanges()
         }
 
         managedObjectContext.saveOrRollback()
     }
 
     internal func recalculateUnreadMessages() {
         WireLogger.badgeCount.info("recalculate unread conversations")
-        syncManagedObjectContext.performGroupedBlock {
+        syncManagedObjectContext.performGroupedBlock { [weak self] in
+            guard let self else { return }
             ZMConversation.recalculateUnreadMessages(in: self.syncManagedObjectContext)
             self.syncManagedObjectContext.saveOrRollback()
             NotificationInContext(name: .calculateBadgeCount, context: self.notificationContext).post()

wire-ios-sync-engine/Source/UserSession/ZMUserSession/ZMUserSession+Logs.swift

@@ -67,8 +67,8 @@ public extension ZMUserSession {
             handler(metadata, type, error, userInfo)
         }
 
-        syncManagedObjectContext.performGroupedBlock {
-            self.syncManagedObjectContext.errorOnSaveCallback = { context, error in
+        syncManagedObjectContext.performGroupedBlock { [weak self] in
+            self?.syncManagedObjectContext.errorOnSaveCallback = { context, error in
                 let type = context.type
 
                 guard

wire-ios-sync-engine/Source/UserSession/ZMUserSession/ZMUserSession+OneOnOne.swift

@@ -37,7 +37,12 @@ public extension ZMUserSession {
             return
         }
 
-        Task {
+        Task { [weak self] in
+            guard let self else {
+                completion(.failure(.userDoesNotExist))
+                return
+            }
+
             do {
                 let syncContext = self.syncContext
 

wire-ios-sync-engine/Source/UserSession/ZMUserSession/ZMUserSession+Push.swift

@@ -85,7 +85,8 @@ public extension ZMUserSession {
     }
 
     internal func registerCurrentPushToken() {
-        managedObjectContext.performGroupedBlock {
+        managedObjectContext.performGroupedBlock { [weak self] in
+            guard let self else { return }
             self.sessionManager?.configurePushToken(session: self)
         }
     }
@@ -166,8 +167,8 @@ extension ZMUserSession {
     ) {
         // foreground notification responder exists on the UI context, so we
         // need to switch to that context
-        managedObjectContext.perform {
-            let responder = self.sessionManager?.foregroundNotificationResponder
+        managedObjectContext.perform { [weak self] in
+            let responder = self?.sessionManager?.foregroundNotificationResponder
             let shouldPresent = responder?.shouldPresentNotification(with: userInfo) ?? true
 
             var options = UNNotificationPresentationOptions()