diff --git a/SignalServiceKit/src/Storage/OWSPrimaryStorage.m b/SignalServiceKit/src/Storage/OWSPrimaryStorage.m index 307ebe66c..0885c8f95 100644 --- a/SignalServiceKit/src/Storage/OWSPrimaryStorage.m +++ b/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. - // - // NOTE: This has to happen on the "registration connections" for this - // database. - NSMutableSet *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(); - }]; - } - }); + [[self registrationConnection] asyncReadWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) { + // Block until all async registrations are complete. + // + // NOTE: This has to happen on the "registration connections" for this + // database. + NSMutableSet *pendingRegistrationConnectionSet = + [[((OWSDatabase *)self.database)clearCollectedRegistrationConnections] mutableCopy]; + DDLogVerbose(@"%@ flushing registration connections: %zd.", self.logTag, pendingRegistrationConnectionSet.count); + + dispatch_async(dispatch_get_main_queue(), ^{ + if (pendingRegistrationConnectionSet.count > 0) { + for (YapDatabaseConnection *dbConnection in pendingRegistrationConnectionSet) { + [dbConnection + flushTransactionsWithCompletionQueue:dispatch_get_main_queue() + completionBlock:^{ + OWSAssertIsOnMainThread(); + + [pendingRegistrationConnectionSet removeObject:dbConnection]; + if (pendingRegistrationConnectionSet.count > 0) { + DDLogVerbose(@"%@ registration connection flushed.", self.logTag); + return; + } + + [self markAsyncRegistrationsAsCompleteWithCompletion: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