fix: prevent main app from blocking notification extension (take 4) - WPB-23511

WireDomain/Sources/WireDomain/Components/ClientSessionComponent.swift

@@ -69,6 +69,7 @@ public final class ClientSessionComponent {
     private let completionHandlers: CompletionHandlers
 
     private let faultyMLSRemovalKeysByDomain: [String: [String]]
+    private let backgroundTaskExecuter: any BackgroundTaskExecuter
 
     public init(
         selfUserID: UUID,
@@ -88,7 +89,8 @@ public final class ClientSessionComponent {
         proteusService: any ProteusServiceInterface,
         coreCryptoProvider: any CoreCryptoProviderProtocol,
         completionHandlers: CompletionHandlers,
-        faultyMLSRemovalKeysByDomain: [String: [String]]
+        faultyMLSRemovalKeysByDomain: [String: [String]],
+        backgroundTaskExecuter: any BackgroundTaskExecuter
     ) {
         self.selfUserID = selfUserID
         self.selfClientID = selfClientID
@@ -108,6 +110,7 @@ public final class ClientSessionComponent {
         self.coreCryptoProvider = coreCryptoProvider
         self.completionHandlers = completionHandlers
         self.faultyMLSRemovalKeysByDomain = faultyMLSRemovalKeysByDomain
+        self.backgroundTaskExecuter = backgroundTaskExecuter
     }
 
     public private(set) lazy var authenticationManager = AuthenticationManager(
@@ -411,7 +414,8 @@ public final class ClientSessionComponent {
         liveBrokenGroupSubject: liveBrokenGroupSubject,
         journal: journal,
         mlsGroupRepairAgent: mlsGroupRepairAgent,
-        earService: earService
+        earService: earService,
+        backgroundTaskExecuter: backgroundTaskExecuter
     )
 
     public lazy var incrementalSyncV2: IncrementalSyncV2 = if let sharedContainerURL {
@@ -430,6 +434,7 @@ public final class ClientSessionComponent {
             journal: journal,
             mlsGroupRepairAgent: mlsGroupRepairAgent,
             earService: earService,
+            backgroundTaskExecuter: backgroundTaskExecuter,
             createPushChannelState: { [selfClientID] in
                 PushChannelState(sharedContainerURL: sharedContainerURL, clientID: selfClientID)
             }

WireDomain/Sources/WireDomain/Components/UserSessionComponent.swift

@@ -45,6 +45,7 @@ public final class UserSessionComponent {
     private let coreCryptoProvider: any CoreCryptoProviderProtocol
 
     private let faultyMLSRemovalKeysByDomain: [String: [String]]
+    private let backgroundTaskExecuter: any BackgroundTaskExecuter
 
     public init(
         currentBuildNumber: String,
@@ -64,7 +65,8 @@ public final class UserSessionComponent {
         mlsDecryptionService: any MLSDecryptionServiceInterface,
         proteusService: any ProteusServiceInterface,
         coreCryptoProvider: any CoreCryptoProviderProtocol,
-        faultyMLSRemovalKeysByDomain: [String: [String]]
+        faultyMLSRemovalKeysByDomain: [String: [String]],
+        backgroundTaskExecuter: any BackgroundTaskExecuter
     ) {
         self.currentBuildNumber = currentBuildNumber
         self.selfUserID = selfUserID
@@ -84,6 +86,7 @@ public final class UserSessionComponent {
         self.coreCryptoProvider = coreCryptoProvider
         self.sharedContainerURL = sharedContainerURL
         self.faultyMLSRemovalKeysByDomain = faultyMLSRemovalKeysByDomain
+        self.backgroundTaskExecuter = backgroundTaskExecuter
     }
 
     private let cookieStorage: any WireNetwork.CookieStorageProtocol
@@ -112,7 +115,8 @@ public final class UserSessionComponent {
             proteusService: proteusService,
             coreCryptoProvider: coreCryptoProvider,
             completionHandlers: completionHandlers,
-            faultyMLSRemovalKeysByDomain: faultyMLSRemovalKeysByDomain
+            faultyMLSRemovalKeysByDomain: faultyMLSRemovalKeysByDomain,
+            backgroundTaskExecuter: backgroundTaskExecuter
         )
     }
 

WireDomain/Sources/WireDomain/Notifications/Components/NSEClientScope.swift

@@ -48,6 +48,8 @@ final class NSEClientScope: Component<NSEClientScopeDependency> {
 
     }
 
+    private let eventID: UUID
+    private let contentHandler: (UNNotificationContent) -> Void
     private let clientID: String
     private let restNetworkService: NetworkService
     private let webSocketNetworkService: NetworkService
@@ -56,12 +58,13 @@ final class NSEClientScope: Component<NSEClientScopeDependency> {
     private let isFederationEnabled: Bool
     private let coreDataStack: CoreDataStack
     private let earService: EARServiceInterface
+    private let backgroundTaskExecuter: any BackgroundTaskExecuter
 
     private let pushChannelCoordinator: AppExtensionPushChannelCoordinator
-    private var currentTask: Task<Void, any Error>?
-    private var monitoringTask: Task<Void, any Error>?
 
     init(
+        eventID: UUID,
+        contentHandler: @escaping (UNNotificationContent) -> Void,
         parent: any Scope,
         clientID: String,
         restNetworkService: NetworkService,
@@ -70,8 +73,11 @@ final class NSEClientScope: Component<NSEClientScopeDependency> {
         localDomain: String,
         isFederationEnabled: Bool,
         coreDataStack: CoreDataStack,
-        earService: EARServiceInterface
+        earService: EARServiceInterface,
+        backgroundTaskExecuter: any BackgroundTaskExecuter
     ) {
+        self.eventID = eventID
+        self.contentHandler = contentHandler
         self.clientID = clientID
         self.restNetworkService = restNetworkService
         self.webSocketNetworkService = webSocketNetworkService
@@ -81,14 +87,12 @@ final class NSEClientScope: Component<NSEClientScopeDependency> {
         self.coreDataStack = coreDataStack
         self.pushChannelCoordinator = AppExtensionPushChannelCoordinator(clientID: clientID)
         self.earService = earService
+        self.backgroundTaskExecuter = backgroundTaskExecuter
 
         super.init(parent: parent)
     }
 
-    func processPayload(
-        eventID: UUID,
-        contentHandler: @escaping (UNNotificationContent) -> Void
-    ) async throws {
+    func processPayload() async throws {
         // Pull pending update events.
         let eventStream: AsyncStream<[UpdateEvent]>
         let publicKeys = try earService.fetchPublicKeys()
@@ -97,56 +101,57 @@ final class NSEClientScope: Component<NSEClientScopeDependency> {
             let (useCase, stream) = syncEventsUseCase()
             eventStream = stream
 
-            // because we might be interrupted when in background, we wrap the sync in an expiringActivity that will
-            // cancel the task (not keeping any file lock in suspend mode)
-            try await withExpiringActivity(reason: "processPayload in NSE") { [weak self] in
-                guard let self else { return }
-                // make sure no pushChannel is open
-                let pushChannelState = PushChannelState(
-                    sharedContainerURL: dependency.appContainerURL,
-                    clientID: clientID
-                )
+            // make sure no pushChannel is open
+            let pushChannelState = PushChannelState(
+                sharedContainerURL: dependency.appContainerURL,
+                clientID: clientID
+            )
+            do {
+                try await pushChannelState.markAsOpen()
+            } catch {
+                throw Failure.pushChannelAlreadyOpened
+            }
+
+            let currentTask = Task<Void, any Error> {
                 do {
-                    try await pushChannelState.markAsOpen()
+                    try Task.checkCancellation()
+                    try await useCase.invoke()
                 } catch {
-                    throw Failure.pushChannelAlreadyOpened
-                }
-
-                monitoringTask = Task { [weak self] in
-                    var request = await self?.pushChannelCoordinator.listenForYieldRequests()
-                    if Task.isCancelled {
-                        return
-                    }
-                    WireLogger.sync.debug("requested to cancel sync", attributes: .incrementalSync, .newNSE)
-                    self?.currentTask?.cancel()
-                    request?.acknowledge()
-                    WireLogger.sync.debug("notified main App to resume sync", attributes: .incrementalSync, .newNSE)
+                    // either we timeout during decrypting/storing events OR an issue
+                    // with the sync. In both cases, we end up with a stream of
+                    // notifications that has not been shown, so we need to continue
+                    // to show them.
+                    WireLogger.sync.warn(
+                        "syncing events via websocket: \(String(describing: error))",
+                        attributes: .incrementalSyncV3, .newNSE
+                    )
+                    await pushChannelState.markAsClosed()
                 }
+            }
 
-                currentTask = Task {
-                    do {
-                        try Task.checkCancellation()
-                        try await useCase.invoke()
-                    } catch {
-                        // either we timeout during decrypting/storing events OR an issue
-                        // with the sync. In both cases, we end up with a stream of
-                        // notifications that has not been shown, so we need to continue
-                        // to show them.
-                        WireLogger.sync.warn(
-                            "syncing events via websocket: \(String(describing: error))",
-                            attributes: .incrementalSyncV3, .newNSE
-                        )
-                        await pushChannelState.markAsClosed()
-                    }
+            let monitoringTask = Task<Void, any Error> { [pushChannelCoordinator] in
+                let request = await pushChannelCoordinator.listenForYieldRequests()
+                if Task.isCancelled {
+                    return
                 }
-                try await currentTask?.value
-                WireLogger.sync.debug("closing push channel")
-                await pushChannelState.markAsClosed()
+                WireLogger.sync.debug("requested to cancel sync", attributes: .incrementalSync, .newNSE)
+                currentTask.cancel()
+                request.acknowledge()
+                WireLogger.sync.debug("notified main App to resume sync", attributes: .incrementalSync, .newNSE)
+            }
 
-                // no need to monitor anymore let's cancel
-                monitoringTask?.cancel()
+            try await withTaskCancellationHandler {
+                try await currentTask.value
+            } onCancel: {
+                currentTask.cancel()
             }
 
+            WireLogger.sync.debug("closing push channel")
+            await pushChannelState.markAsClosed()
+
+            // no need to monitor anymore let's cancel
+            monitoringTask.cancel()
+
         } else {
             eventStream = try await pullEventsUseCase.invoke(publicKeys: publicKeys)
         }
@@ -260,7 +265,7 @@ final class NSEClientScope: Component<NSEClientScopeDependency> {
                 coreCryptoKeyMigrationManager: coreCryptoMigrationManager,
                 allowCreation: false,
                 localDomain: localDomain,
-                backgroundTaskManager: nil
+                backgroundTaskExecuter: backgroundTaskExecuter
             )
         }
     }

WireDomain/Sources/WireDomain/Notifications/Components/NSEFlow.swift

@@ -46,6 +46,7 @@ final class NSEFlow: BootstrapComponent {
     // TODO: [WPB-19778] use to check app version migration
     private let currentAppVersion: String
     private var scopesByAccount = [Account: NSEUserScope]()
+    private let backgroundTaskExecuter = PassthroughTaskExecuter()
 
     init(
         currentAppVersion: String,
@@ -100,7 +101,8 @@ final class NSEFlow: BootstrapComponent {
         shared {
             NSEUserScope(
                 parent: self,
-                account: account
+                account: account,
+                backgroundTaskExecuter: backgroundTaskExecuter
             )
         }
     }

WireDomain/Sources/WireDomain/Notifications/Components/NSEUserScope.swift

@@ -61,6 +61,8 @@ final class NSEUserScope: Component<NSEUserScopeDependency> {
         account.userIdentifier
     }
 
+    let backgroundTaskExecuter: any BackgroundTaskExecuter
+
     public var userAccountDataURL: URL {
         dependency.accountDataURL.appending(path: accountID.uuidString)
     }
@@ -90,9 +92,11 @@ final class NSEUserScope: Component<NSEUserScopeDependency> {
 
     init(
         parent: any Scope,
-        account: Account
+        account: Account,
+        backgroundTaskExecuter: any BackgroundTaskExecuter
     ) {
         self.account = account
+        self.backgroundTaskExecuter = backgroundTaskExecuter
         super.init(parent: parent)
     }
 
@@ -167,6 +171,8 @@ final class NSEUserScope: Component<NSEUserScopeDependency> {
 
         // Continue with client.
         let clientScope = clientScope(
+            eventID: eventID,
+            contentHandler: contentHandler,
             clientID: clientID,
             restNetworkService: networkServices.rest,
             webSocketNetworkService: networkServices.webSocket,
@@ -177,10 +183,7 @@ final class NSEUserScope: Component<NSEUserScopeDependency> {
             earService: earService
         )
 
-        try await clientScope.processPayload(
-            eventID: eventID,
-            contentHandler: contentHandler
-        )
+        try await clientScope.processPayload()
     }
 
     private func fetchBackendEnvironment() throws -> BackendEnvironment2? {
@@ -309,6 +312,8 @@ final class NSEUserScope: Component<NSEUserScopeDependency> {
     // MARK: - Children
 
     private func clientScope(
+        eventID: UUID,
+        contentHandler: @escaping (UNNotificationContent) -> Void,
         clientID: String,
         restNetworkService: NetworkService,
         webSocketNetworkService: NetworkService,
@@ -319,6 +324,8 @@ final class NSEUserScope: Component<NSEUserScopeDependency> {
         earService: EARServiceInterface
     ) -> NSEClientScope {
         NSEClientScope(
+            eventID: eventID,
+            contentHandler: contentHandler,
             parent: self,
             clientID: clientID,
             restNetworkService: restNetworkService,
@@ -327,7 +334,8 @@ final class NSEUserScope: Component<NSEUserScopeDependency> {
             localDomain: localDomain,
             isFederationEnabled: isFederationEnabled,
             coreDataStack: coreDataStack,
-            earService: earService
+            earService: earService,
+            backgroundTaskExecuter: backgroundTaskExecuter
         )
     }
 

WireDomain/Sources/WireDomain/Notifications/NotificationServiceExtension.swift

@@ -31,7 +31,7 @@ import WireUtilitiesPackage
 ///
 /// These sequential steps represents the NSE dependency graph (using Needle).
 
-public final class NotificationServiceExtension: NotificationServiceProtocol {
+public final class NotificationServiceExtension {
 
     // MARK: - Properties
 
@@ -129,9 +129,13 @@ public final class NotificationServiceExtension: NotificationServiceProtocol {
         }
     }
 
-    public func serviceExtensionTimeWillExpire() {
-        logger.warn("new notification service will expire", attributes: .newNSE, .safePublic)
+    public var hasOnGoingTask: Bool {
+        onGoingTask != nil
+    }
+
+    public func cancel() async {
         onGoingTask?.cancel()
+        await onGoingTask?.value
     }
 }
 
@@ -223,7 +227,7 @@ extension NotificationServiceExtension {
 
     private func logUserError(_ error: NSEUserScope.Failure) {
         switch error {
-        case let .mainAppRequired(message):
+        case let .mainAppRequired(message, _):
             logger.warn(
                 "Main app required, need to open main app: \(message)",
                 attributes: .newNSE, .safePublic