From f511ca35df2d50de88e2da1659f4af35d570d727 Mon Sep 17 00:00:00 2001
From: Matthew Chen <matthew@signal.org>
Date: Fri, 24 Sep 2021 15:25:58 -0300
Subject: [PATCH] Refine NSE completion.

---
 .../NotificationService.swift                 | 19 +++++++++----------
 .../src/Messages/MessageFetcherJob.swift      |  6 ++++++
 .../src/Messages/MessageProcessor.swift       | 10 +++++-----
 SignalServiceKit/src/Messages/MessageSender.h |  3 +++
 .../src/Messages/MessageSender.swift          |  6 ++++++
 5 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/NotificationServiceExtension/NotificationService.swift b/NotificationServiceExtension/NotificationService.swift
index 5d7fc6d24f..50622ae65f 100644
--- a/NotificationServiceExtension/NotificationService.swift
+++ b/NotificationServiceExtension/NotificationService.swift
@@ -58,14 +58,7 @@ class NotificationService: UNNotificationServiceExtension {
         }
         Logger.flush()
 
-        let isSync = timeHasExpired
-        if isSync {
-            contentHandler(content)
-        } else {
-            DispatchQueue.global().asyncAfter(deadline: .now() + .milliseconds(1000)) {
-                contentHandler(content)
-            }
-        }
+        contentHandler(content)
     }
 
     override func didReceive(_ request: UNNotificationRequest, withContentHandler contentHandler: @escaping (UNNotificationContent) -> Void) {
@@ -150,11 +143,17 @@ class NotificationService: UNNotificationServiceExtension {
         }.catch { _ in
             // Do nothing, Promise.timeout() will log timeouts.
         }
-        fetchPromise.then { [weak self] () -> Promise<Void> in
+        fetchPromise.then(on: .global()) { [weak self] () -> Promise<Void> in
             Logger.info("Waiting for processing to complete.")
             guard let self = self else { return Promise.value(()) }
             let processingCompletePromise = self.messageProcessor.processingCompletePromise()
-            processingCompletePromise.timeout(seconds: 20, description: "Message Processing Timeout.") {
+            processingCompletePromise.then(on: .global()) {
+                // Wait until all ACKs are enqueued.
+                Self.messageFetcherJob.pendingAcksPromise()
+            }.then(on: .global()) {
+                // Wait until all outgoing messages are sent.
+                Self.messageSender.pendingSendsPromise()
+            }.timeout(seconds: 20, description: "Message Processing Timeout.") {
                 NotificationServiceError.timeout
             }.catch { _ in
                 // Do nothing, Promise.timeout() will log timeouts.
diff --git a/SignalServiceKit/src/Messages/MessageFetcherJob.swift b/SignalServiceKit/src/Messages/MessageFetcherJob.swift
index a8fa7cb7ed..687979697e 100644
--- a/SignalServiceKit/src/Messages/MessageFetcherJob.swift
+++ b/SignalServiceKit/src/Messages/MessageFetcherJob.swift
@@ -273,6 +273,12 @@ public class MessageFetcherJob: NSObject {
         ackOperationQueue.addOperation(ackOperation)
     }
 
+    public func pendingAcksPromise() -> Promise<Void> {
+        firstly(on: .global()) {
+            self.ackOperationQueue.waitUntilAllOperationsAreFinished()
+        }
+    }
+
     // MARK: -
 
     typealias EnvelopeJob = MessageProcessor.EnvelopeJob
diff --git a/SignalServiceKit/src/Messages/MessageProcessor.swift b/SignalServiceKit/src/Messages/MessageProcessor.swift
index 76aa0c8fb6..7e303a45c0 100644
--- a/SignalServiceKit/src/Messages/MessageProcessor.swift
+++ b/SignalServiceKit/src/Messages/MessageProcessor.swift
@@ -654,15 +654,15 @@ private struct EncryptedEnvelope: PendingEnvelope, Dependencies {
         guard let other = other as? EncryptedEnvelope else {
             return false
         }
-        // serverDeliveryTimestamp is a cheaper comparison and is likely
-        // to eliminate most candidates.
-        guard self.serverDeliveryTimestamp == other.serverDeliveryTimestamp else {
+        guard let serverGuid = encryptedEnvelope.serverGuid else {
+            owsFailDebug("Missing serverGuid.")
             return false
         }
-        guard self.encryptedEnvelopeData == other.encryptedEnvelopeData else {
+        guard let otherServerGuid = other.encryptedEnvelope.serverGuid else {
+            owsFailDebug("Missing other.serverGuid.")
             return false
         }
-        return true
+        return serverGuid == otherServerGuid
     }
 }
 
diff --git a/SignalServiceKit/src/Messages/MessageSender.h b/SignalServiceKit/src/Messages/MessageSender.h
index 50145ca3dc..d6e74a5948 100644
--- a/SignalServiceKit/src/Messages/MessageSender.h
+++ b/SignalServiceKit/src/Messages/MessageSender.h
@@ -101,6 +101,9 @@ NS_SWIFT_NAME(OutgoingAttachmentInfo)
 // TODO: Make this private.
 - (void)sendMessageToRecipient:(OWSMessageSend *)messageSend;
 
+// TODO: Make this private.
++ (NSOperationQueue *)globalSendingQueue;
+
 @end
 
 #pragma mark -
diff --git a/SignalServiceKit/src/Messages/MessageSender.swift b/SignalServiceKit/src/Messages/MessageSender.swift
index cb5803d063..74937ea934 100644
--- a/SignalServiceKit/src/Messages/MessageSender.swift
+++ b/SignalServiceKit/src/Messages/MessageSender.swift
@@ -152,6 +152,12 @@ extension MessageSender {
         }
         return promises
     }
+
+    public func pendingSendsPromise() -> Promise<Void> {
+        firstly(on: .global()) {
+            Self.globalSendingQueue().waitUntilAllOperationsAreFinished()
+        }
+    }
 }
 
 // MARK: -