From 59f480d5c473825eb894572e967c51f598f4ec83 Mon Sep 17 00:00:00 2001
From: Matthew Chen <charlesmchen@gmail.com>
Date: Tue, 27 Feb 2018 08:42:30 -0500
Subject: [PATCH] Use background tasks during storage registration and
 disappearing message jobs.

---
 .../src/Messages/OWSDisappearingMessagesJob.m | 13 ++++++++++++
 SignalServiceKit/src/Storage/OWSStorage.m     | 20 ++++++++++++++-----
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/SignalServiceKit/src/Messages/OWSDisappearingMessagesJob.m b/SignalServiceKit/src/Messages/OWSDisappearingMessagesJob.m
index 16c058ef4..fd0e3c17a 100644
--- a/SignalServiceKit/src/Messages/OWSDisappearingMessagesJob.m
+++ b/SignalServiceKit/src/Messages/OWSDisappearingMessagesJob.m
@@ -8,6 +8,7 @@
 #import "ContactsManagerProtocol.h"
 #import "NSDate+OWS.h"
 #import "NSTimer+OWS.h"
+#import "OWSBackgroundTask.h"
 #import "OWSDisappearingConfigurationUpdateInfoMessage.h"
 #import "OWSDisappearingMessagesConfiguration.h"
 #import "OWSDisappearingMessagesFinder.h"
@@ -88,6 +89,8 @@ NS_ASSUME_NONNULL_BEGIN
 {
     uint64_t now = [NSDate ows_millisecondTimeStamp];
 
+    OWSBackgroundTask *_Nullable backgroundTask = [OWSBackgroundTask backgroundTaskWithLabelStr:__PRETTY_FUNCTION__];
+
     __block uint expirationCount = 0;
     [self.databaseConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *_Nonnull transaction) {
         [self.disappearingMessagesFinder enumerateExpiredMessagesWithBlock:^(TSMessage *message) {
@@ -106,6 +109,8 @@ NS_ASSUME_NONNULL_BEGIN
     }];
 
     DDLogDebug(@"%@ Removed %u expired messages", self.logTag, expirationCount);
+
+    backgroundTask = nil;
 }
 
 // This method should only be called on the serialQueue.
@@ -203,6 +208,8 @@ NS_ASSUME_NONNULL_BEGIN
 // This method should only be called on the serialQueue.
 - (void)setExpirationsForThread:(TSThread *)thread
 {
+    OWSBackgroundTask *_Nullable backgroundTask = [OWSBackgroundTask backgroundTaskWithLabelStr:__PRETTY_FUNCTION__];
+
     uint64_t now = [NSDate ows_millisecondTimeStamp];
     [self.databaseConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *_Nonnull transaction) {
         [self.disappearingMessagesFinder
@@ -220,6 +227,8 @@ NS_ASSUME_NONNULL_BEGIN
                                                  }
                                            transaction:transaction];
     }];
+
+    backgroundTask = nil;
 }
 
 + (void)becomeConsistentWithConfigurationForMessage:(TSMessage *)message
@@ -234,6 +243,8 @@ NS_ASSUME_NONNULL_BEGIN
     OWSAssert(message);
     OWSAssert(contactsManager);
 
+    OWSBackgroundTask *_Nullable backgroundTask = [OWSBackgroundTask backgroundTaskWithLabelStr:__PRETTY_FUNCTION__];
+
     dispatch_async(OWSDisappearingMessagesJob.serialQueue, ^{
         // Become eventually consistent in the case that the remote changed their settings at the same time.
         // Also in case remote doesn't support expiring messages
@@ -280,6 +291,8 @@ NS_ASSUME_NONNULL_BEGIN
                                                                         configuration:disappearingMessagesConfiguration]
                 save];
         }
+
+        backgroundTask = nil;
     });
 }
 
diff --git a/SignalServiceKit/src/Storage/OWSStorage.m b/SignalServiceKit/src/Storage/OWSStorage.m
index 6687f5e57..cbf424f5c 100644
--- a/SignalServiceKit/src/Storage/OWSStorage.m
+++ b/SignalServiceKit/src/Storage/OWSStorage.m
@@ -350,16 +350,23 @@ typedef NSData *_Nullable (^CreateDatabaseMetadataBlock)(void);
 
 + (void)setupStorage
 {
+    __block OWSBackgroundTask *_Nullable backgroundTask =
+        [OWSBackgroundTask backgroundTaskWithLabelStr:__PRETTY_FUNCTION__];
+
     for (OWSStorage *storage in self.allStorages) {
         [storage runSyncRegistrations];
     }
 
     for (OWSStorage *storage in self.allStorages) {
         [storage runAsyncRegistrationsWithCompletion:^{
-            
-            [self postRegistrationCompleteNotificationIfPossible];
+            if ([self postRegistrationCompleteNotificationIfPossible]) {
+                // If all registrations are complete, clean up the
+                // registration process.
+
+                ((OWSDatabase *)storage.database).registrationConnectionCached = nil;
 
-            ((OWSDatabase *)storage.database).registrationConnectionCached = nil;
+                backgroundTask = nil;
+            }
         }];
     }
 }
@@ -369,10 +376,11 @@ typedef NSData *_Nullable (^CreateDatabaseMetadataBlock)(void);
     return self.database.registrationConnection;
 }
 
-+ (void)postRegistrationCompleteNotificationIfPossible
+// Returns YES IFF all registrations are complete.
++ (BOOL)postRegistrationCompleteNotificationIfPossible
 {
     if (!self.isStorageReady) {
-        return;
+        return NO;
     }
 
     static dispatch_once_t onceToken;
@@ -381,6 +389,8 @@ typedef NSData *_Nullable (^CreateDatabaseMetadataBlock)(void);
                                                                  object:nil
                                                                userInfo:nil];
     });
+
+    return YES;
 }
 
 + (BOOL)isStorageReady