Don't allow list media integrity check to race with backup uploads

Author: harry-signal

Signal/src/ViewControllers/AppSettings/Internal/InternalListMediaViewController.swift

@@ -14,11 +14,22 @@ class InternalListMediaViewController: OWSTableViewController2 {
         super.viewDidLoad()
 
         title = "Backup media debug"
-
-        updateTableContents()
     }
 
-    private var lastResult: ListMediaIntegrityCheckResult?
+    override func viewDidAppear(_ animated: Bool) {
+        super.viewDidAppear(animated)
+
+        ModalActivityIndicatorViewController.present(
+            fromViewController: self,
+            asyncBlock: { [weak self] modal in
+                try? await DependenciesBridge.shared.backupListMediaManager.queryListMediaIfNeeded()
+                await MainActor.run {
+                    self?.updateTableContents()
+                }
+                modal.dismiss(animated: true)
+            }
+        )
+    }
 
     func updateTableContents() {
         let contents = OWSTableContents()
@@ -28,6 +39,7 @@ class InternalListMediaViewController: OWSTableViewController2 {
             pendingUploadThumbnailCount,
             pendingOrphanDeleteCount,
             lastListMediaFailure,
+            lastListMediaResult,
         ) = DependenciesBridge.shared.db.read { tx in
             return (
                 try! QueuedBackupAttachmentUpload
@@ -37,7 +49,8 @@ class InternalListMediaViewController: OWSTableViewController2 {
                     .filter(Column(QueuedBackupAttachmentUpload.CodingKeys.isFullsize) == false)
                     .fetchCount(tx.database),
                 try! OrphanedBackupAttachment.fetchCount(tx.database),
-                try! DependenciesBridge.shared.backupListMediaManager.getLastFailingIntegrityCheckResult(tx: tx)
+                try! DependenciesBridge.shared.backupListMediaManager.getLastFailingIntegrityCheckResult(tx: tx),
+                try! DependenciesBridge.shared.backupListMediaManager.getMostRecentIntegrityCheckResult(tx: tx)
             )
         }
 
@@ -54,24 +67,23 @@ class InternalListMediaViewController: OWSTableViewController2 {
         contents.add(lastFailureSection)
 
         let lastResultSection = OWSTableSection(title: "Latest result")
-        if let lastResult {
-            Self.populate(section: lastResultSection, with: lastResult)
+        if let lastListMediaResult {
+            Self.populate(section: lastResultSection, with: lastListMediaResult)
         }
         lastResultSection.add(.actionItem(withText: "Perform remote integrity check", actionBlock: { [weak self] in
             guard let self else { return }
-            ModalActivityIndicatorViewController.present(
-                fromViewController: self,
-                asyncBlock: { [weak self] modal in
-                    guard let result = try? await DependenciesBridge.shared.backupListMediaManager.performUploadIntegrityCheck() else {
-                        return
-                    }
-                    await MainActor.run {
-                        self?.lastResult = result
-                        self?.updateTableContents()
-                    }
-                    modal.dismiss(animated: true)
-                }
+            let vc = ActionSheetController(
+                title: "This will schedule the integrity check to run on next app launch, then exit the app. "
+                    + "After tapping \"Okay\", please relaunch the app and return to this screen to check the results."
             )
+            vc.addAction(.init(title: "Okay", handler: { _ in
+                DependenciesBridge.shared.db.write { tx in
+                    DependenciesBridge.shared.backupListMediaManager.setManualNeedsListMedia(tx: tx)
+                }
+                exit(0)
+            }))
+            vc.addAction(.cancel)
+            present(vc, animated: true)
         }))
         contents.add(lastResultSection)
 

SignalServiceKit/Backups/Attachments/BackupListMediaManager.swift

@@ -79,7 +79,9 @@ public protocol BackupListMediaManager {
 
     func getLastFailingIntegrityCheckResult(tx: DBReadTransaction) throws -> ListMediaIntegrityCheckResult?
 
-    func performUploadIntegrityCheck() async throws -> ListMediaIntegrityCheckResult
+    func getMostRecentIntegrityCheckResult(tx: DBReadTransaction) throws -> ListMediaIntegrityCheckResult?
+
+    func setManualNeedsListMedia(tx: DBWriteTransaction)
 }
 
 public class BackupListMediaManagerImpl: BackupListMediaManager {
@@ -151,8 +153,17 @@ public class BackupListMediaManagerImpl: BackupListMediaManager {
         try kvStore.getCodableValue(forKey: Constants.lastNonEmptyIntegrityCheckResultKey, transaction: tx)
     }
 
+    public func getMostRecentIntegrityCheckResult(tx: DBReadTransaction) throws -> ListMediaIntegrityCheckResult? {
+        try kvStore.getCodableValue(forKey: Constants.lastIntegrityCheckResultKey, transaction: tx)
+    }
+
     private let taskQueue = ConcurrentTaskQueue(concurrentLimit: 1)
 
+    /// Nil if we have not run list media this app launch (only held in memory).
+    /// Set to the upload era where we ran list media this app launch.
+    /// Should only be accessed from the taskQueue for locking purposes.
+    private var lastListMediaUploadEraThisAppSession: String?
+
     public func queryListMediaIfNeeded() async throws {
         let task = Task {
             // Enqueue in a concurrent(1) task queue; we only want to run one of these at a time.
@@ -175,13 +186,7 @@ public class BackupListMediaManagerImpl: BackupListMediaManager {
         )
     }
 
-    public func performUploadIntegrityCheck() async throws -> ListMediaIntegrityCheckResult {
-        return try await taskQueue.run {
-            try await self._queryListMediaIfNeeded(forceIfNotNeeded: true)
-        }
-    }
-
-    private func _queryListMediaIfNeeded(forceIfNotNeeded: Bool = false) async throws -> ListMediaIntegrityCheckResult {
+    private func _queryListMediaIfNeeded() async throws -> ListMediaIntegrityCheckResult {
         guard FeatureFlags.Backups.supported else {
             return .empty(listMediaStartTimestamp: 0)
         }
@@ -216,7 +221,7 @@ public class BackupListMediaManagerImpl: BackupListMediaManager {
                 integrityCheckResult
             )
         }
-        guard needsToQuery || forceIfNotNeeded else {
+        guard needsToQuery else {
             return .empty(listMediaStartTimestamp: 0)
         }
 
@@ -1115,6 +1120,16 @@ public class BackupListMediaManagerImpl: BackupListMediaManager {
         case .free:
             return false
         case .paid, .paidExpiringSoon, .paidAsTester:
+            // Only query once per app session per upload era; this overrides the manual
+            // toggle and the date-based checks.
+            if lastListMediaUploadEraThisAppSession == currentUploadEra {
+                return false
+            }
+
+            if kvStore.getBool(Constants.manuallySetNeedsListMediaKey, defaultValue: false, transaction: tx) {
+                return true
+            }
+
             // If paid tier, query periodically as a catch-all to ensure local state
             // stays in sync with the server.
             let nowMs = dateProvider().ows_millisecondsSince1970
@@ -1157,16 +1172,19 @@ public class BackupListMediaManagerImpl: BackupListMediaManager {
         self.kvStore.setBool(true, key: Constants.hasEverRunListMediaKey, transaction: tx)
         if let uploadEra = kvStore.getString(Constants.inProgressUploadEraKey, transaction: tx) {
             self.kvStore.setString(uploadEra, key: Constants.lastListMediaUploadEraKey, transaction: tx)
+            self.lastListMediaUploadEraThisAppSession = uploadEra
             self.kvStore.removeValue(forKey: Constants.inProgressUploadEraKey, transaction: tx)
         } else {
             owsFailDebug("Missing in progress upload era?")
         }
         self.kvStore.setUInt64(startTimestamp, key: Constants.lastListMediaStartTimestampKey, transaction: tx)
+        self.kvStore.setBool(false, key: Constants.manuallySetNeedsListMediaKey, transaction: tx)
         kvStore.removeValue(forKey: Constants.inProgressListMediaStartTimestampKey, transaction: tx)
 
         if integrityCheckResult.hasFailures {
             try kvStore.setCodable(integrityCheckResult, key: Constants.lastNonEmptyIntegrityCheckResultKey, transaction: tx)
         }
+        try kvStore.setCodable(integrityCheckResult, key: Constants.lastIntegrityCheckResultKey, transaction: tx)
         kvStore.removeValue(forKey: Constants.inProgressIntegrityCheckResultKey, transaction: tx)
 
         self.kvStore.setBool(false, key: Constants.hasCompletedListingMediaKey, transaction: tx)
@@ -1184,10 +1202,15 @@ public class BackupListMediaManagerImpl: BackupListMediaManager {
             // Rotate the last integrity check failure when disabled
             db.write { tx in
                 kvStore.removeValue(forKey: Constants.lastNonEmptyIntegrityCheckResultKey, transaction: tx)
+                kvStore.removeValue(forKey: Constants.lastIntegrityCheckResultKey, transaction: tx)
             }
         }
     }
 
+    public func setManualNeedsListMedia(tx: DBWriteTransaction) {
+        kvStore.setBool(true, key: Constants.manuallySetNeedsListMediaKey, transaction: tx)
+    }
+
     private enum Constants {
         /// Maps to the upload era (active subscription) when we last queried the list media
         /// endpoint, or nil if its never been queried.
@@ -1197,6 +1220,8 @@ public class BackupListMediaManagerImpl: BackupListMediaManager {
         static let lastListMediaStartTimestampKey = "lastListMediaTimestamp"
         static let inProgressListMediaStartTimestampKey = "inProgressListMediaTimestamp"
 
+        static let manuallySetNeedsListMediaKey = "manuallySetNeedsListMediaKey"
+
         /// True if we've ever run list media in the lifetime of this app.
         static let hasEverRunListMediaKey = "hasEverRunListMedia"
 
@@ -1221,6 +1246,7 @@ public class BackupListMediaManagerImpl: BackupListMediaManager {
         static let hasCompletedEnumeratingAttachmentsKey = "hasCompletedEnumeratingAttachmentsKey"
 
         static let lastNonEmptyIntegrityCheckResultKey = "lastNonEmptyIntegrityCheckResultKey"
+        static let lastIntegrityCheckResultKey = "lastIntegrityCheckResultKey"
         static let inProgressIntegrityCheckResultKey = "inProgressIntegrityCheckResultKey"
     }
 }
@@ -1238,8 +1264,12 @@ class MockBackupListMediaManager: BackupListMediaManager {
         nil
     }
 
-    func performUploadIntegrityCheck() async throws -> ListMediaIntegrityCheckResult {
-        return .empty(listMediaStartTimestamp: 0)
+    func getMostRecentIntegrityCheckResult(tx: DBReadTransaction) throws -> ListMediaIntegrityCheckResult? {
+        nil
+    }
+
+    func setManualNeedsListMedia(tx: DBWriteTransaction) {
+        // Nothing
     }
 }