From f1646b6cbcfee0d1fa2f15ae682b79cba90ca9f0 Mon Sep 17 00:00:00 2001
From: Matthew Chen <matthew@signal.org>
Date: Tue, 16 Oct 2018 13:00:50 -0400
Subject: [PATCH] Move db connections to environments.

---
 .../migrations/OWSDatabaseMigration.m         |  9 ++----
 .../Messages/Interactions/TSOutgoingMessage.m |  7 +----
 SignalServiceKit/src/SSKEnvironment.h         |  3 ++
 SignalServiceKit/src/SSKEnvironment.m         | 31 +++++++++++++++++++
 .../OWSPrimaryStorage+SessionStore.m          | 12 ++-----
 5 files changed, 40 insertions(+), 22 deletions(-)

diff --git a/SignalMessaging/environment/migrations/OWSDatabaseMigration.m b/SignalMessaging/environment/migrations/OWSDatabaseMigration.m
index 56cf5eb6c..1a7c7702f 100644
--- a/SignalMessaging/environment/migrations/OWSDatabaseMigration.m
+++ b/SignalMessaging/environment/migrations/OWSDatabaseMigration.m
@@ -4,6 +4,7 @@
 
 #import "OWSDatabaseMigration.h"
 #import <SignalServiceKit/OWSPrimaryStorage.h>
+#import <SignalServiceKit/SSKEnvironment.h>
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -81,13 +82,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 + (YapDatabaseConnection *)dbReadWriteConnection
 {
-    static dispatch_once_t onceToken;
-    static YapDatabaseConnection *sharedDBConnection;
-    dispatch_once(&onceToken, ^{
-        sharedDBConnection = [OWSPrimaryStorage sharedManager].newDatabaseConnection;
-    });
-
-    return sharedDBConnection;
+    return SSKEnvironment.shared.migrationDBConnection;
 }
 
 - (YapDatabaseConnection *)dbReadConnection
diff --git a/SignalServiceKit/src/Messages/Interactions/TSOutgoingMessage.m b/SignalServiceKit/src/Messages/Interactions/TSOutgoingMessage.m
index 8bd41456e..1597449e3 100644
--- a/SignalServiceKit/src/Messages/Interactions/TSOutgoingMessage.m
+++ b/SignalServiceKit/src/Messages/Interactions/TSOutgoingMessage.m
@@ -220,12 +220,7 @@ NSString *NSStringForOutgoingMessageRecipientState(OWSOutgoingMessageRecipientSt
 
 + (YapDatabaseConnection *)dbMigrationConnection
 {
-    static YapDatabaseConnection *connection = nil;
-    static dispatch_once_t onceToken;
-    dispatch_once(&onceToken, ^{
-        connection = [[OWSPrimaryStorage sharedManager] newDatabaseConnection];
-    });
-    return connection;
+    return SSKEnvironment.shared.migrationDBConnection;
 }
 
 + (instancetype)outgoingMessageInThread:(nullable TSThread *)thread
diff --git a/SignalServiceKit/src/SSKEnvironment.h b/SignalServiceKit/src/SSKEnvironment.h
index 53495deae..02d6e04c9 100644
--- a/SignalServiceKit/src/SSKEnvironment.h
+++ b/SignalServiceKit/src/SSKEnvironment.h
@@ -90,6 +90,9 @@ NS_ASSUME_NONNULL_BEGIN
 @property (atomic, nullable) id<NotificationsProtocol> notificationsManager;
 
 @property (atomic, readonly) YapDatabaseConnection *objectReadWriteConnection;
+@property (atomic, readonly) YapDatabaseConnection *sessionStoreDBConnection;
+@property (atomic, readonly) YapDatabaseConnection *migrationDBConnection;
+@property (atomic, readonly) YapDatabaseConnection *analyticsDBConnection;
 
 - (BOOL)isComplete;
 
diff --git a/SignalServiceKit/src/SSKEnvironment.m b/SignalServiceKit/src/SSKEnvironment.m
index 8bcb13e83..8aefcc811 100644
--- a/SignalServiceKit/src/SSKEnvironment.m
+++ b/SignalServiceKit/src/SSKEnvironment.m
@@ -42,6 +42,9 @@ static SSKEnvironment *sharedSSKEnvironment;
 @synthesize callMessageHandler = _callMessageHandler;
 @synthesize notificationsManager = _notificationsManager;
 @synthesize objectReadWriteConnection = _objectReadWriteConnection;
+@synthesize sessionStoreDBConnection = _sessionStoreDBConnection;
+@synthesize migrationDBConnection = _migrationDBConnection;
+@synthesize analyticsDBConnection = _analyticsDBConnection;
 
 - (instancetype)initWithContactsManager:(id<ContactsManagerProtocol>)contactsManager
                           messageSender:(OWSMessageSender *)messageSender
@@ -188,6 +191,34 @@ static SSKEnvironment *sharedSSKEnvironment;
         return _objectReadWriteConnection;
     }
 }
+
+- (YapDatabaseConnection *)sessionStoreDBConnection {
+    @synchronized(self) {
+        if (!_sessionStoreDBConnection) {
+            _sessionStoreDBConnection = self.primaryStorage.newDatabaseConnection;
+        }
+        return _sessionStoreDBConnection;
+    }
+}
+
+- (YapDatabaseConnection *)migrationDBConnection {
+    @synchronized(self) {
+        if (!_migrationDBConnection) {
+            _migrationDBConnection = self.primaryStorage.newDatabaseConnection;
+        }
+        return _migrationDBConnection;
+    }
+}
+
+- (YapDatabaseConnection *)analyticsDBConnection {
+    @synchronized(self) {
+        if (!_analyticsDBConnection) {
+            _analyticsDBConnection = self.primaryStorage.newDatabaseConnection;
+        }
+        return _analyticsDBConnection;
+    }
+}
+
 @end
 
 NS_ASSUME_NONNULL_END
diff --git a/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+SessionStore.m b/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+SessionStore.m
index ec4ff7413..86acee97f 100644
--- a/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+SessionStore.m
+++ b/SignalServiceKit/src/Storage/AxolotlStore/OWSPrimaryStorage+SessionStore.m
@@ -2,8 +2,9 @@
 //  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
 //
 
-#import "OWSPrimaryStorage+SessionStore.h"
 #import "OWSFileSystem.h"
+#import "OWSPrimaryStorage+SessionStore.h"
+#import "SSKEnvironment.h"
 #import "YapDatabaseConnection+OWS.h"
 #import "YapDatabaseTransaction+OWS.h"
 #import <AxolotlKit/SessionRecord.h>
@@ -23,14 +24,7 @@ NSString *const kSessionStoreDBConnectionKey = @"kSessionStoreDBConnectionKey";
  */
 + (YapDatabaseConnection *)sessionStoreDBConnection
 {
-    static dispatch_once_t onceToken;
-    static YapDatabaseConnection *sessionStoreDBConnection;
-    dispatch_once(&onceToken, ^{
-        sessionStoreDBConnection = [OWSPrimaryStorage sharedManager].newDatabaseConnection;
-        sessionStoreDBConnection.objectCacheEnabled = NO;
-    });
-
-    return sessionStoreDBConnection;
+    return SSKEnvironment.shared.sessionStoreDBConnection;
 }
 
 - (YapDatabaseConnection *)sessionStoreDBConnection