Put interaction-index manipulations inside a Bencher

Author: sashaweiss-signal

.swiftlint.yml

@@ -8,6 +8,7 @@ excluded:
 - Signal/Emoji/EmojiWithSkinTones+String.swift
 disabled_rules:
 - block_based_kvo
+- closure_parameter_position
 - compiler_protocol_init
 - control_statement
 - cyclomatic_complexity

SignalServiceKit/Environment/AppSetup.swift

@@ -1081,6 +1081,7 @@ public class AppSetup {
             appVersion: appVersion,
             attachmentDownloadManager: attachmentDownloadManager,
             attachmentUploadManager: attachmentUploadManager,
+            avatarFetcher: messageBackupAvatarFetcher,
             backupAttachmentDownloadManager: backupAttachmentDownloadManager,
             backupAttachmentUploadManager: backupAttachmentUploadManager,
             backupRequestManager: messageBackupRequestManager,

SignalServiceKit/MessageBackup/MessageBackup+Bench.swift

@@ -16,7 +16,8 @@ extension MessageBackup {
         private let startDate: MonotonicDate
 
         private var totalFramesProcessed: UInt64 = 0
-        private var metrics = [FrameType: Metrics]()
+        private var preFrameMetrics = [PreFrameRestoreAction: Metrics]()
+        private var frameMetrics = [FrameType: Metrics]()
         private var postFrameMetrics = [PostFrameRestoreAction: Metrics]()
 
         init(
@@ -50,18 +51,6 @@ extension MessageBackup {
             return try block(frameBencher)
         }
 
-        /// Measures the clock time spent in the provided block.
-        func benchPostFrameAction(_ action: PostFrameRestoreAction, _ block: () throws -> Void) rethrows -> Void {
-            let startDate = dateProvider()
-            try block()
-            let durationMs = dateProvider().millisSince(startDate)
-            var metrics = postFrameMetrics[action] ?? Metrics()
-            metrics.frameCount += 1
-            metrics.totalDurationMs += durationMs
-            metrics.maxDurationMs = max(durationMs, metrics.maxDurationMs)
-            postFrameMetrics[action] = metrics
-        }
-
         /// Given a block that does an enumeration over db objects, wraps that enumeration to instead take
         /// a closure with a FrameBencher that also measures the time spent enumerating.
         func wrapEnumeration<Input, T, Output>(
@@ -102,6 +91,33 @@ extension MessageBackup {
             }
         }
 
+        func benchPreFrameAction<T>(_ action: PreFrameRestoreAction, _ block: () throws -> T) rethrows -> T {
+            return try benchAction(action, actionMetricsKeyPath: \.preFrameMetrics, block: block)
+        }
+
+        func benchPostFrameAction<T>(_ action: PostFrameRestoreAction, _ block: () throws -> T) rethrows -> T {
+            return try benchAction(action, actionMetricsKeyPath: \.postFrameMetrics, block: block)
+        }
+
+        /// Measures the clock time spent in the provided block.
+        private func benchAction<Action: Hashable, T>(
+            _ action: Action,
+            actionMetricsKeyPath: ReferenceWritableKeyPath<Bencher, [Action: Metrics]>,
+            block: () throws -> T
+        ) rethrows -> T {
+            let startDate = dateProvider()
+            let result = try block()
+            let durationMs = dateProvider().millisSince(startDate)
+
+            var metrics = self[keyPath: actionMetricsKeyPath][action] ?? Metrics()
+            metrics.frameCount += 1
+            metrics.totalDurationMs += durationMs
+            metrics.maxDurationMs = max(durationMs, metrics.maxDurationMs)
+            self[keyPath: actionMetricsKeyPath][action] = metrics
+
+            return result
+        }
+
         class DBFileSizeBencher {
             private let dateProvider: DateProviderMonotonic
             private let dbFileSizeProvider: DBFileSizeProvider
@@ -165,7 +181,7 @@ extension MessageBackup {
                 let durationMs = bencher.dateProvider().millisSince(startDate)
                 bencher.totalFramesProcessed += 1
 
-                var metrics = bencher.metrics[frameType] ?? Metrics()
+                var metrics = bencher.frameMetrics[frameType] ?? Metrics()
                 metrics.frameCount += 1
                 metrics.totalDurationMs += durationMs
                 metrics.maxDurationMs = max(durationMs, metrics.maxDurationMs)
@@ -182,12 +198,12 @@ extension MessageBackup {
                     metrics.totalEnumerationDurationMs += startDate.millisSince(beforeEnumerationStartDate)
                 }
 
-                bencher.metrics[frameType] = metrics
+                bencher.frameMetrics[frameType] = metrics
             }
         }
 
         func logResults() {
-            let totalFrameCount = metrics.reduce(0, { $0 + $1.value.frameCount })
+            let totalFrameCount = frameMetrics.reduce(0, { $0 + $1.value.frameCount })
             Logger.info("Processed \(loggableCountString(totalFrameCount)) frames in \(dateProvider().millisSince(startDate))ms")
 
             func logMetrics(_ metrics: Metrics, typeString: String) {
@@ -212,9 +228,17 @@ extension MessageBackup {
                 Logger.info(logString)
             }
 
-            for (frameType, metrics) in self.metrics.sorted(by: { $0.value.totalDurationMs > $1.value.totalDurationMs }) {
+            Logger.info("Pre-Frame Metrics:")
+            for (action, metrics) in self.preFrameMetrics.sorted(by: { $0.value.totalDurationMs > $1.value.totalDurationMs }) {
+                logMetrics(metrics, typeString: action.rawValue)
+            }
+
+            Logger.info("Frame Metrics:")
+            for (frameType, metrics) in self.frameMetrics.sorted(by: { $0.value.totalDurationMs > $1.value.totalDurationMs }) {
                 logMetrics(metrics, typeString: frameType.rawValue)
             }
+
+            Logger.info("Post-Frame Metrics:")
             for (action, metrics) in self.postFrameMetrics.sorted(by: { $0.value.totalDurationMs > $1.value.totalDurationMs }) {
                 logMetrics(metrics, typeString: action.rawValue)
             }
@@ -280,7 +304,7 @@ extension MessageBackup {
             }
         }
 
-        private enum FrameType: String, CaseIterable {
+        private enum FrameType: String {
             case AccountData
 
             case Recipient_Contact
@@ -414,12 +438,17 @@ extension MessageBackup {
             }
         }
 
-        enum PostFrameRestoreAction: String, CaseIterable {
+        enum PreFrameRestoreAction: String {
+            case DropInteractionIndexes
+        }
+
+        enum PostFrameRestoreAction: String {
             case InsertContactHiddenInfoMessage
             case InsertPhoneNumberMissingAci
             case UpdateThreadMetadata
             case EnqueueAvatarFetch
             case IndexThreads
+            case RecreateInteractionIndexes
         }
     }
 }

SignalServiceKit/MessageBackup/MessageBackupManagerImpl.swift

@@ -36,6 +36,7 @@ public class MessageBackupManagerImpl: MessageBackupManager {
     private let appVersion: AppVersion
     private let attachmentDownloadManager: AttachmentDownloadManager
     private let attachmentUploadManager: AttachmentUploadManager
+    private let avatarFetcher: MessageBackupAvatarFetcher
     private let backupAttachmentDownloadManager: BackupAttachmentDownloadManager
     private let backupAttachmentUploadManager: BackupAttachmentUploadManager
     private let backupRequestManager: MessageBackupRequestManager
@@ -72,6 +73,7 @@ public class MessageBackupManagerImpl: MessageBackupManager {
         appVersion: AppVersion,
         attachmentDownloadManager: AttachmentDownloadManager,
         attachmentUploadManager: AttachmentUploadManager,
+        avatarFetcher: MessageBackupAvatarFetcher,
         backupAttachmentDownloadManager: BackupAttachmentDownloadManager,
         backupAttachmentUploadManager: BackupAttachmentUploadManager,
         backupRequestManager: MessageBackupRequestManager,
@@ -105,6 +107,7 @@ public class MessageBackupManagerImpl: MessageBackupManager {
         self.appVersion = appVersion
         self.attachmentDownloadManager = attachmentDownloadManager
         self.attachmentUploadManager = attachmentUploadManager
+        self.avatarFetcher = avatarFetcher
         self.backupAttachmentDownloadManager = backupAttachmentDownloadManager
         self.backupAttachmentUploadManager = backupAttachmentUploadManager
         self.backupRequestManager = backupRequestManager
@@ -675,14 +678,6 @@ public class MessageBackupManagerImpl: MessageBackupManager {
     ) async throws {
         let result: Result<BackupProto_BackupInfo, Error> = await db.awaitableWriteWithTxCompletion { tx in
             do {
-                /// Drops all indexes on the `TSInteraction` table before doing
-                /// the import, which dramatically speeds up the import. We'll
-                /// then recreate all these indexes in bulk afterwards.
-                let interactionIndexes = try dropAllIndexes(
-                    forTable: InteractionRecord.databaseTableName,
-                    tx: tx
-                )
-
                 let backupInfo = try Bench(
                     title: benchTitle,
                     memorySamplerRatio: FeatureFlags.messageBackupMemorySamplerRatio,
@@ -709,19 +704,6 @@ public class MessageBackupManagerImpl: MessageBackupManager {
                     }
                 }
 
-                let timeBeforeCreatingIndexes = dateProviderMonotonic()
-
-                /// Now that we've imported successfully, we want to recreate
-                /// the indexes we temporarily dropped.
-                try createIndexes(
-                    interactionIndexes,
-                    onTable: InteractionRecord.databaseTableName,
-                    tx: tx
-                )
-
-                let timeAfterCreatingIndexes = dateProviderMonotonic()
-                Logger.info("Created indexes in \(timeAfterCreatingIndexes.millisSince(timeBeforeCreatingIndexes))ms")
-
                 return .commit(.success(backupInfo))
             } catch let error {
                 return .rollback(.failure(error))
@@ -750,6 +732,16 @@ public class MessageBackupManagerImpl: MessageBackupManager {
             throw OWSAssertionError("Restoring from backup twice!")
         }
 
+        // Drops all indexes on the `TSInteraction` table before doing the
+        // import, which dramatically speeds up the import. We'll then recreate
+        // all these indexes in bulk afterwards.
+        let interactionIndexes = try bencher.benchPreFrameAction(.DropInteractionIndexes) {
+            try dropAllIndexes(
+                forTable: InteractionRecord.databaseTableName,
+                tx: tx
+            )
+        }
+
         var frameErrors = [LoggableErrorAndProto]()
         let result = Result<BackupProto_BackupInfo, Error>(catching: {
 
@@ -1042,26 +1034,50 @@ public class MessageBackupManagerImpl: MessageBackupManager {
 
             stream.closeFileStream()
 
-            /// Take any necessary post-frame-restore actions.
+            // Now that we've imported successfully, we want to recreate the
+            // the indexes we temporarily dropped.
+            try bencher.benchPostFrameAction(.RecreateInteractionIndexes) {
+                try createIndexes(
+                    interactionIndexes,
+                    onTable: InteractionRecord.databaseTableName,
+                    tx: tx
+                )
+            }
+
+            // Take any necessary post-frame-restore actions.
             try postFrameRestoreActionManager.performPostFrameRestoreActions(
                 recipientActions: contexts.recipient.postFrameRestoreActions,
                 chatActions: contexts.chat.postFrameRestoreActions,
                 bencher: bencher,
                 chatItemContext: contexts.chatItem
             )
 
-            // Index threads synchronously
+            // Index threads synchronously, since that should be fast.
             bencher.benchPostFrameAction(.IndexThreads) {
                 fullTextSearchIndexer.indexThreads(tx: tx)
             }
-            // Schedule message indexing asynchronously
+
+            // Schedule background message indexing, since that'll be slow.
             try fullTextSearchIndexer.scheduleMessagesJob(tx: tx)
 
-            tx.addAsyncCompletion(on: DispatchQueue.global()) { [backupAttachmentDownloadManager, disappearingMessagesJob] in
+            // Record that we've restored a Backup!
+            kvStore.setBool(true, key: Constants.keyValueStoreHasRestoredBackupKey, transaction: tx)
+
+            tx.addAsyncCompletion(on: DispatchQueue.global()) { [
+                avatarFetcher,
+                backupAttachmentDownloadManager,
+                disappearingMessagesJob
+            ] in
+                Task {
+                    // Kick off avatar fetches enqueued during restore.
+                    try await avatarFetcher.runIfNeeded()
+                }
+
                 Task {
-                    // Enqueue downloads for all the attachments.
+                    // Kick off attachment downloads enqueued during restore.
                     try await backupAttachmentDownloadManager.restoreAttachmentsIfNeeded()
                 }
+
                 // Start ticking down for disappearing messages.
                 disappearingMessagesJob.startIfNecessary()
             }
@@ -1071,10 +1087,9 @@ public class MessageBackupManagerImpl: MessageBackupManager {
             Logger.info("Backup first app version: \(backupInfo.firstAppVersion.nilIfEmpty ?? "Missing!")")
             bencher.logResults()
 
-            kvStore.setBool(true, key: Constants.keyValueStoreHasRestoredBackupKey, transaction: tx)
-
             return backupInfo
         })
+
         processErrors(errors: frameErrors, didFail: result.isSuccess.negated, tx: tx)
         return try result.get()
     }