From 0683cbe9529329ecd3c3805f86b40c4ff05349bf Mon Sep 17 00:00:00 2001
From: nielsandriesse <andriesseniels@gmail.com>
Date: Thu, 11 Jun 2020 08:20:31 +1000
Subject: [PATCH] Improve transaction handling

---
 .../Loki/Database/Storage+Collections.swift   | 19 ++++++
 .../src/Loki/Database/Storage.swift           | 59 ++++++++++++++-----
 2 files changed, 62 insertions(+), 16 deletions(-)
 create mode 100644 SignalServiceKit/src/Loki/Database/Storage+Collections.swift

diff --git a/SignalServiceKit/src/Loki/Database/Storage+Collections.swift b/SignalServiceKit/src/Loki/Database/Storage+Collections.swift
new file mode 100644
index 000000000..00b868328
--- /dev/null
+++ b/SignalServiceKit/src/Loki/Database/Storage+Collections.swift
@@ -0,0 +1,19 @@
+
+@objc public extension Storage {
+
+    // TODO: Add remaining collections
+
+    @objc func getDeviceLinkCollection(for masterPublicKey: String) -> String {
+        return "LokiDeviceLinkCollection-\(masterPublicKey)"
+    }
+
+    @objc public static func getSwarmCollection(for publicKey: String) -> String {
+        return "LokiSwarmCollection-\(publicKey)"
+    }
+
+    @objc public static let onionRequestPathCollection = "LokiOnionRequestPathCollection"
+    @objc public static let openGroupCollection = "LokiPublicChatCollection"
+    @objc public static let openGroupUserCountCollection = "LokiPublicChatUserCountCollection"
+    @objc public static let sessionRequestTimestampCollection = "LokiSessionRequestTimestampCollection"
+    @objc public static let snodePoolCollection = "LokiSnodePoolCollection"
+}
diff --git a/SignalServiceKit/src/Loki/Database/Storage.swift b/SignalServiceKit/src/Loki/Database/Storage.swift
index 4f3d702b2..4f27ddba5 100644
--- a/SignalServiceKit/src/Loki/Database/Storage.swift
+++ b/SignalServiceKit/src/Loki/Database/Storage.swift
@@ -1,17 +1,23 @@
+import PromiseKit
+
+// Some important notes about YapDatabase:
+//
+// • Connections are thread-safe.
+// • Executing a write transaction from within a write transaction is NOT allowed.
 
-/// Some important notes about YapDatabase:
-///
-/// • Connections are thread-safe.
-/// • Executing a write transaction from within a write transaction is NOT allowed.
 @objc(LKStorage)
 public final class Storage : NSObject {
+    private static let queue = DispatchQueue(label: "Storage.queue", qos: .userInitiated)
 
     private static var owsStorage: OWSPrimaryStorage { OWSPrimaryStorage.shared() }
 
-    /// Some important points regarding reading from the database:
-    ///
-    /// • Background threads should use `OWSPrimaryStorage`'s `dbReadPool`, whereas the main thread should use `OWSPrimaryStorage`'s `uiDatabaseConnection` (see the `YapDatabaseConnectionPool` documentation for more information).
-    /// • Multiple read transactions can safely be executed at the same time.
+    // MARK: Reading
+
+    // Some important points regarding reading from the database:
+    //
+    // • Background threads should use `OWSPrimaryStorage`'s `dbReadConnection`, whereas the main thread should use `OWSPrimaryStorage`'s `uiDatabaseConnection` (see the `YapDatabaseConnectionPool` documentation for more information).
+    // • Multiple read transactions can safely be executed at the same time.
+
     @objc(readWithBlock:)
     public static func read(with block: @escaping (YapDatabaseReadTransaction) -> Void) {
         let isMainThread = Thread.current.isMainThread
@@ -19,14 +25,35 @@ public final class Storage : NSObject {
         connection.read(block)
     }
 
-    /// Some important points regarding writing to the database:
-    ///
-    /// • There can only be a single write transaction per database at any one time, so all write transactions must use `OWSPrimaryStorage`'s `dbReadWriteConnection`.
-    /// • Executing a write transaction from within a write transaction causes a deadlock and must be avoided.
+    // MARK: Writing
+
+    // Some important points regarding writing to the database:
+    //
+    // • There can only be a single write transaction per database at any one time, so all write transactions must use `OWSPrimaryStorage`'s `dbReadWriteConnection`.
+    // • Executing a write transaction from within a write transaction causes a deadlock and must be avoided.
+
     @objc(writeWithBlock:)
-    public static func write(with block: @escaping (YapDatabaseReadWriteTransaction) -> Void) {
-        // TODO: Right now this is kind of pointless, but the idea is to eventually
-        // somehow manage nested write transactions in this class.
-        owsStorage.dbReadWriteConnection.readWrite(block)
+    public static func objc_write(with block: @escaping (YapDatabaseReadWriteTransaction) -> Void) -> AnyPromise {
+        return AnyPromise.from(write(with: block))
+    }
+
+    public static func write(with block: @escaping (YapDatabaseReadWriteTransaction) -> Void) -> Promise<Void> {
+        let (promise, seal) = Promise<Void>.pending()
+        queue.async { // TODO: There are cases where this isn't necessary
+            owsStorage.dbReadWriteConnection.readWrite(block)
+            seal.fulfill(())
+        }
+        return promise
+    }
+
+    /// Blocks the calling thread until the write has finished.
+    @objc(syncWriteWithBlock:error:)
+    public static func objc_syncWrite(with block: @escaping (YapDatabaseReadWriteTransaction) -> Void) throws {
+        try syncWrite(with: block)
+    }
+
+    /// Blocks the calling thread until the write has finished.
+    public static func syncWrite(with block: @escaping (YapDatabaseReadWriteTransaction) -> Void) throws {
+        try write(with: block).wait()
     }
 }