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

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

@@ -259,8 +259,7 @@ final class NSEClientScope: Component<NSEClientScopeDependency> {
                 syncContext: coreDataStack.syncContext,
                 coreCryptoKeyMigrationManager: coreCryptoMigrationManager,
                 allowCreation: false,
-                localDomain: localDomain,
-                backgroundTaskManager: nil
+                localDomain: localDomain
             )
         }
     }

WireDomain/Tests/WireDomainTests/Synchronization/IncrementalSyncV2Tests.swift

@@ -62,10 +62,7 @@ final class IncrementalSyncV2Tests: XCTestCase {
         coreCrypto = MockCoreCryptoProtocol()
         coreCrypto.mockTransaction(context: coreCryptoContext)
         coreCryptoProvider = MockCoreCryptoProviderProtocol()
-        coreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(
-            backgroundTaskManager: nil,
-            coreCrypto: coreCrypto
-        )
+        coreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(coreCrypto: coreCrypto)
         journal = Journal(
             userID: UUID(),
             storage: UserDefaults.temporary()

WireDomain/Tests/WireDomainTests/Synchronization/PullPendingUpdateEventsSyncV2Tests.swift

@@ -67,10 +67,7 @@ class PullPendingUpdateEventsSyncV2Tests: XCTestCase {
         )
 
         // Setup mocks
-        coreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(
-            backgroundTaskManager: nil,
-            coreCrypto: coreCrypto
-        )
+        coreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(coreCrypto: coreCrypto)
         decryptor.decryptEventsInContext_MockMethod = { envelope, _ in
             EventDecryptorResult(events: envelope.events, brokenMLSGroupIDs: [Scaffolding.mlsGroupID])
         }

wire-ios-data-model/Source/Core Crypto/CoreCryptoProvider.swift

@@ -75,8 +75,6 @@ public actor CoreCryptoProvider: CoreCryptoProviderProtocol {
     private var epochObserver: WireCoreCryptoUniffi.EpochObserver?
     private let localDomain: String?
 
-    private let backgroundTaskManager: (any BackgroundTaskManager)?
-
     public init(
         selfUserID: UUID,
         sharedContainerURL: URL,
@@ -85,8 +83,7 @@ public actor CoreCryptoProvider: CoreCryptoProviderProtocol {
         syncContext: NSManagedObjectContext,
         coreCryptoKeyMigrationManager: CoreCryptoKeyMigrationManagerProtocol,
         allowCreation: Bool = true,
-        localDomain: String?,
-        backgroundTaskManager: (any BackgroundTaskManager)?
+        localDomain: String?
     ) {
         self.selfUserID = selfUserID
         self.sharedContainerURL = sharedContainerURL
@@ -97,16 +94,12 @@ public actor CoreCryptoProvider: CoreCryptoProviderProtocol {
         self.coreCryptoKeyMigrationManager = coreCryptoKeyMigrationManager
         self.featureRespository = LegacyFeatureRepository(context: syncContext)
         self.localDomain = localDomain
-        self.backgroundTaskManager = backgroundTaskManager
     }
 
     public func coreCrypto() async throws -> SafeCoreCrypto {
         let coreCrypto = try await getCoreCrypto()
         try await registerMlsTransportIfNecessary(with: coreCrypto)
-        return SafeCoreCrypto(
-            backgroundTaskManager: backgroundTaskManager,
-            coreCrypto: coreCrypto
-        )
+        return SafeCoreCrypto(coreCrypto: coreCrypto)
     }
 
     public func initialiseMLSWithBasicCredentials(mlsClientID: MLSClientID) async throws {

wire-ios-data-model/Source/Core Crypto/SafeCoreCrypto.swift

@@ -23,37 +23,19 @@ import WireLogging
 /// within a background task (if a task manager is provided).
 public final class SafeCoreCrypto {
 
-    private let backgroundTaskManager: (any BackgroundTaskManager)?
     let coreCrypto: any CoreCryptoProtocol
 
-    public init(
-        backgroundTaskManager: (any BackgroundTaskManager)?,
-        coreCrypto: any CoreCryptoProtocol
-    ) {
-        self.backgroundTaskManager = backgroundTaskManager
+    public init(coreCrypto: any CoreCryptoProtocol) {
         self.coreCrypto = coreCrypto
     }
 
     public func registerEpochObserver(_ epochObserver: any EpochObserver) async throws {
         try await coreCrypto.registerEpochObserver(epochObserver)
     }
 
-    public func transaction<Result>(
-        block: @escaping (any CoreCryptoContextProtocol) async throws -> Result
-    ) async throws -> Result {
-        if let backgroundTaskManager {
-            try await transaction(
-                backgroundTaskManager: backgroundTaskManager,
-                block: block
-            )
-        } else {
-            try await coreCrypto.transaction(block)
-        }
-    }
-
-    /// Perform a transaction within the context of a background task to
-    /// ensure the transaction has additional time to complete while the
-    /// app transitions to a suspended state.
+    /// Perform a transaction within an expiring activity to ensure the
+    /// transaction has additional time to complete while the app transitions
+    /// to a suspended state.
     ///
     /// This is particularly important when running in the main app because
     /// if the app is suspended during a transaction, then core crypto will
@@ -62,30 +44,13 @@ public final class SafeCoreCrypto {
     /// own transactions. This can lead to the Notification Service Extension
     /// being blocked and not process any notifications until the main app
     /// is resumed and the transaction is completed and the file lock released.
-    private func transaction<Result>(
-        backgroundTaskManager: any BackgroundTaskManager,
-        block: @escaping (any CoreCryptoContextProtocol) async throws -> Result
+    ///
+    public func transaction<Result>(
+        block: @escaping @Sendable (any CoreCryptoContextProtocol) async throws -> Result
     ) async throws -> Result {
-        let transactionTask = Task {
+        try await withExpiringActivity(reason: "core crypto transaction") { [coreCrypto] in
             try await coreCrypto.transaction(block)
         }
-
-        let taskID = backgroundTaskManager.beginBackgroundTask(
-            withName: "core crypto transaction",
-            expirationHandler: {
-                WireLogger.coreCrypto.warn(
-                    "Background task expiring, cancelling core crypto transaction",
-                    attributes: .safePublic
-                )
-                transactionTask.cancel()
-            }
-        )
-
-        defer {
-            backgroundTaskManager.endBackgroundTask(taskID)
-        }
-
-        return try await transactionTask.value
     }
 
 }

wire-ios-data-model/Support/Sources/CoreCryptoMocksEnvelope.swift

@@ -37,10 +37,7 @@ public class CoreCryptoMocksEnvelope {
         self.coreCrypto = MockCoreCryptoProtocol()
         self.coreCryptoProvider = MockCoreCryptoProviderProtocol()
 
-        coreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(
-            backgroundTaskManager: nil,
-            coreCrypto: coreCrypto
-        )
+        coreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(coreCrypto: coreCrypto)
         coreCrypto.mockTransaction(context: coreCryptoContext)
     }
 

wire-ios-data-model/Tests/E2EIdentity/E2EIServiceTests.swift

@@ -38,10 +38,7 @@ final class E2EIServiceTests: ZMConversationTestsBase {
         mockCoreCrypto = MockCoreCryptoProtocol()
         mockCoreCrypto.mockTransaction(context: mockCoreCryptoContext)
         mockCoreCryptoProvider = MockCoreCryptoProviderProtocol()
-        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(
-            backgroundTaskManager: nil,
-            coreCrypto: mockCoreCrypto
-        )
+        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(coreCrypto: mockCoreCrypto)
         sut = E2EIService(
             e2eIdentity: mockE2eIdentity,
             coreCryptoProvider: mockCoreCryptoProvider,

wire-ios-data-model/Tests/E2EIdentity/E2EIVerificationStatusServiceTests.swift

@@ -35,10 +35,7 @@ final class E2EIVerificationStatusServiceTests: XCTestCase {
         mockCoreCrypto = MockCoreCryptoProtocol()
         mockCoreCrypto.mockTransaction(context: mockCoreCryptoContext)
         mockCoreCryptoProvider = MockCoreCryptoProviderProtocol()
-        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(
-            backgroundTaskManager: nil,
-            coreCrypto: mockCoreCrypto
-        )
+        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(coreCrypto: mockCoreCrypto)
         sut = E2EIVerificationStatusService(coreCryptoProvider: mockCoreCryptoProvider)
     }
 

wire-ios-data-model/Tests/MLS/MLSActionExecutorTests.swift

@@ -41,10 +41,7 @@ class MLSActionExecutorTests: ZMBaseManagedObjectTest {
         mockCoreCrypto = MockCoreCryptoProtocol()
         mockCoreCrypto.mockTransaction(context: mockCoreCryptoContext)
         mockCoreCryptoProvider = MockCoreCryptoProviderProtocol()
-        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(
-            backgroundTaskManager: nil,
-            coreCrypto: mockCoreCrypto
-        )
+        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(coreCrypto: mockCoreCrypto)
         mockLegacyFeatureRepository = MockLegacyFeatureRepositoryInterface()
 
         sut = MLSActionExecutor(

wire-ios-data-model/Tests/MLS/MLSEncryptionServiceTests.swift

@@ -37,10 +37,7 @@ final class MLSEncryptionServiceTests: XCTestCase {
         mockCoreCrypto = MockCoreCryptoProtocol()
         mockCoreCrypto.mockTransaction(context: mockCoreCryptoContext)
         mockCoreCryptoProvider = MockCoreCryptoProviderProtocol()
-        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(
-            backgroundTaskManager: nil,
-            coreCrypto: mockCoreCrypto
-        )
+        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(coreCrypto: mockCoreCrypto)
         sut = MLSEncryptionService(coreCryptoProvider: mockCoreCryptoProvider)
     }
 

wire-ios-data-model/Tests/MLS/MLSServiceTests.swift

@@ -57,10 +57,7 @@ final class MLSServiceTests: ZMConversationTestsBase, MLSServiceDelegate {
         mockCoreCrypto = MockCoreCryptoProtocol()
         mockCoreCrypto.mockTransaction(context: mockCoreCryptoContext)
         mockCoreCryptoProvider = MockCoreCryptoProviderProtocol()
-        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(
-            backgroundTaskManager: nil,
-            coreCrypto: mockCoreCrypto
-        )
+        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(coreCrypto: mockCoreCrypto)
         mockEncryptionService = MockMLSEncryptionServiceInterface()
         mockDecryptionService = MockMLSDecryptionServiceInterface()
         mockMLSActionExecutor = MockMLSActionExecutor()

wire-ios-data-model/Tests/Proteus/ProteusServiceTests.swift

@@ -41,10 +41,7 @@ class ProteusServiceTests: XCTestCase {
         mockCoreCrypto = MockCoreCryptoProtocol()
         mockCoreCrypto.mockTransaction(context: mockCoreCryptoContext)
         mockCoreCryptoProvider = MockCoreCryptoProviderProtocol()
-        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(
-            backgroundTaskManager: nil,
-            coreCrypto: mockCoreCrypto
-        )
+        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(coreCrypto: mockCoreCrypto)
         sut = ProteusService(coreCryptoProvider: mockCoreCryptoProvider)
     }
 

wire-ios-data-model/Tests/UseCases/IsUserE2EICertifiedUseCaseTests.swift

@@ -50,10 +50,7 @@ final class IsUserE2EICertifiedUseCaseTests: ZMBaseManagedObjectTest {
         mockCoreCrypto = MockCoreCryptoProtocol()
         mockCoreCrypto.mockTransaction(context: mockCoreCryptoContext)
         mockCoreCryptoProvider = MockCoreCryptoProviderProtocol()
-        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(
-            backgroundTaskManager: nil,
-            coreCrypto: mockCoreCrypto
-        )
+        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(coreCrypto: mockCoreCrypto)
         mockLegacyFeatureRepository = .init()
         mockLegacyFeatureRepository.fetchE2EI_MockValue = .init(status: .enabled, config: .init())
         sut = .init(

wire-ios-request-strategy/Sources/E2EIdentity/E2EIKeyPackageRotatorTests.swift

@@ -36,10 +36,7 @@ class E2EIKeyPackageRotatorTests: MessagingTestBase {
 
         mockCoreCrypto = MockCoreCryptoProtocol()
         mockCoreCryptoProvider = MockCoreCryptoProviderProtocol()
-        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(
-            backgroundTaskManager: nil,
-            coreCrypto: mockCoreCrypto
-        )
+        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(coreCrypto: mockCoreCrypto)
         mockLegacyFeatureRepository = .init()
 
         sut = E2EIKeyPackageRotator(

wire-ios-request-strategy/Tests/Helpers/MessagingTestBase.swift

@@ -569,8 +569,7 @@ extension MessagingTestBase {
             syncContext: syncMOC,
             coreCryptoKeyMigrationManager: mockKeyMigrationManager,
             allowCreation: true,
-            localDomain: owningDomain,
-            backgroundTaskManager: nil
+            localDomain: owningDomain
         )
 
         // Initialize CoreCrypto (this calls proteusInit internally)

wire-ios-request-strategy/Tests/Helpers/ProteusClientSimulator.swift

@@ -83,8 +83,7 @@ final class ProteusClientSimulator {
             syncContext: syncMOC,
             coreCryptoKeyMigrationManager: mockKeyMigrationManager,
             allowCreation: true,
-            localDomain: owningDomain,
-            backgroundTaskManager: nil
+            localDomain: owningDomain
         )
 
         // Initialize CoreCrypto

wire-ios-share-engine/Sources/SharingSession.swift

@@ -236,8 +236,7 @@ public final class SharingSession {
             syncContext: coreDataStack.syncContext,
             coreCryptoKeyMigrationManager: CoreCryptoKeyMigrationManager(journal: journal),
             allowCreation: false,
-            localDomain: localDomain,
-            backgroundTaskManager: nil
+            localDomain: localDomain
         )
         let featureRepository = LegacyFeatureRepository(context: coreDataStack.syncContext)
         let mlsActionExecutor = MLSActionExecutor(

wire-ios-share-engine/Sources/SharingSessionLoader.swift

@@ -335,8 +335,7 @@ public struct SharingSessionLoader {
             syncContext: coreDataStack.syncContext,
             coreCryptoKeyMigrationManager: CoreCryptoKeyMigrationManager(journal: journal),
             allowCreation: false,
-            localDomain: backendMetadata.domain,
-            backgroundTaskManager: nil
+            localDomain: backendMetadata.domain
         )
         let featureRepository = LegacyFeatureRepository(context: coreDataStack.syncContext)
         let mlsActionExecutor = MLSActionExecutor(

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

@@ -437,8 +437,7 @@ final class UserSessionLoader {
             sharedUserDefaults: sharedUserDefaults,
             syncContext: coreDataStack.syncContext,
             coreCryptoKeyMigrationManager: coreCryptoKeyMigrationManager,
-            localDomain: backendMetadata.domain,
-            backgroundTaskManager: UIApplication.shared
+            localDomain: backendMetadata.domain
         )
 
         let lastEventIDRepository = LastEventIDRepository(

wire-ios-sync-engine/Tests/Source/E2EI/CertificateRevocationListsCheckerTests.swift

@@ -43,10 +43,7 @@ final class CertificateRevocationListsCheckerTests: XCTestCase {
         let coreCrypto = MockCoreCryptoProtocol()
         coreCrypto.mockTransaction(context: mockCoreCryptoContext)
         let provider = MockCoreCryptoProviderProtocol()
-        provider.coreCrypto_MockValue = SafeCoreCrypto(
-            backgroundTaskManager: nil,
-            coreCrypto: coreCrypto
-        )
+        provider.coreCrypto_MockValue = SafeCoreCrypto(coreCrypto: coreCrypto)
 
         coreDataHelper = CoreDataStackHelper()
         mockCoreDataStack = try await coreDataHelper.createStack()

wire-ios-sync-engine/Tests/Source/Use cases/CheckOneOnOneConversationIsReadyUseCaseTests.swift

@@ -44,10 +44,7 @@ class CheckOneOnOneConversationIsReadyUseCaseTests: XCTestCase {
         mockCoreCrypto = MockCoreCryptoProtocol()
         mockCoreCrypto.mockTransaction(context: mockCoreCryptoContext)
         mockCoreCryptoProvider = MockCoreCryptoProviderProtocol()
-        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(
-            backgroundTaskManager: nil,
-            coreCrypto: mockCoreCrypto
-        )
+        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(coreCrypto: mockCoreCrypto)
 
         sut = CheckOneOnOneConversationIsReadyUseCase(
             context: syncMOC,

wire-ios-sync-engine/Tests/Source/UserSession/ZMUserSessionTestsBase.swift

@@ -156,10 +156,7 @@ class ZMUserSessionTestsBase: MessagingTest {
         let mockCoreCrypto = MockCoreCryptoProtocol()
         mockCoreCrypto.registerEpochObserver_MockMethod = { _ in }
         mockCoreCryptoProvider = MockCoreCryptoProviderProtocol()
-        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(
-            backgroundTaskManager: nil,
-            coreCrypto: mockCoreCrypto
-        )
+        mockCoreCryptoProvider.coreCrypto_MockValue = SafeCoreCrypto(coreCrypto: mockCoreCrypto)
 
         let mockContextStorable = MockLAContextStorable()
         mockContextStorable.clear_MockMethod = {}