Fix edge cases in migrations.

Signal/src/ViewControllers/DebugUI/DebugUIBackup.m

@@ -90,7 +90,26 @@ NS_ASSUME_NONNULL_BEGIN
 {
     DDLogInfo(@"%@ tryToImportBackup.", self.logTag);
 
-    [OWSBackup.sharedManager tryToImportBackup];
+    UIAlertController *controller =
+        [UIAlertController alertControllerWithTitle:@"Restore CloudKit Backup"
+                                            message:@"This will delete all of your database contents."
+                                     preferredStyle:UIAlertControllerStyleAlert];
+
+    [controller addAction:[UIAlertAction
+                              actionWithTitle:@"Restore"
+                                        style:UIAlertActionStyleDefault
+                                      handler:^(UIAlertAction *_Nonnull action) {
+                                          [[OWSPrimaryStorage.sharedManager newDatabaseConnection]
+                                              readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
+                                                  [transaction removeAllObjectsInCollection:[TSThread collection]];
+                                                  [transaction removeAllObjectsInCollection:[TSInteraction collection]];
+                                                  [transaction removeAllObjectsInCollection:[TSAttachment collection]];
+                                              }];
+                                          [OWSBackup.sharedManager tryToImportBackup];
+                                      }]];
+    [controller addAction:[OWSAlerts cancelAction]];
+    UIViewController *fromViewController = [[UIApplication sharedApplication] frontmostViewController];
+    [fromViewController presentViewController:controller animated:YES completion:nil];
 }
 
 @end

Signal/src/util/OWSBackupExportJob.m

@@ -11,6 +11,7 @@
 #import <SignalServiceKit/OWSBackupStorage.h>
 #import <SignalServiceKit/OWSError.h>
 #import <SignalServiceKit/OWSFileSystem.h>
+#import <SignalServiceKit/TSAttachment.h>
 #import <SignalServiceKit/TSAttachmentStream.h>
 #import <SignalServiceKit/TSMessage.h>
 #import <SignalServiceKit/TSThread.h>
@@ -268,7 +269,7 @@ NSString *const kOWSBackup_ExportDatabaseKeySpec = @"kOWSBackup_ExportDatabaseKe
 
             // Copy attachments.
             [srcTransaction
-                enumerateKeysAndObjectsInCollection:[TSAttachmentStream collection]
+                enumerateKeysAndObjectsInCollection:[TSAttachment collection]
                                          usingBlock:^(NSString *key, id object, BOOL *stop) {
                                              if (self.isComplete) {
                                                  *stop = YES;

Signal/src/util/OWSBackupImportJob.m

@@ -359,7 +359,6 @@ NSString *const kOWSBackup_ImportDatabaseKeySpec = @"kOWSBackup_ImportDatabaseKe
         [backupStorage runAsyncRegistrationsWithCompletion:^{
             dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
                 [weakSelf restoreDatabaseContents:backupStorage completion:completion];
-                completion(YES);
             });
         }];
     });
@@ -387,113 +386,85 @@ NSString *const kOWSBackup_ImportDatabaseKeySpec = @"kOWSBackup_ImportDatabaseKe
         return completion(NO);
     }
 
-    __block unsigned long long copiedThreads = 0;
+    NSDictionary<NSString *, Class> *collectionTypeMap = @{
-    __block unsigned long long copiedInteractions = 0;
+        [TSThread collection] : [TSThread class],
+        [TSAttachment collection] : [TSAttachment class],
+        [TSInteraction collection] : [TSInteraction class],
+        [OWSDatabaseMigration collection] : [OWSDatabaseMigration class],
+    };
+    // Order matters here.
+    NSArray<NSString *> *collectionsToRestore = @[
+        [TSThread collection],
+        [TSAttachment collection],
+        // Interactions refer to threads and attachments,
+        // so copy them afterward.
+        [TSInteraction collection],
+        [OWSDatabaseMigration collection],
+    ];
+    NSMutableDictionary<NSString *, NSNumber *> *restoredEntityCounts = [NSMutableDictionary new];
     __block unsigned long long copiedEntities = 0;
