Rework flush of registration connection(s).

SignalServiceKit/src/Storage/OWSPrimaryStorage.h

@@ -6,9 +6,6 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
-void runSyncRegistrationsForStorage(OWSStorage *storage);
-void runAsyncRegistrationsForStorage(OWSStorage *storage);
-
 @interface OWSPrimaryStorage : OWSStorage
 
 - (instancetype)init NS_UNAVAILABLE;

SignalServiceKit/src/Storage/OWSPrimaryStorage.m

@@ -30,9 +30,10 @@ void runSyncRegistrationsForStorage(OWSStorage *storage)
     [TSDatabaseView registerCrossProcessNotifier:storage];
 }
 
-void runAsyncRegistrationsForStorage(OWSStorage *storage)
+void runAsyncRegistrationsForStorage(OWSStorage *storage, dispatch_block_t completion)
 {
     OWSCAssert(storage);
+    OWSCAssert(completion);
 
     // Asynchronously register other extensions.
     //
@@ -57,7 +58,9 @@ void runAsyncRegistrationsForStorage(OWSStorage *storage)
     [OWSFailedMessagesJob asyncRegisterDatabaseExtensionsWithPrimaryStorage:storage];
     [OWSFailedAttachmentDownloadsJob asyncRegisterDatabaseExtensionsWithPrimaryStorage:storage];
     [OWSMediaGalleryFinder asyncRegisterDatabaseExtensionsWithPrimaryStorage:storage];
-    [TSDatabaseView asyncRegisterLazyRestoreAttachmentsDatabaseView:storage];
+    // NOTE: Always pass the completion to the _LAST_ of the async database
+    // view registrations.
+    [TSDatabaseView asyncRegisterLazyRestoreAttachmentsDatabaseView:storage completion:completion];
 }
 
 #pragma mark -
@@ -132,27 +135,19 @@ void runAsyncRegistrationsForStorage(OWSStorage *storage)
 {
     OWSAssert(completion);
 
-    runAsyncRegistrationsForStorage(self);
+    DDLogVerbose(@"%@ async registrations enqueuing.", self.logTag);
 
-    DDLogVerbose(@"%@ async registrations enqueued.", self.logTag);
+    runAsyncRegistrationsForStorage(self, ^{
+        OWSAssertIsOnMainThread();
 
-    // Use an empty read/write transaction to to ensure all async registrations have completed.
+        OWSAssert(!self.areAsyncRegistrationsComplete);
-    [[self newDatabaseConnection] asyncReadWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
-        // Do nothing.
-        //
-        // We can't use flushTransactionsWithCompletionQueue because it
-        // doesn't flush the YapDatabase.writeQueue.
-    }
-        completionQueue:dispatch_get_main_queue()
-        completionBlock:^{
-            OWSAssert(!self.areAsyncRegistrationsComplete);
 
-            DDLogVerbose(@"%@ async registrations complete.", self.logTag);
+        DDLogVerbose(@"%@ async registrations complete.", self.logTag);
 
-            self.areAsyncRegistrationsComplete = YES;
+        self.areAsyncRegistrationsComplete = YES;
 
-            completion();
+        completion();
-        }];
+    });
 }
 
 + (void)protectFiles

SignalServiceKit/src/Storage/OWSStorage.h

@@ -71,6 +71,9 @@ extern NSString *const StorageIsReadyNotification;
 #endif
 
 - (void)asyncRegisterExtension:(YapDatabaseExtension *)extension withName:(NSString *)extensionName;
+- (void)asyncRegisterExtension:(YapDatabaseExtension *)extension
+                      withName:(NSString *)extensionName
+                    completion:(nullable dispatch_block_t)completion;
 
 - (nullable id)registeredExtension:(NSString *)extensionName;
 

SignalServiceKit/src/Storage/OWSStorage.m

@@ -487,6 +487,13 @@ typedef NSData *_Nullable (^CreateDatabaseMetadataBlock)(void);
 
 - (void)asyncRegisterExtension:(YapDatabaseExtension *)extension
                       withName:(NSString *)extensionName
+{
+    [self asyncRegisterExtension:extension withName:extensionName completion:nil];
+}
+
+- (void)asyncRegisterExtension:(YapDatabaseExtension *)extension
+                      withName:(NSString *)extensionName
+                    completion:(nullable dispatch_block_t)completion
 {
     [self.database asyncRegisterExtension:extension
                                  withName:extensionName
@@ -496,6 +503,12 @@ typedef NSData *_Nullable (^CreateDatabaseMetadataBlock)(void);
                               } else {
                                   DDLogVerbose(@"%@ asyncRegisterExtension succeeded: %@", self.logTag, extensionName);
                               }
+
+                              dispatch_async(dispatch_get_main_queue(), ^{
+                                  if (completion) {
+                                      completion();
+                                  }
+                              });
                           }];
 }
 

SignalServiceKit/src/Storage/TSDatabaseView.h

@@ -58,6 +58,7 @@ extern NSString *const TSLazyRestoreAttachmentsDatabaseViewExtensionName;
 
 + (void)asyncRegisterSecondaryDevicesDatabaseView:(OWSStorage *)storage;
 
-+ (void)asyncRegisterLazyRestoreAttachmentsDatabaseView:(OWSStorage *)storage;
++ (void)asyncRegisterLazyRestoreAttachmentsDatabaseView:(OWSStorage *)storage
+                                             completion:(nullable dispatch_block_t)completion;
 
 @end

SignalServiceKit/src/Storage/TSDatabaseView.m

@@ -343,6 +343,7 @@ NSString *const TSLazyRestoreAttachmentsGroup = @"TSLazyRestoreAttachmentsGroup"
 }
 
 + (void)asyncRegisterLazyRestoreAttachmentsDatabaseView:(OWSStorage *)storage
+                                             completion:(nullable dispatch_block_t)completion
 {
     YapDatabaseViewGrouping *viewGrouping = [YapDatabaseViewGrouping withObjectBlock:^NSString *_Nullable(
         YapDatabaseReadTransaction *transaction, NSString *collection, NSString *key, id object) {
@@ -391,7 +392,9 @@ NSString *const TSLazyRestoreAttachmentsGroup = @"TSLazyRestoreAttachmentsGroup"
         [[YapWhitelistBlacklist alloc] initWithWhitelist:[NSSet setWithObject:[TSAttachment collection]]];
     YapDatabaseView *view =
         [[YapDatabaseAutoView alloc] initWithGrouping:viewGrouping sorting:viewSorting versionTag:@"3" options:options];
-    [storage asyncRegisterExtension:view withName:TSLazyRestoreAttachmentsDatabaseViewExtensionName];
+    [storage asyncRegisterExtension:view
+                           withName:TSLazyRestoreAttachmentsDatabaseViewExtensionName
+                         completion:completion];
 }
 
 + (id)unseenDatabaseViewExtension:(YapDatabaseReadTransaction *)transaction