From e1138df771045c682c4041cc8f43781b535638bd Mon Sep 17 00:00:00 2001
From: Matthew Chen <matthew@signal.org>
Date: Wed, 18 Apr 2018 16:24:18 -0400
Subject: [PATCH] Rework flush of registration connection(s).

---
 .../src/Storage/OWSPrimaryStorage.m           | 52 ++++++-------------
 SignalServiceKit/src/Storage/OWSStorage.h     |  6 ---
 SignalServiceKit/src/Storage/OWSStorage.m     | 22 --------
 3 files changed, 17 insertions(+), 63 deletions(-)

diff --git a/SignalServiceKit/src/Storage/OWSPrimaryStorage.m b/SignalServiceKit/src/Storage/OWSPrimaryStorage.m
index b15cbdc9f..5a2701434 100644
--- a/SignalServiceKit/src/Storage/OWSPrimaryStorage.m
+++ b/SignalServiceKit/src/Storage/OWSPrimaryStorage.m
@@ -132,52 +132,34 @@ void runAsyncRegistrationsForStorage(OWSStorage *storage)
 {
     OWSAssert(completion);
 
-    [((OWSDatabase *)self.database)collectRegistrationConnections];
-
     runAsyncRegistrationsForStorage(self);
 
     DDLogVerbose(@"%@ async registrations enqueued.", self.logTag);
 
     // Flush the write queue to ensure all async registrations have begun.
-    [[self registrationConnection]
+    [[self newDatabaseConnection]
         flushTransactionsWithCompletionQueue:dispatch_get_main_queue()
                              completionBlock:^{
                                  // Block until all async registrations are complete.
                                  //
                                  // NOTE: This has to happen on the "registration connections" for this
                                  //       database.
-                                 NSMutableSet<YapDatabaseConnection *> *pendingRegistrationConnectionSet = [
-                                     [((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) {
-                                         [dbConnection
-                                             flushTransactionsWithCompletionQueue:dispatch_get_main_queue()
-                                                                  completionBlock:^{
-                                                                      OWSAssertIsOnMainThread();
-                                                                      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();
-                                                                  }];
-                                     }
-                                 });
+
+                                 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();
                              }];
 }
 
diff --git a/SignalServiceKit/src/Storage/OWSStorage.h b/SignalServiceKit/src/Storage/OWSStorage.h
index 5beab984e..bff225e68 100644
--- a/SignalServiceKit/src/Storage/OWSStorage.h
+++ b/SignalServiceKit/src/Storage/OWSStorage.h
@@ -44,12 +44,6 @@ extern NSString *const StorageIsReadyNotification;
            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 -
diff --git a/SignalServiceKit/src/Storage/OWSStorage.m b/SignalServiceKit/src/Storage/OWSStorage.m
index e79cc8069..c3a3418c0 100644
--- a/SignalServiceKit/src/Storage/OWSStorage.m
+++ b/SignalServiceKit/src/Storage/OWSStorage.m
@@ -134,8 +134,6 @@ typedef NSData *_Nullable (^CreateDatabaseMetadataBlock)(void);
 
 @property (atomic, weak) id<OWSDatabaseConnectionDelegate> delegate;
 
-@property (nonatomic, readonly, nullable) NSMutableSet<YapDatabaseConnection *> *registrationConnectionSet;
-
 @end
 
 #pragma mark -
@@ -185,29 +183,9 @@ 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 -