-    __block unsigned long long copiedAttachments = 0;
+    __block BOOL aborted = NO;
-    __block unsigned long long copiedMigrations = 0;
-
     [tempDBConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *srcTransaction) {
         [primaryDBConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *dstTransaction) {
-            // Copy threads.
+            for (NSString *collection in collectionsToRestore) {
-            [srcTransaction
+                if ([collection isEqualToString:[OWSDatabaseMigration collection]]) {
-                enumerateKeysAndObjectsInCollection:[TSThread collection]
+                    // It's okay if there are existing migrations; we'll clear those
-                                         usingBlock:^(NSString *key, id object, BOOL *stop) {
+                    // before restoring.
-                                             if (self.isComplete) {
+                    continue;
-                                                 *stop = YES;
+                }
-                                                 return;
+                if ([dstTransaction numberOfKeysInCollection:collection] > 0) {
-                                             }
+                    DDLogError(@"%@ cannot restore into non-empty database (%@).", self.logTag, collection);
-                                             if (![object isKindOfClass:[TSThread class]]) {
+                    aborted = YES;
-                                                 OWSProdLogAndFail(
+                    return completion(NO);
-                                                     @"%@ unexpected class: %@", self.logTag, [object class]);
+                }
-                                                 return;
+            }
-                                             }
-                                             TSThread *thread = object;
-                                             [thread saveWithTransaction:dstTransaction];
-                                             copiedThreads++;
-                                             copiedEntities++;
-                                         }];
-
-            // Copy attachments.
-            [srcTransaction
-                enumerateKeysAndObjectsInCollection:[TSAttachmentStream collection]
-                                         usingBlock:^(NSString *key, id object, BOOL *stop) {
-                                             if (self.isComplete) {
-                                                 *stop = YES;
-                                                 return;
-                                             }
-                                             if (![object isKindOfClass:[TSAttachment class]]) {
-                                                 OWSProdLogAndFail(
-                                                     @"%@ unexpected class: %@", self.logTag, [object class]);
-                                                 return;
-                                             }
-                                             TSAttachment *attachment = object;
-                                             [attachment saveWithTransaction:dstTransaction];
-                                             copiedAttachments++;
-                                             copiedEntities++;
-                                         }];
-
-            // Copy interactions.
-            //
-            // Interactions refer to threads and attachments, so copy the last.
-            [srcTransaction
-                enumerateKeysAndObjectsInCollection:[TSInteraction collection]
-                                         usingBlock:^(NSString *key, id object, BOOL *stop) {
-                                             if (self.isComplete) {
-                                                 *stop = YES;
-                                                 return;
-                                             }
-                                             if (![object isKindOfClass:[TSInteraction class]]) {
-                                                 OWSProdLogAndFail(
-                                                     @"%@ unexpected class: %@", self.logTag, [object class]);
-                                                 return;
-                                             }
-                                             // Ignore disappearing messages.
-                                             if ([object isKindOfClass:[TSMessage class]]) {
-                                                 TSMessage *message = object;
-                                                 if (message.isExpiringMessage) {
-                                                     return;
-                                                 }
-                                             }
-                                             TSInteraction *interaction = object;
-                                             // Ignore dynamic interactions.
-                                             if (interaction.isDynamicInteraction) {
-                                                 return;
-                                             }
-                                             [interaction saveWithTransaction:dstTransaction];
-                                             copiedInteractions++;
-                                             copiedEntities++;
-                                         }];
 
             // Clear existing migrations.
+            //
+            // This is safe since we only ever import into an empty database.
+            // Non-database migrations should be idempotent.
             [dstTransaction removeAllObjectsInCollection:[OWSDatabaseMigration collection]];
 
-            // Copy migrations.
+            // Copy database entities.
-            [srcTransaction
+            for (NSString *collection in collectionsToRestore) {
-                enumerateKeysAndObjectsInCollection:[OWSDatabaseMigration collection]
+                [srcTransaction enumerateKeysAndObjectsInCollection:collection
-                                         usingBlock:^(NSString *key, id object, BOOL *stop) {
+                                                         usingBlock:^(NSString *key, id object, BOOL *stop) {
-                                             if (self.isComplete) {
+                                                             if (self.isComplete) {
-                                                 *stop = YES;
+                                                                 *stop = YES;
-                                                 return;
+                                                                 aborted = YES;
-                                             }
+                                                                 return;
-                                             if (![object isKindOfClass:[OWSDatabaseMigration class]]) {
+                                                             }
-                                                 OWSProdLogAndFail(
+                                                             Class expectedType = collectionTypeMap[collection];
-                                                     @"%@ unexpected class: %@", self.logTag, [object class]);
+                                                             OWSAssert(expectedType);
-                                                 return;
+                                                             if (![object isKindOfClass:expectedType]) {
-                                             }
+                                                                 OWSProdLogAndFail(@"%@ unexpected class: %@ != %@",
-                                             OWSDatabaseMigration *migration = object;
+                                                                     self.logTag,
-                                             [migration saveWithTransaction:dstTransaction];
+                                                                     [object class],
-                                             copiedMigrations++;
+                                                                     expectedType);
-                                             copiedEntities++;
+                                                                 return;
-                                         }];
+                                                             }
+                                                             TSYapDatabaseObject *databaseObject = object;
+                                                             [databaseObject saveWithTransaction:dstTransaction];
+
+                                                             NSUInteger count
+                                                                 = restoredEntityCounts[collection].unsignedIntValue;
+                                                             restoredEntityCounts[collection] = @(count + 1);
+                                                             copiedEntities++;
+                                                         }];
+            }
         }];
     }];
 
