chore: Fix parallel requests for NSE

Sources/MessagingPush/MessagingPushImplementation.swift

@@ -5,11 +5,14 @@ import UserNotifications
 #endif
 
 class MessagingPushImplementation: MessagingPushInstance {
-    let moduleConfig: MessagingPushConfigOptions
-    let logger: Logger
     let pushLogger: PushNotificationLogger
-    let jsonAdapter: JsonAdapter
-    let eventBusHandler: EventBusHandler
+    private let moduleConfig: MessagingPushConfigOptions
+    private let logger: Logger
+    private let jsonAdapter: JsonAdapter
+    private let eventBusHandler: EventBusHandler
+
+    /// Current NSE coordinator when a push is being processed; cleared when run completes or on expiry so `contentHandler` is only ever called once.
+    var nseCoordinator: NSEPushCoordinator?
 
     init(diGraph: DIGraphShared, moduleConfig: MessagingPushConfigOptions) {
         self.moduleConfig = moduleConfig
@@ -31,24 +34,6 @@ class MessagingPushImplementation: MessagingPushInstance {
         eventBusHandler.postEvent(TrackMetricEvent(deliveryID: deliveryID, event: event.rawValue, deviceToken: deviceToken))
     }
 
-    func trackMetricFromNSE(
-        deliveryID: String,
-        event: Metric,
-        deviceToken: String
-    ) {
-        // Access richPushDeliveryTracker from DIGraphShared.shared directly as it is only required for NSE.
-        // Keeping it as class property results in initialization of UserAgentUtil before SDK client is overridden by wrapper SDKs.
-        // In future, we can improve how we access SdkClient so that we don't need to worry about initialization order.
-        DIGraphShared.shared.richPushDeliveryTracker.trackMetric(token: deviceToken, event: event, deliveryId: deliveryID, timestamp: nil) { result in
-            switch result {
-            case .success:
-                self.pushLogger.logPushMetricTracked(deliveryId: deliveryID, event: event.rawValue)
-            case .failure(let error):
-                self.pushLogger.logPushMetricTrackingFailed(deliveryId: deliveryID, event: event.rawValue, error: error)
-            }
-        }
-    }
-
     #if canImport(UserNotifications)
     /**
      - returns:
@@ -63,44 +48,47 @@ class MessagingPushImplementation: MessagingPushInstance {
         let push = UNNotificationWrapper(notificationRequest: request)
         pushLogger.logReceivedPushMessage(notification: push)
 
-        guard let pushCioDeliveryInfo = push.cioDelivery else {
+        guard push.cioDelivery != nil else {
             pushLogger.logReceivedNonCioPushMessage()
             return false
         }
 
         pushLogger.logReceivedCioPushMessage()
-
         if moduleConfig.autoTrackPushEvents {
-            pushLogger.logTrackingPushMessageDelivered(deliveryId: pushCioDeliveryInfo.id)
-
-            trackMetricFromNSE(deliveryID: pushCioDeliveryInfo.id, event: .delivered, deviceToken: pushCioDeliveryInfo.token)
+            pushLogger.logTrackingPushMessageDelivered(deliveryId: push.cioDelivery!.id)
         } else {
             pushLogger.logPushMetricsAutoTrackingDisabled()
         }
 
-        RichPushRequestHandler.shared.startRequest(
-            push: push
-        ) { composedRichPush in
-            self.logger.debug("rich push was composed \(composedRichPush).")
-
-            // This conditional will only work in production and not in automated tests. But this file cannot be in automated tests so this conditional is OK for now.
-            if let composedRichPush = composedRichPush as? UNNotificationWrapper {
-                self.logger.info("Customer.io push processing is done!")
-                contentHandler(composedRichPush.notificationContent)
-            }
+        // Access from DIGraphShared at NSE handling time so wrapper SDKs can override sdkClient first.
+        // Storing as instance properties would initialize UserAgentUtil (via httpClient) in init, before overrides.
+        let coordinator = NSEPushCoordinator(
+            deliveryTracker: DIGraphShared.shared.richPushDeliveryTracker,
+            pushLogger: pushLogger,
+            richPushHandler: RichPushRequestHandler.shared,
+            httpClient: DIGraphShared.shared.httpClient
+        )
+        nseCoordinator = coordinator
+        Task { [weak self] in
+            await coordinator.handle(
+                request: request,
+                withContentHandler: contentHandler,
+                autoTrackDelivery: self?.moduleConfig.autoTrackPushEvents ?? false
+            )
+            self?.nseCoordinator = nil
         }
-
         return true
     }
 
     /**
      iOS telling the notification service to hurry up and stop modifying the push notifications.
-     Stop all network requests and modifying and show the push for what it looks like now.
      */
     func serviceExtensionTimeWillExpire() {
         logger.info("notification service time will expire. Stopping all notification requests early.")
-
-        RichPushRequestHandler.shared.stopAll()
+        if let coordinator = nseCoordinator {
+            Task { await coordinator.cancel() }
+            nseCoordinator = nil
+        }
     }
     #endif
 }

Sources/MessagingPush/PushHandling/NSEPushCoordinator.swift

@@ -0,0 +1,143 @@
+import CioInternalCommon
+import Foundation
+#if canImport(UserNotifications)
+import UserNotifications
+#endif
+
+/// Abstraction for rich push (start and stop) so the coordinator can be tested with a mock.
+protocol RichPushRequestHandling {
+    func startRequest(push: PushNotification, completionHandler: @escaping (PushNotification) -> Void)
+    func stopAll()
+}
+
+/// Coordinates delivery metric request and rich push modification in parallel for the Notification Service Extension.
+/// Holds all NSE logic: delivery tracking (with logging), rich push, parallel run, single contentHandler call, and cancel on expiry.
+/// Create one coordinator and call `handle(request:withContentHandler:autoTrackDelivery:)` to process a notification; call `cancel()` on expiry.
+/// Dependencies are injected for easier testing.
+actor NSEPushCoordinator {
+    private let deliveryTracker: RichPushDeliveryTracker
+    private let pushLogger: PushNotificationLogger
+    private let richPushHandler: RichPushRequestHandling
+    private let httpClient: HttpClient
+
+    private var originalContent: UNNotificationContent?
+    private var contentHandler: ((UNNotificationContent) -> Void)?
+    private var contentHandlerCalled = false
+    private var richPushResult: PushNotification?
+
+    init(
+        deliveryTracker: RichPushDeliveryTracker,
+        pushLogger: PushNotificationLogger,
+        richPushHandler: RichPushRequestHandling,
+        httpClient: HttpClient
+    ) {
+        self.deliveryTracker = deliveryTracker
+        self.pushLogger = pushLogger
+        self.richPushHandler = richPushHandler
+        self.httpClient = httpClient
+    }
+
+    /// Handles the notification: runs delivery (if autoTrackDelivery) and rich push in parallel, waits for both, then calls `contentHandler` exactly once.
+    /// Call this once per notification. Expects a CIO push (cioDelivery non-nil); if not, calls contentHandler with request.content and returns.
+    /// Final content: rich push modified content if present, otherwise original.
+    func handle(
+        request: UNNotificationRequest,
+        withContentHandler contentHandler: @escaping (UNNotificationContent) -> Void,
+        autoTrackDelivery: Bool
+    ) async {
+        let push = UNNotificationWrapper(notificationRequest: request)
+        guard let info = push.cioDelivery else {
+            contentHandler(request.content)
+            return
+        }
+
+        self.originalContent = request.content
+        self.contentHandler = contentHandler
+
+        async let delivery: Void = trackDeliveryIfNeeded(
+            autoTrackDelivery: autoTrackDelivery,
+            token: info.token,
+            deliveryID: info.id
+        )
+
+        async let richPush: PushNotification? = runRichPush(push)
+
+        let result = await richPush
+        richPushResult = result
+
+        _ = await delivery
+
+        guard let originalContent = originalContent else { return }
+        finishWithContent(contentFrom(result, originalContent: originalContent))
+    }
+
+    /// Call when NSE time will expire. Signals cancellation to both operations and calls `contentHandler` exactly once with available content.
+    func cancel() {
+        guard !contentHandlerCalled else { return }
+        guard let originalContent = originalContent else { return }
+
+        contentHandlerCalled = true
+        richPushHandler.stopAll()
+        httpClient.cancel(finishTasks: false)
+
+        contentHandler?(contentFrom(richPushResult, originalContent: originalContent))
+    }
+
+    private func trackDeliveryIfNeeded(
+        autoTrackDelivery: Bool,
+        token: String,
+        deliveryID: String
+    ) async {
+        guard autoTrackDelivery else { return }
+
+        await withCheckedContinuation { (cont: CheckedContinuation<Void, Never>) in
+            deliveryTracker.trackMetric(
+                token: token,
+                event: .delivered,
+                deliveryId: deliveryID,
+                timestamp: nil
+            ) { result in
+                switch result {
+                case .success:
+                    self.pushLogger.logPushMetricTracked(
+                        deliveryId: deliveryID,
+                        event: Metric.delivered.rawValue
+                    )
+                case .failure(let error):
+                    self.pushLogger.logPushMetricTrackingFailed(
+                        deliveryId: deliveryID,
+                        event: Metric.delivered.rawValue,
+                        error: error
+                    )
+                }
+                cont.resume()
+            }
+        }
+    }
+
+    private func runRichPush(_ push: PushNotification) async -> PushNotification? {
+        await withCheckedContinuation { cont in
+            richPushHandler.startRequest(push: push) { composed in
+                cont.resume(returning: composed)
+            }
+        }
+    }
+
+    private func finishWithContent(_ content: UNNotificationContent) {
+        guard !contentHandlerCalled else { return }
+        contentHandlerCalled = true
+        contentHandler?(content)
+    }
+
+    private func contentFrom(
+        _ push: PushNotification?,
+        originalContent: UNNotificationContent
+    ) -> UNNotificationContent {
+        if let wrapper = push as? UNNotificationWrapper {
+            return wrapper.notificationContent
+        }
+        return originalContent
+    }
+}
+
+extension RichPushRequestHandler: RichPushRequestHandling {}

Sources/MessagingPush/RichPush/RichPushRequest.swift

@@ -1,41 +1,51 @@
 import CioInternalCommon
 import Foundation
 
+/// Request to download and attach a rich push image. Only created by the handler when the push
+/// has a valid image URL, so this request always uses the shared HTTP client and cancel() always cancels it.
 class RichPushRequest {
     private let completionHandler: (PushNotification) -> Void
     private var push: PushNotification
+    private let imageURL: URL
     private let httpClient: HttpClient
 
+    private var isCompleted = false
+
     init(
         push: PushNotification,
+        imageURL: URL,
         httpClient: HttpClient,
         completionHandler: @escaping (PushNotification) -> Void
     ) {
         self.completionHandler = completionHandler
         self.push = push
+        self.imageURL = imageURL
         self.httpClient = httpClient
     }
 
     func start() {
-        guard let image = push.cioImage?.url else {
-            // no async operations or modifications to the notification to do. Therefore, let's just finish.
-            return finishImmediately()
-        }
-
-        httpClient.downloadFile(url: image, fileType: .richPushImage) { [weak self] localFilePath in
+        httpClient.downloadFile(url: imageURL, fileType: .richPushImage) { [weak self] localFilePath in
             guard let self = self else { return }
 
             if let localFilePath = localFilePath {
                 self.push.cioRichPushImageFile = localFilePath
             }
 
-            self.finishImmediately()
+            self.completeOnce()
         }
     }
 
-    func finishImmediately() {
+    /// Call when the request must be aborted (e.g. NSE expiry). Cancels the download and completes.
+    func cancel() {
+        guard !isCompleted else { return }
+        isCompleted = true
         httpClient.cancel(finishTasks: false)
+        completionHandler(push)
+    }
 
+    private func completeOnce() {
+        guard !isCompleted else { return }
+        isCompleted = true
         completionHandler(push)
     }
 }

Sources/MessagingPush/RichPush/RichPushRequestHandler.swift

@@ -12,6 +12,14 @@ class RichPushRequestHandler {
         push: PushNotification,
         completionHandler: @escaping (PushNotification) -> Void
     ) {
+        guard let imageURLString = push.cioImage,
+              !imageURLString.isEmpty,
+              let imageURL = URL(string: imageURLString)
+        else {
+            completionHandler(push)
+            return
+        }
+
         let requestId = push.pushId
 
         let existingRequest = requests[requestId]
@@ -22,6 +30,7 @@ class RichPushRequestHandler {
 
         let newRequest = RichPushRequest(
             push: push,
+            imageURL: imageURL,
             httpClient: httpClient,
             completionHandler: completionHandler
         )
@@ -32,7 +41,7 @@ class RichPushRequestHandler {
 
     func stopAll() {
         requests.forEach {
-            $0.value.finishImmediately()
+            $0.value.cancel()
         }
 
         requests = [:]