Rework flush of registration connection(s).

SignalServiceKit/src/Storage/OWSPrimaryStorage.m

@@ -138,36 +138,50 @@ void runAsyncRegistrationsForStorage(OWSStorage *storage)
 
     DDLogVerbose(@"%@ async registrations enqueued.", self.logTag);
 
-    // Block until all async registrations are complete.
+    [[self registrationConnection] asyncReadWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
-    //
+        // Block until all async registrations are complete.
-    // NOTE: This has to happen on the "registration connections" for this
+        //
-    //       database.
+        // NOTE: This has to happen on the "registration connections" for this
-    NSMutableSet<YapDatabaseConnection *> *pendingRegistrationConnectionSet =
+        //       database.
-        [[((OWSDatabase *)self.database)clearCollectedRegistrationConnections] mutableCopy];
+        NSMutableSet<YapDatabaseConnection *> *pendingRegistrationConnectionSet =
-    DDLogVerbose(@"%@ flushing registration connections: %zd.", self.logTag, pendingRegistrationConnectionSet.count);
+            [[((OWSDatabase *)self.database)clearCollectedRegistrationConnections] mutableCopy];
-
+        DDLogVerbose(@"%@ flushing registration connections: %zd.", self.logTag, pendingRegistrationConnectionSet.count);
-    dispatch_async(dispatch_get_main_queue(), ^{
+
-        for (YapDatabaseConnection *dbConnection in pendingRegistrationConnectionSet) {
+        dispatch_async(dispatch_get_main_queue(), ^{
-            [dbConnection
+            if (pendingRegistrationConnectionSet.count > 0) {
-                flushTransactionsWithCompletionQueue:dispatch_get_main_queue()
+                for (YapDatabaseConnection *dbConnection in pendingRegistrationConnectionSet) {
-                                     completionBlock:^{
+                    [dbConnection
-                                         OWSAssertIsOnMainThread();
+                        flushTransactionsWithCompletionQueue:dispatch_get_main_queue()
-                                         OWSAssert(!self.areAsyncRegistrationsComplete);
+                                             completionBlock:^{
-
+                                                 OWSAssertIsOnMainThread();
-                                         [pendingRegistrationConnectionSet removeObject:dbConnection];
+
-                                         if (pendingRegistrationConnectionSet.count > 0) {
+                                                 [pendingRegistrationConnectionSet removeObject:dbConnection];
-                                             DDLogVerbose(@"%@ registration connection flushed.", self.logTag);
+                                                 if (pendingRegistrationConnectionSet.count > 0) {
-                                             return;
+                                                     DDLogVerbose(@"%@ registration connection flushed.", self.logTag);
-                                         }
+                                                     return;
-
+                                                 }
-                                         DDLogVerbose(@"%@ async registrations complete.", self.logTag);
+
-
+                                                 [self markAsyncRegistrationsAsCompleteWithCompletion:completion];
-                                         self.areAsyncRegistrationsComplete = YES;
+                                             }];
-
+                }
-                                         completion();
+            } else {
-                                     }];
+                [self markAsyncRegistrationsAsCompleteWithCompletion:completion];
-        }
+            }
-    });
+        });
+    }];
+}
+
+- (void)markAsyncRegistrationsAsCompleteWithCompletion:(void (^_Nonnull)(void))completion
+{
+    OWSAssertIsOnMainThread();
+    OWSAssert(!self.areAsyncRegistrationsComplete);
+    OWSAssert(completion);
+
+    DDLogVerbose(@"%@ async registrations complete.", self.logTag);
+
+    self.areAsyncRegistrationsComplete = YES;
+
+    completion();
 }
 
 + (void)protectFiles