diff --git a/SignalServiceKit/src/Storage/OWSPrimaryStorage.m b/SignalServiceKit/src/Storage/OWSPrimaryStorage.m
index 307ebe66c..ddff2a59a 100644
--- a/SignalServiceKit/src/Storage/OWSPrimaryStorage.m
+++ b/SignalServiceKit/src/Storage/OWSPrimaryStorage.m
@@ -132,42 +132,26 @@ 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 connections" for this
+    // NOTE: This has to happen on the "registration connection" 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();
-                                     }];
-        }
-    });
+    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();
+                                       }];
 }
 
 + (void)protectFiles
diff --git a/SignalServiceKit/src/Storage/OWSStorage.h b/SignalServiceKit/src/Storage/OWSStorage.h
index 5beab984e..75f710db1 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/YapDatabase.h>
+#import <YapDatabase/YapDatabaseConnection.h>
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -34,26 +34,6 @@ 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 e79cc8069..73fc3680d 100644
--- a/SignalServiceKit/src/Storage/OWSStorage.m
+++ b/SignalServiceKit/src/Storage/OWSStorage.m
@@ -130,11 +130,17 @@ typedef NSData *_Nullable (^CreateDatabaseMetadataBlock)(void);
 
 #pragma mark -
 
-@interface OWSDatabase ()
+@interface OWSDatabase : YapDatabase
 
 @property (atomic, weak) id<OWSDatabaseConnectionDelegate> delegate;
 
-@property (nonatomic, readonly, nullable) NSMutableSet<YapDatabaseConnection *> *registrationConnectionSet;
+- (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;
 
 @end
 
@@ -185,29 +191,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 -