-    DDLogInfo(@"%@ copiedThreads: %llu", self.logTag, copiedThreads);
+    if (aborted) {
-    DDLogInfo(@"%@ copiedMessages: %llu", self.logTag, copiedInteractions);
+        return;
+    }
+
+    for (NSString *collection in collectionsToRestore) {
+        Class expectedType = collectionTypeMap[collection];
+        OWSAssert(expectedType);
+        DDLogInfo(@"%@ copied %@ (%@): %@", self.logTag, expectedType, collection, restoredEntityCounts[collection]);
+    }
     DDLogInfo(@"%@ copiedEntities: %llu", self.logTag, copiedEntities);
-    DDLogInfo(@"%@ copiedAttachments: %llu", self.logTag, copiedAttachments);
-    DDLogInfo(@"%@ copiedMigrations: %llu", self.logTag, copiedMigrations);
 
     [backupStorage logFileSizes];
 
@@ -510,9 +481,15 @@ NSString *const kOWSBackup_ImportDatabaseKeySpec = @"kOWSBackup_ImportDatabaseKe
 
     DDLogVerbose(@"%@ %s", self.logTag, __PRETTY_FUNCTION__);
 
-    [[[OWSDatabaseMigrationRunner alloc] initWithPrimaryStorage:self.primaryStorage] runAllOutstandingWithCompletion:^{
+    // It's okay that we do this in a separate transaction from the
-        completion(YES);
+    // restoration of backup contents.  If some of migrations don't
-    }];
+    // complete, they'll be run the next time the app launches.
+    dispatch_async(dispatch_get_main_queue(), ^{
+        [[[OWSDatabaseMigrationRunner alloc] initWithPrimaryStorage:self.primaryStorage]
+            runAllOutstandingWithCompletion:^{
+                completion(YES);
+            }];
+    });
 }
 
 - (BOOL)restoreFileWithRecordName:(NSString *)recordName

Signal/src/util/OWSBackupJob.m

@@ -21,6 +21,7 @@ NSString *const kOWSBackup_KeychainService = @"kOWSBackup_KeychainService";
 @property (nonatomic, weak) id<OWSBackupJobDelegate> delegate;
 
 @property (atomic) BOOL isComplete;
+@property (atomic) BOOL hasSucceeded;
 
 @property (nonatomic) OWSPrimaryStorage *primaryStorage;
 
@@ -96,6 +97,12 @@ NSString *const kOWSBackup_KeychainService = @"kOWSBackup_KeychainService";
             return;
         }
         self.isComplete = YES;
+
+        // There's a lot of asynchrony in these backup jobs;
+        // ensure we only end up finishing these jobs once.
+        OWSAssert(!self.hasSucceeded);
+        self.hasSucceeded = YES;
+
         [self.delegate backupJobDidSucceed:self];
     });
 }

SignalMessaging/environment/migrations/OWS106EnsureProfileComplete.swift

