diff --git a/SignalServiceKit/src/Loki/API/LokiFileServerAPI.swift b/SignalServiceKit/src/Loki/API/LokiFileServerAPI.swift index c5cd758df..1aed913e2 100644 --- a/SignalServiceKit/src/Loki/API/LokiFileServerAPI.swift +++ b/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> { - return getDeviceLinks(associatedWith: [ hexEncodedPublicKey ], in: transaction) + public static func getDeviceLinks(associatedWith hexEncodedPublicKey: String, usingCache: Bool = false) -> Promise> { + 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, in transaction: YapDatabaseReadWriteTransaction? = nil) -> Promise> { + public static func getDeviceLinks(associatedWith hexEncodedPublicKeys: Set, usingCache: Bool = false) -> Promise> { let hexEncodedPublicKeysDescription = "[ \(hexEncodedPublicKeys.joined(separator: ", ")) ]" print("[Loki] Getting device links for: \(hexEncodedPublicKeysDescription).") return getAuthToken(for: server).then(on: DispatchQueue.global()) { token -> Promise> in @@ -86,18 +86,19 @@ public final class LokiFileServerAPI : LokiDotNetAPI { }) }.then(on: DispatchQueue.global()) { deviceLinks -> Promise> in let (promise, seal) = Promise>.pending() - // 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 + if usingCache { + storage.setDeviceLinksInCache(deviceLinks) + seal.fulfill(deviceLinks) + } else { + // Dispatch async on the main queue to avoid nested write transactions + DispatchQueue.main.async { + 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) } - // 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 } diff --git a/SignalServiceKit/src/Loki/Database/OWSPrimaryStorage+Loki.swift b/SignalServiceKit/src/Loki/Database/OWSPrimaryStorage+Loki.swift index 6cc95236b..f63274b5b 100644 --- a/SignalServiceKit/src/Loki/Database/OWSPrimaryStorage+Loki.swift +++ b/SignalServiceKit/src/Loki/Database/OWSPrimaryStorage+Loki.swift @@ -4,6 +4,17 @@ public extension OWSPrimaryStorage { private func getDeviceLinkCollection(for masterHexEncodedPublicKey: String) -> String { return "LokiDeviceLinkCollection-\(masterHexEncodedPublicKey)" } + + public func setDeviceLinksInCache(_ deviceLinks: Set) { + 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, in transaction: YapDatabaseReadWriteTransaction) { // TODO: Clear collections first? @@ -21,6 +32,9 @@ public extension OWSPrimaryStorage { } public func getDeviceLinks(for masterHexEncodedPublicKey: String, in transaction: YapDatabaseReadTransaction) -> Set { + 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 = [] diff --git a/SignalServiceKit/src/Loki/Protocol/Multi Device/MultiDeviceProtocol.swift b/SignalServiceKit/src/Loki/Protocol/Multi Device/MultiDeviceProtocol.swift index 6b22f113a..39234c2b0 100644 --- a/SignalServiceKit/src/Loki/Protocol/Multi Device/MultiDeviceProtocol.swift +++ b/SignalServiceKit/src/Loki/Protocol/Multi Device/MultiDeviceProtocol.swift @@ -144,11 +144,16 @@ public final class MultiDeviceProtocol : NSObject { } } - @objc(updateDeviceLinksIfNeededForHexEncodedPublicKey:in:) - public static func updateDeviceLinksIfNeeded(for hexEncodedPublicKey: String, in transaction: YapDatabaseReadTransaction) -> AnyPromise { - let promise = getMultiDeviceDestinations(for: hexEncodedPublicKey, in: transaction) + @objc(updateDeviceLinksIfNeededForHexEncodedPublicKey:in:usingCache:) + public static func updateDeviceLinksIfNeeded(for hexEncodedPublicKey: String, in transaction: YapDatabaseReadTransaction, usingCache: Bool = false) -> AnyPromise { + 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:) @@ -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> { + fileprivate static func getMultiDeviceDestinations(for hexEncodedPublicKey: String, in transaction: YapDatabaseReadTransaction, usingCache: Bool = false) -> Promise> { let (promise, seal) = Promise>.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 diff --git a/SignalServiceKit/src/Messages/OWSMessageManager.m b/SignalServiceKit/src/Messages/OWSMessageManager.m index 501b1f94f..4302b62df 100644 --- a/SignalServiceKit/src/Messages/OWSMessageManager.m +++ b/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: ======== diff --git a/SignalServiceKit/src/Storage/OWSPrimaryStorage.h b/SignalServiceKit/src/Storage/OWSPrimaryStorage.h index dd1cf65b3..96d2bc6bf 100644 --- a/SignalServiceKit/src/Storage/OWSPrimaryStorage.h +++ b/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 *deviceLinkCache; - (void)updateUIDatabaseConnectionToLatest; diff --git a/SignalServiceKit/src/Storage/OWSPrimaryStorage.m b/SignalServiceKit/src/Storage/OWSPrimaryStorage.m index f05fcc8b9..a72507286 100644 --- a/SignalServiceKit/src/Storage/OWSPrimaryStorage.m +++ b/SignalServiceKit/src/Storage/OWSPrimaryStorage.m @@ -75,6 +75,8 @@ void VerifyRegistrationsForPrimaryStorage(OWSStorage *storage) if (self) { [self loadDatabase]; + + _deviceLinkCache = [NSSet set]; _dbReadPool = [[YapDatabaseConnectionPool alloc] initWithDatabase:self.database]; _dbReadWriteConnection = [self newDatabaseConnection];