Fix flush of registration connections.

SignalServiceKit/src/Storage/OWSPrimaryStorage.m

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

SignalServiceKit/src/Storage/OWSStorage.h

@@ -2,7 +2,7 @@
 //  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
 //
 
-#import <YapDatabase/YapDatabaseConnection.h>
+#import <YapDatabase/YapDatabase.h>
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -34,6 +34,26 @@ extern NSString *const StorageIsReadyNotification;
 
 #pragma mark -
 
+@interface OWSDatabase : YapDatabase
+
+- (instancetype)init NS_UNAVAILABLE;
+
+- (id)initWithPath:(NSString *)inPath
+        serializer:(nullable YapDatabaseSerializer)inSerializer
+      deserializer:(YapDatabaseDeserializer)inDeserializer
+           options:(YapDatabaseOptions *)inOptions
+          delegate:(id<OWSDatabaseConnectionDelegate>)delegate NS_DESIGNATED_INITIALIZER;
+
+// Starts collecting references to the registration connections.
+- (void)collectRegistrationConnections;
+// Stops collecting references to the registration connections and returns
+// all collected connections.
+- (NSSet<YapDatabaseConnection *> *)clearCollectedRegistrationConnections;
+
+@end
+
+#pragma mark -
+
 @interface OWSStorage : NSObject
 
 - (instancetype)init NS_UNAVAILABLE;

SignalServiceKit/src/Storage/OWSStorage.m

@@ -130,17 +130,11 @@ typedef NSData *_Nullable (^CreateDatabaseMetadataBlock)(void);
 
 #pragma mark -
 
-@interface OWSDatabase : YapDatabase
+@interface OWSDatabase ()
 
 @property (atomic, weak) id<OWSDatabaseConnectionDelegate> delegate;
 
-- (instancetype)init NS_UNAVAILABLE;
+@property (nonatomic, readonly, nullable) NSMutableSet<YapDatabaseConnection *> *registrationConnectionSet;
-
-- (id)initWithPath:(NSString *)inPath
-        serializer:(nullable YapDatabaseSerializer)inSerializer
-      deserializer:(YapDatabaseDeserializer)inDeserializer
-           options:(YapDatabaseOptions *)inOptions
-          delegate:(id<OWSDatabaseConnectionDelegate>)delegate NS_DESIGNATED_INITIALIZER;
 
 @end
 
@@ -191,9 +185,29 @@ typedef NSData *_Nullable (^CreateDatabaseMetadataBlock)(void);
     ((OWSDatabaseConnection *)connection).canWriteBeforeStorageReady = YES;
 #endif
 
+    [self.registrationConnectionSet addObject:connection];
+
     return connection;
 }
 
+- (void)collectRegistrationConnections
+{
+    OWSAssert(!self.registrationConnectionSet);
+
+    _registrationConnectionSet = [NSMutableSet set];
+}
+
+- (NSSet<YapDatabaseConnection *> *)clearCollectedRegistrationConnections
+{
+    OWSAssert(self.registrationConnectionSet);
+
+    NSSet<YapDatabaseConnection *> *registrationConnectionSetCopy = [self.registrationConnectionSet copy];
+
+    _registrationConnectionSet = nil;
+
+    return registrationConnectionSetCopy;
+}
+
 @end
 
 #pragma mark -