@@ -21,11 +21,6 @@ public class OWS106EnsureProfileComplete: OWSDatabaseMigration {
     // Overriding runUp since we have some specific completion criteria which
     // is more likely to fail since it involves network requests.
     override public func runUp(completion:@escaping ((Void)) -> Void) {
-        guard type(of: self).sharedCompleteRegistrationFixerJob == nil else {
-            owsFail("\(self.TAG) should only be called once.")
-            return
-        }
-
         let job = CompleteRegistrationFixerJob(completionHandler: {
             Logger.info("\(self.TAG) Completed. Saving.")
             self.save()

SignalMessaging/environment/migrations/OWSDatabaseMigrationRunner.m

@@ -56,52 +56,48 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)runAllOutstandingWithCompletion:(OWSDatabaseMigrationCompletion)completion
 {
-    [self runMigrations:self.allMigrations completion:completion];
+    [self runMigrations:[self.allMigrations mutableCopy] completion:completion];
 }
 
-- (void)runMigrations:(NSArray<OWSDatabaseMigration *> *)migrations
+// Run migrations serially to:
+//
+// * Ensure predictable ordering.
+// * Prevent them from interfering with each other (e.g. deadlock).
+- (void)runMigrations:(NSMutableArray<OWSDatabaseMigration *> *)migrations
            completion:(OWSDatabaseMigrationCompletion)completion
 {
     OWSAssert(migrations);
     OWSAssert(completion);
 
-    NSMutableArray<OWSDatabaseMigration *> *migrationsToRun = [NSMutableArray new];
+    // TODO: Remove.
+    DDLogInfo(@"%@ Considering migrations: %zd", self.logTag, migrations.count);
     for (OWSDatabaseMigration *migration in migrations) {
-        if ([OWSDatabaseMigration fetchObjectWithUniqueID:migration.uniqueId] == nil) {
+        DDLogInfo(@"%@ Considering migrations: %@", self.logTag, migration.class);
-            [migrationsToRun addObject:migration];
-        }
     }
 
-    if (migrationsToRun.count < 1) {
+    // If there are no more migrations to run, complete.
+    if (migrations.count < 1) {
         dispatch_async(dispatch_get_main_queue(), ^{
             completion();
         });
         return;
     }
 
-    NSUInteger totalMigrationCount = migrationsToRun.count;
+    // Pop next migration from front of queue.
-    __block NSUInteger completedMigrationCount = 0;
+    OWSDatabaseMigration *migration = migrations.firstObject;
-    // Call the completion exactly once, when the last migration completes.
+    [migrations removeObjectAtIndex:0];
-    void (^checkMigrationCompletion)(void) = ^{
+
-        @synchronized(self)
+    // If migration has already been run, skip it.
-        {
+    if ([OWSDatabaseMigration fetchObjectWithUniqueID:migration.uniqueId] != nil) {
-            completedMigrationCount++;
+        [self runMigrations:migrations completion:completion];
-            if (completedMigrationCount == totalMigrationCount) {
+        return;
-                dispatch_async(dispatch_get_main_queue(), ^{
-                    completion();
-                });
-            }
-        }
-    };
-
-    for (OWSDatabaseMigration *migration in migrationsToRun) {
-        if ([OWSDatabaseMigration fetchObjectWithUniqueID:migration.uniqueId]) {
-            DDLogDebug(@"%@ Skipping previously run migration: %@", self.logTag, migration);
-        } else {
-            DDLogWarn(@"%@ Running migration: %@", self.logTag, migration);
-            [migration runUpWithCompletion:checkMigrationCompletion];
-        }
     }
+
+    DDLogInfo(@"%@ Running migration: %@", self.logTag, migration);
+    [migration runUpWithCompletion:^{
+        DDLogInfo(@"%@ Migration complete: %@", self.logTag, migration);
+        [self runMigrations:migrations completion:completion];
+    }];
 }
 
 @end

SignalServiceKit/src/Storage/OWSStorage.m

@@ -74,7 +74,7 @@ typedef NSData *_Nullable (^CreateDatabaseMetadataBlock)(void);
 {
     id<OWSDatabaseConnectionDelegate> delegate = self.delegate;
     OWSAssert(delegate);
-    OWSAssert(delegate.areAllRegistrationsComplete || self.canWriteBeforeStorageReady);
+    //    OWSAssert(delegate.areAllRegistrationsComplete || self.canWriteBeforeStorageReady);
 
     OWSBackgroundTask *_Nullable backgroundTask = nil;
     if (CurrentAppContext().isMainApp) {