fix the issue with cache

SignalServiceKit/src/Loki/API/LokiFileServerAPI.swift

@@ -24,8 +24,8 @@ public final class LokiFileServerAPI : LokiDotNetAPI {
 
     /// Gets the device links associated with the given hex encoded public key from the
     /// server and stores and returns the valid ones.
-    public static func getDeviceLinks(associatedWith hexEncodedPublicKey: String, in transaction: YapDatabaseReadWriteTransaction? = nil) -> Promise<Set<DeviceLink>> {
+    public static func getDeviceLinks(associatedWith hexEncodedPublicKey: String, usingCache: Bool = false) -> Promise<Set<DeviceLink>> {
-        return getDeviceLinks(associatedWith: [ hexEncodedPublicKey ], in: transaction)
+        return getDeviceLinks(associatedWith: [ hexEncodedPublicKey ], usingCache: usingCache)
     }
 
     @objc(getDeviceLinksAssociatedWithHexEncodedPublicKeys:)
@@ -35,7 +35,7 @@ public final class LokiFileServerAPI : LokiDotNetAPI {
     
     /// Gets the device links associated with the given hex encoded public keys from the
     /// server and stores and returns the valid ones.
-    public static func getDeviceLinks(associatedWith hexEncodedPublicKeys: Set<String>, in transaction: YapDatabaseReadWriteTransaction? = nil) -> Promise<Set<DeviceLink>> {
+    public static func getDeviceLinks(associatedWith hexEncodedPublicKeys: Set<String>, usingCache: Bool = false) -> Promise<Set<DeviceLink>> {
         let hexEncodedPublicKeysDescription = "[ \(hexEncodedPublicKeys.joined(separator: ", ")) ]"
         print("[Loki] Getting device links for: \(hexEncodedPublicKeysDescription).")
         return getAuthToken(for: server).then(on: DispatchQueue.global()) { token -> Promise<Set<DeviceLink>> in
@@ -86,19 +86,20 @@ public final class LokiFileServerAPI : LokiDotNetAPI {
                 })
             }.then(on: DispatchQueue.global()) { deviceLinks -> Promise<Set<DeviceLink>> in
                 let (promise, seal) = Promise<Set<DeviceLink>>.pending()
+                if usingCache {
+                    storage.setDeviceLinksInCache(deviceLinks)
+                    seal.fulfill(deviceLinks)
+                } else {
                     // Dispatch async on the main queue to avoid nested write transactions
                     DispatchQueue.main.async {
-                    if let trans = transaction {
-                        storage.setDeviceLinks(deviceLinks, in: trans)
-                    } else {
                         storage.dbReadWriteConnection.readWrite{ transaction in
                             storage.setDeviceLinks(deviceLinks, in: transaction)
                         }
-                    }
                         // We have to wait for the device links to be stored because a lot of our logic relies
                         // on them being in the database
                         seal.fulfill(deviceLinks)
                     }
+                }
                 return promise
             }
         }

SignalServiceKit/src/Loki/Database/OWSPrimaryStorage+Loki.swift

@@ -5,6 +5,17 @@ public extension OWSPrimaryStorage {
         return "LokiDeviceLinkCollection-\(masterHexEncodedPublicKey)"
     }
     
+    public func setDeviceLinksInCache(_ deviceLinks: Set<DeviceLink>) {
+        self.deviceLinkCache = deviceLinks
+    }
+    
+    public func syncDeviceLinkCacheToDatabase(in transaction: YapDatabaseReadWriteTransaction) {
+        if !self.deviceLinkCache.isEmpty {
+            self.setDeviceLinks(self.deviceLinkCache, in: transaction)
+            self.deviceLinkCache.removeAll()
+        }
+    }
+
     public func setDeviceLinks(_ deviceLinks: Set<DeviceLink>, in transaction: YapDatabaseReadWriteTransaction) {
         // TODO: Clear collections first?
         deviceLinks.forEach { addDeviceLink($0, in: transaction) } // TODO: Check the performance impact of this
@@ -21,6 +32,9 @@ public extension OWSPrimaryStorage {
     }
     
     public func getDeviceLinks(for masterHexEncodedPublicKey: String, in transaction: YapDatabaseReadTransaction) -> Set<DeviceLink> {
+        if !self.deviceLinkCache.isEmpty {
+            return self.deviceLinkCache
+        }
         let collection = getDeviceLinkCollection(for: masterHexEncodedPublicKey)
         guard !transaction.allKeys(inCollection: collection).isEmpty else { return [] } // Fixes a crash that used to occur on Josh's device
         var result: Set<DeviceLink> = []

SignalServiceKit/src/Loki/Protocol/Multi Device/MultiDeviceProtocol.swift

@@ -144,12 +144,17 @@ public final class MultiDeviceProtocol : NSObject {
         }
     }
 
-    @objc(updateDeviceLinksIfNeededForHexEncodedPublicKey:in:)
+    @objc(updateDeviceLinksIfNeededForHexEncodedPublicKey:in:usingCache:)
-    public static func updateDeviceLinksIfNeeded(for hexEncodedPublicKey: String, in transaction: YapDatabaseReadTransaction) -> AnyPromise {
+    public static func updateDeviceLinksIfNeeded(for hexEncodedPublicKey: String, in transaction: YapDatabaseReadTransaction, usingCache: Bool = false) -> AnyPromise {
-        let promise = getMultiDeviceDestinations(for: hexEncodedPublicKey, in: transaction)
+        let promise = getMultiDeviceDestinations(for: hexEncodedPublicKey, in: transaction, usingCache: usingCache)
         return AnyPromise.from(promise)
     }
     
+    @objc(syncDeviceLinkCacheToDatabaseInTransaction:)
+    public static func syncDeviceLinkCacheToDatabase(in transaction: YapDatabaseReadWriteTransaction) {
+        storage.syncDeviceLinkCacheToDatabase(in: transaction)
+    }
+
     /// See [Auto-Generated Friend Requests](https://github.com/loki-project/session-protocol-docs/wiki/Auto-Generated-Friend-Requests) for more information.
     @objc(getAutoGeneratedMultiDeviceFRMessageForHexEncodedPublicKey:in:)
     public static func getAutoGeneratedMultiDeviceFRMessage(for hexEncodedPublicKey: String, in transaction: YapDatabaseReadWriteTransaction) -> FriendRequestMessage {
@@ -284,7 +289,7 @@ public final class MultiDeviceProtocol : NSObject {
 // Here (in a non-@objc extension) because it doesn't interoperate well with Obj-C
 public extension MultiDeviceProtocol {
 
-    fileprivate static func getMultiDeviceDestinations(for hexEncodedPublicKey: String, in transaction: YapDatabaseReadTransaction) -> Promise<Set<MultiDeviceDestination>> {
+    fileprivate static func getMultiDeviceDestinations(for hexEncodedPublicKey: String, in transaction: YapDatabaseReadTransaction, usingCache: Bool = false) -> Promise<Set<MultiDeviceDestination>> {
         let (promise, seal) = Promise<Set<MultiDeviceDestination>>.pending()
         func getDestinations(in transaction: YapDatabaseReadTransaction? = nil) {
             storage.dbReadConnection.read { transaction in
@@ -306,7 +311,7 @@ public extension MultiDeviceProtocol {
         }
         if timeSinceLastUpdate > deviceLinkUpdateInterval {
             let masterHexEncodedPublicKey = storage.getMasterHexEncodedPublicKey(for: hexEncodedPublicKey, in: transaction) ?? hexEncodedPublicKey
-            LokiFileServerAPI.getDeviceLinks(associatedWith: masterHexEncodedPublicKey, in: transaction as! YapDatabaseReadWriteTransaction).done(on: DispatchQueue.global()) { _ in
+            LokiFileServerAPI.getDeviceLinks(associatedWith: masterHexEncodedPublicKey, usingCache: usingCache).done(on: DispatchQueue.global()) { _ in
                 getDestinations()
                 lastDeviceLinkUpdate[hexEncodedPublicKey] = Date()
             }.catch(on: DispatchQueue.global()) { error in

SignalServiceKit/src/Messages/OWSMessageManager.m

@@ -1266,17 +1266,18 @@ NS_ASSUME_NONNULL_BEGIN
     // Loki: Update device links in a blocking way
     // FIXME: This is horrible for performance
     // FIXME: ========
-    // FIX: Using the same transaction for write to avoid deadlock
+    // FIX: Using the cache for write to avoid deadlock
     // The envelope source is set during UD decryption
     if ([ECKeyPair isValidHexEncodedPublicKeyWithCandidate:envelope.source] && dataMessage.publicChatInfo == nil) { // Handled in LokiPublicChatPoller for open group messages
         dispatch_semaphore_t semaphore = dispatch_semaphore_create(0);
         dispatch_queue_t queue = dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0);
-        [[LKMultiDeviceProtocol updateDeviceLinksIfNeededForHexEncodedPublicKey:envelope.source in:transaction].ensureOn(queue, ^() {
+        [[LKMultiDeviceProtocol updateDeviceLinksIfNeededForHexEncodedPublicKey:envelope.source in:transaction usingCache: true].ensureOn(queue, ^() {
             dispatch_semaphore_signal(semaphore);
         }).catchOn(queue, ^(NSError *error) {
             dispatch_semaphore_signal(semaphore);
         }) retainUntilComplete];
         dispatch_semaphore_wait(semaphore, dispatch_time(DISPATCH_TIME_NOW, 10 * NSEC_PER_SEC));
+        [LKMultiDeviceProtocol syncDeviceLinkCacheToDatabaseInTransaction:transaction];
     }
     // FIXME: ========
 

SignalServiceKit/src/Storage/OWSPrimaryStorage.h

@@ -12,6 +12,8 @@ extern NSString *const OWSUIDatabaseConnectionWillUpdateExternallyNotification;
 extern NSString *const OWSUIDatabaseConnectionDidUpdateExternallyNotification;
 extern NSString *const OWSUIDatabaseConnectionNotificationsKey;
 
+@class LKDeviceLink;
+
 @interface OWSPrimaryStorage : OWSStorage
 
 - (instancetype)init NS_UNAVAILABLE;
@@ -23,6 +25,7 @@ extern NSString *const OWSUIDatabaseConnectionNotificationsKey;
 @property (nonatomic, readonly) YapDatabaseConnection *uiDatabaseConnection;
 @property (nonatomic, readonly) YapDatabaseConnection *dbReadConnection;
 @property (nonatomic, readonly) YapDatabaseConnection *dbReadWriteConnection;
+@property (nonatomic) NSSet<LKDeviceLink*> *deviceLinkCache;
 
 - (void)updateUIDatabaseConnectionToLatest;
 

SignalServiceKit/src/Storage/OWSPrimaryStorage.m

@@ -76,6 +76,8 @@ void VerifyRegistrationsForPrimaryStorage(OWSStorage *storage)
     if (self) {
         [self loadDatabase];
         
+        _deviceLinkCache = [NSSet set];
+
         _dbReadPool = [[YapDatabaseConnectionPool alloc] initWithDatabase:self.database];
         _dbReadWriteConnection = [self newDatabaseConnection];
         _uiDatabaseConnection = [self newDatabaseConnection];