From 27e5c836b70ed8839bc7cb75306c86b5d20e9e1e Mon Sep 17 00:00:00 2001
From: Matthew Chen <charlesmchen@gmail.com>
Date: Tue, 20 Jun 2017 12:04:59 -0400
Subject: [PATCH] Refine observation of async registration completion.

// FREEBIE
---
 src/Storage/TSDatabaseView.h   |  3 +++
 src/Storage/TSDatabaseView.m   | 40 ++++++++++++++++------------------
 src/Storage/TSStorageManager.m |  5 +++++
 3 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/src/Storage/TSDatabaseView.h b/src/Storage/TSDatabaseView.h
index 97500b29e..92279916c 100644
--- a/src/Storage/TSDatabaseView.h
+++ b/src/Storage/TSDatabaseView.h
@@ -54,4 +54,7 @@ extern NSString *const TSSecondaryDevicesDatabaseViewExtensionName;
 // NOTE: It is not safe to call this method while hasPendingViewRegistrations is YES.
 + (id)threadSpecialMessagesDatabaseView:(YapDatabaseReadTransaction *)transaction;
 
+// This method should be called _after_ all async database registrations have been started.
++ (void)asyncRegistrationCompletion;
+
 @end
diff --git a/src/Storage/TSDatabaseView.m b/src/Storage/TSDatabaseView.m
index a5efe6435..a11185252 100644
--- a/src/Storage/TSDatabaseView.m
+++ b/src/Storage/TSDatabaseView.m
@@ -31,7 +31,7 @@ NSString *const TSSecondaryDevicesDatabaseViewExtensionName = @"TSSecondaryDevic
 
 @interface TSDatabaseView ()
 
-@property (nonatomic) int pendingViewRegistrations;
+@property (nonatomic) BOOL areAllAsyncRegistrationsComplete;
 
 @end
 
@@ -62,24 +62,11 @@ NSString *const TSSecondaryDevicesDatabaseViewExtensionName = @"TSSecondaryDevic
     return self;
 }
 
-- (void)setPendingViewRegistrations:(int)pendingViewRegistrations
-{
-    OWSAssert([NSThread isMainThread]);
-
-    _pendingViewRegistrations = pendingViewRegistrations;
-
-    if (pendingViewRegistrations == 0) {
-        [[NSNotificationCenter defaultCenter] postNotificationName:kNSNotificationName_DatabaseViewRegistrationComplete
-                                                            object:nil
-                                                          userInfo:nil];
-    }
-}
-
 + (BOOL)hasPendingViewRegistrations
 {
     OWSAssert([NSThread isMainThread]);
 
-    return [TSDatabaseView sharedInstance].pendingViewRegistrations > 0;
+    return ![TSDatabaseView sharedInstance].areAllAsyncRegistrationsComplete;
 }
 
 + (void)registerMessageDatabaseViewWithName:(NSString *)viewName
@@ -107,8 +94,6 @@ NSString *const TSSecondaryDevicesDatabaseViewExtensionName = @"TSSecondaryDevic
         [[YapDatabaseView alloc] initWithGrouping:viewGrouping sorting:viewSorting versionTag:version options:options];
 
     if (async) {
-        TSDatabaseView.sharedInstance.pendingViewRegistrations++;
-
         [[TSStorageManager sharedManager].database
             asyncRegisterExtension:view
                           withName:viewName
@@ -116,10 +101,6 @@ NSString *const TSSecondaryDevicesDatabaseViewExtensionName = @"TSSecondaryDevic
                        OWSCAssert(ready);
 
                        DDLogInfo(@"%@ asyncRegisterExtension: %@ -> %d", self.tag, viewName, ready);
-
-                       dispatch_async(dispatch_get_main_queue(), ^{
-                           TSDatabaseView.sharedInstance.pendingViewRegistrations--;
-                       });
                    }];
     } else {
         [[TSStorageManager sharedManager].database registerExtension:view withName:viewName];
@@ -417,6 +398,23 @@ NSString *const TSSecondaryDevicesDatabaseViewExtensionName = @"TSSecondaryDevic
     return result;
 }
 
++ (void)asyncRegistrationCompletion
+{
+    OWSAssert([NSThread isMainThread]);
+
+    // All async registrations are complete when writes are unblocked.
+    [[TSStorageManager sharedManager].newDatabaseConnection
+        asyncReadWriteWithBlock:^(YapDatabaseReadWriteTransaction *_Nonnull transaction) {
+            dispatch_async(dispatch_get_main_queue(), ^{
+                TSDatabaseView.sharedInstance.areAllAsyncRegistrationsComplete = YES;
+                [[NSNotificationCenter defaultCenter]
+                    postNotificationName:kNSNotificationName_DatabaseViewRegistrationComplete
+                                  object:nil
+                                userInfo:nil];
+            });
+        }];
+}
+
 #pragma mark - Logging
 
 + (NSString *)tag
diff --git a/src/Storage/TSStorageManager.m b/src/Storage/TSStorageManager.m
index 6b392dda3..c429b12f0 100644
--- a/src/Storage/TSStorageManager.m
+++ b/src/Storage/TSStorageManager.m
@@ -225,6 +225,11 @@ static NSString *keychainDBPassAccount    = @"TSDatabasePass";
     OWSFailedAttachmentDownloadsJob *failedAttachmentDownloadsMessagesJob =
         [[OWSFailedAttachmentDownloadsJob alloc] initWithStorageManager:self];
     [failedAttachmentDownloadsMessagesJob asyncRegisterDatabaseExtensions];
+
+    // NOTE: [TSDatabaseView asyncRegistrationCompletion] ensures that
+    // kNSNotificationName_DatabaseViewRegistrationComplete is not fired until all
+    // of the async registrations are complete.
+    [TSDatabaseView asyncRegistrationCompletion];
 }
 
 - (void)protectSignalFiles {