diff --git a/SignalServiceKit/src/Storage/OWSPrimaryStorage.m b/SignalServiceKit/src/Storage/OWSPrimaryStorage.m index e05947567..ea8e8b3d7 100644 --- a/SignalServiceKit/src/Storage/OWSPrimaryStorage.m +++ b/SignalServiceKit/src/Storage/OWSPrimaryStorage.m @@ -131,26 +131,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; - OWSAssert(self.registrationConnection); - [dbConnection flushTransactionsWithCompletionQueue:dispatch_get_main_queue() - completionBlock:^{ - OWSAssert(!self.areAsyncRegistrationsComplete); - - DDLogVerbose(@"%@ async registrations complete.", self.logTag); - - self.areAsyncRegistrationsComplete = YES; - - completion(); - }]; + 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(); + }]; + } + }); } + (void)protectFiles diff --git a/SignalServiceKit/src/Storage/OWSStorage.h b/SignalServiceKit/src/Storage/OWSStorage.h index 75f710db1..5beab984e 100644 --- a/SignalServiceKit/src/Storage/OWSStorage.h +++ b/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; diff --git a/SignalServiceKit/src/Storage/OWSStorage.m b/SignalServiceKit/src/Storage/OWSStorage.m index 73fc3680d..e79cc8069 100644 --- a/SignalServiceKit/src/Storage/OWSStorage.m +++ b/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; - -- (id)initWithPath:(NSString *)inPath - serializer:(nullable YapDatabaseSerializer)inSerializer - deserializer:(YapDatabaseDeserializer)inDeserializer - options:(YapDatabaseOptions *)inOptions - delegate:(id<OWSDatabaseConnectionDelegate>)delegate NS_DESIGNATED_INITIALIZER; +@property (nonatomic, readonly, nullable) NSMutableSet<YapDatabaseConnection *> *registrationConnectionSet; @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 -