From d561ba4c62c51cfee4230f434d1b338de1ac78dc Mon Sep 17 00:00:00 2001
From: Matthew Chen <matthew@signal.org>
Date: Wed, 11 Apr 2018 13:09:44 -0400
Subject: [PATCH] Reduce 0xdead10cc crashes.

---
 .../contacts/SystemContactsFetcher.swift      | 16 +++++++++-
 .../profiles/ProfileFetcherJob.swift          | 14 +++++++++
 .../src/Messages/OWSBatchMessageProcessor.m   | 29 ++++++++++++++-----
 3 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/SignalMessaging/contacts/SystemContactsFetcher.swift b/SignalMessaging/contacts/SystemContactsFetcher.swift
index 4c6246899..42bac1828 100644
--- a/SignalMessaging/contacts/SystemContactsFetcher.swift
+++ b/SignalMessaging/contacts/SystemContactsFetcher.swift
@@ -285,10 +285,24 @@ public class SystemContactsFetcher: NSObject {
     private func updateContacts(completion completionParam: ((Error?) -> Void)?, isUserRequested: Bool = false) {
         AssertIsOnMainThread()
 
+        var backgroundTask: OWSBackgroundTask? = OWSBackgroundTask(label: "\(#function)", completionBlock: { [weak self] status in
+            AssertIsOnMainThread()
+
+            guard status == .expired else {
+                return
+            }
+
+            guard let _ = self else {
+                return
+            }
+            Logger.error("background task time ran out contacts fetch completed.")
+        })
+
         // Ensure completion is invoked on main thread.
-        let completion = { error in
+        let completion: (Error?) -> Void = { error in
             DispatchMainThreadSafe({
                 completionParam?(error)
+                backgroundTask = nil
             })
         }
 
diff --git a/SignalMessaging/profiles/ProfileFetcherJob.swift b/SignalMessaging/profiles/ProfileFetcherJob.swift
index 642b84211..600827282 100644
--- a/SignalMessaging/profiles/ProfileFetcherJob.swift
+++ b/SignalMessaging/profiles/ProfileFetcherJob.swift
@@ -19,6 +19,8 @@ public class ProfileFetcherJob: NSObject {
 
     let ignoreThrottling: Bool
 
+    var backgroundTask: OWSBackgroundTask?
+
     @objc
     public class func run(thread: TSThread, networkManager: TSNetworkManager) {
         ProfileFetcherJob(networkManager: networkManager).run(recipientIds: thread.recipientIdentifiers)
@@ -38,6 +40,18 @@ public class ProfileFetcherJob: NSObject {
     public func run(recipientIds: [String]) {
         AssertIsOnMainThread()
 
+        backgroundTask = OWSBackgroundTask(label: "\(#function)", completionBlock: { [weak self] status in
+            AssertIsOnMainThread()
+
+            guard status == .expired else {
+                return
+            }
+            guard let _ = self else {
+                return
+            }
+            Logger.error("background task time ran out before profile fetch completed.")
+        })
+
         if (!CurrentAppContext().isMainApp) {
             // Only refresh profiles in the MainApp to decrease the chance of missed SN notifications
             // in the AppExtension for our users who choose not to verify contacts.
diff --git a/SignalServiceKit/src/Messages/OWSBatchMessageProcessor.m b/SignalServiceKit/src/Messages/OWSBatchMessageProcessor.m
index 510caec14..35cacb16f 100644
--- a/SignalServiceKit/src/Messages/OWSBatchMessageProcessor.m
+++ b/SignalServiceKit/src/Messages/OWSBatchMessageProcessor.m
@@ -310,9 +310,9 @@ NSString *const OWSMessageContentJobFinderExtensionGroup = @"OWSMessageContentJo
     // We want a value that is just high enough to yield perf benefits.
     const NSUInteger kIncomingMessageBatchSize = 32;
 
-    NSArray<OWSMessageContentJob *> *jobs = [self.finder nextJobsForBatchSize:kIncomingMessageBatchSize];
-    OWSAssert(jobs);
-    if (jobs.count < 1) {
+    NSArray<OWSMessageContentJob *> *batchJobs = [self.finder nextJobsForBatchSize:kIncomingMessageBatchSize];
+    OWSAssert(batchJobs);
+    if (batchJobs.count < 1) {
         self.isDrainingQueue = NO;
         DDLogVerbose(@"%@ Queue is drained", self.logTag);
         return;
@@ -320,15 +320,16 @@ NSString *const OWSMessageContentJobFinderExtensionGroup = @"OWSMessageContentJo
 
     OWSBackgroundTask *backgroundTask = [OWSBackgroundTask backgroundTaskWithLabelStr:__PRETTY_FUNCTION__];
 
-    [self processJobs:jobs];
+    NSArray<OWSMessageContentJob *> *processedJobs = [self processJobs:batchJobs];
 
-    [self.finder removeJobsWithIds:jobs.uniqueIds];
+    [self.finder removeJobsWithIds:processedJobs.uniqueIds];
 
     backgroundTask = nil;
 
-    DDLogVerbose(@"%@ completed %zd jobs. %zd jobs left.",
+    DDLogVerbose(@"%@ completed %zd/%zd jobs. %zd jobs left.",
         self.logTag,
-        jobs.count,
+        processedJobs.count,
+        batchJobs.count,
         [OWSMessageContentJob numberOfKeysInCollection]);
 
     // Wait a bit in hopes of increasing the batch size.
@@ -340,17 +341,29 @@ NSString *const OWSMessageContentJobFinderExtensionGroup = @"OWSMessageContentJo
     });
 }
 
-- (void)processJobs:(NSArray<OWSMessageContentJob *> *)jobs
+- (NSArray<OWSMessageContentJob *> *)processJobs:(NSArray<OWSMessageContentJob *> *)jobs
 {
     AssertOnDispatchQueue(self.serialQueue);
 
+    NSMutableArray<OWSMessageContentJob *> *processedJobs = [NSMutableArray new];
     [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
         for (OWSMessageContentJob *job in jobs) {
             [self.messagesManager processEnvelope:job.envelopeProto
                                     plaintextData:job.plaintextData
                                       transaction:transaction];
+            [processedJobs addObject:job];
+
+            if (CurrentAppContext().isInBackground) {
+                // If the app is in the background, stop processing this batch.
+                //
+                // Since this check is done after processing jobs, we'll continue
+                // to process jobs in batches of 1.  This reduces the cost of
+                // being interrupted and rolled back if app is suspended.
+                break;
+            }
         }
     }];
+    return processedJobs;
 }
 
 @end