Track completed unit count for backup attachment upload progress

Author: harry-signal

SignalServiceKit/Backups/Attachments/BackupAttachmentUploadProgress.swift

@@ -82,6 +82,7 @@ public actor BackupAttachmentUploadProgressImpl: BackupAttachmentUploadProgress
         let queueSnapshot = try self.computeRemainingUnuploadedByteCount()
         let sink = OWSProgress.createSink(block)
         let source = await sink.addSource(withLabel: "", unitCount: queueSnapshot.totalByteCount)
+        source.incrementCompletedUnitCount(by: queueSnapshot.completedByteCount)
         let observer = Observer(
             queueSnapshot: queueSnapshot,
             sink: sink,
@@ -226,6 +227,7 @@ public actor BackupAttachmentUploadProgressImpl: BackupAttachmentUploadProgress
 
     fileprivate struct UploadQueueSnapshot {
         let totalByteCount: UInt64
+        let completedByteCount: UInt64
         // We want to ignore updates from uploads that were scheduled after
         // we started observing. Take advantage of sequential row ids by
         // ignoring updates from ids that came after initial setup.
@@ -234,24 +236,46 @@ public actor BackupAttachmentUploadProgressImpl: BackupAttachmentUploadProgress
 
     private nonisolated func computeRemainingUnuploadedByteCount() throws -> UploadQueueSnapshot {
         return try db.read { tx in
-            var totalByteCount: UInt64 = 0
+            var remainingByteCount: UInt64 = 0
+            var completedByteCount: UInt64 = 0
             var maxRowId: Int64?
 
-            let cursor = try QueuedBackupAttachmentUpload
+            var cursor = try QueuedBackupAttachmentUpload
+                .filter(
+                    Column(QueuedBackupAttachmentUpload.CodingKeys.state)
+                        == QueuedBackupAttachmentUpload.State.ready.rawValue
+                )
                 // Don't coun't thumbnails in progress
                 .filter(Column(QueuedBackupAttachmentUpload.CodingKeys.isFullsize) == true)
                 .fetchCursor(tx.database)
 
             while let uploadRecord = try cursor.next() {
-                totalByteCount += UInt64(uploadRecord.estimatedByteCount)
+                remainingByteCount += UInt64(uploadRecord.estimatedByteCount)
                 if let existingMaxRowId = maxRowId {
                     maxRowId = max(existingMaxRowId, uploadRecord.id!)
                 } else {
                     maxRowId = uploadRecord.id
                 }
             }
 
-            return UploadQueueSnapshot(totalByteCount: totalByteCount, maxRowId: maxRowId)
+            cursor = try QueuedBackupAttachmentUpload
+                .filter(
+                    Column(QueuedBackupAttachmentUpload.CodingKeys.state)
+                        == QueuedBackupAttachmentUpload.State.done.rawValue
+                )
+                // Don't coun't thumbnails in progress
+                .filter(Column(QueuedBackupAttachmentUpload.CodingKeys.isFullsize) == true)
+                .fetchCursor(tx.database)
+
+            while let uploadRecord = try cursor.next() {
+                completedByteCount += UInt64(uploadRecord.estimatedByteCount)
+            }
+
+            return UploadQueueSnapshot(
+                totalByteCount: remainingByteCount + completedByteCount,
+                completedByteCount: completedByteCount,
+                maxRowId: maxRowId
+            )
         }
     }
 }
@@ -277,6 +301,7 @@ open class BackupAttachmentUploadProgressMock: BackupAttachmentUploadProgress {
         return BackupAttachmentUploadProgressObserver(
             queueSnapshot: .init(
                 totalByteCount: 100,
+                completedByteCount: 0,
                 maxRowId: nil
             ),
             sink: sink,

SignalServiceKit/Backups/Attachments/BackupAttachmentUploadQueueRunner.swift

@@ -756,7 +756,9 @@ class BackupAttachmentUploadQueueRunnerImpl: BackupAttachmentUploadQueueRunner {
         }
 
         func removeRecord(_ record: TaskRecord, tx: DBWriteTransaction) throws {
-            try backupAttachmentUploadStore.removeQueuedUpload(
+            // We don't actually delete records when finishing; we just mark
+            // them done so we can still keep track of their byte count.
+            try backupAttachmentUploadStore.markUploadDone(
                 for: record.record.attachmentRowId,
                 fullsize: record.record.isFullsize,
                 tx: tx

SignalServiceKit/Backups/Attachments/BackupAttachmentUploadStore.swift

@@ -13,7 +13,7 @@ public protocol BackupAttachmentUploadStore {
     /// If the same attachment is already enqueued, updates it to the greater of the old and new owner's timestamp.
     ///
     /// Doesn't actually trigger an upload; callers must later call `fetchNextUpload`, complete the upload of
-    /// both the fullsize and thumbnail as needed, and then call `removeQueuedUpload` once finished.
+    /// both the fullsize and thumbnail as needed, and then call `markUploadDone` once finished.
     /// Note that the upload operation can (and will) be separately durably enqueued in AttachmentUploadQueue,
     /// that's fine and doesn't change how this queue works.
     func enqueue(
@@ -36,7 +36,7 @@ public protocol BackupAttachmentUploadStore {
     /// Remove the upload from the queue. Should be called once uploaded (or permanently failed).
     /// - returns the removed record, if any.
     @discardableResult
-    func removeQueuedUpload(
+    func markUploadDone(
         for attachmentId: Attachment.IDType,
         fullsize: Bool,
         tx: DBWriteTransaction
@@ -79,8 +79,12 @@ public class BackupAttachmentUploadStoreImpl: BackupAttachmentUploadStore {
             .fetchOne(db)
 
         if var existingRecord {
-            // Only update if the new one has higher priority; otherwise leave untouched.
-            if newRecord.highestPriorityOwnerType.isHigherPriority(than: existingRecord.highestPriorityOwnerType) {
+            // Only update if done or the new one has higher priority; otherwise leave untouched.
+            let shouldUpdate = switch existingRecord.state {
+            case .done: true
+            case .ready: newRecord.highestPriorityOwnerType.isHigherPriority(than: existingRecord.highestPriorityOwnerType)
+            }
+            if shouldUpdate {
                 existingRecord.highestPriorityOwnerType = newRecord.highestPriorityOwnerType
                 try existingRecord.update(db)
             }
@@ -104,26 +108,29 @@ public class BackupAttachmentUploadStoreImpl: BackupAttachmentUploadStore {
                 tx.database,
                 sql: """
                     SELECT * FROM \(QueuedBackupAttachmentUpload.databaseTableName)
-                    WHERE \(QueuedBackupAttachmentUpload.CodingKeys.isFullsize.rawValue) = ?
+                    WHERE
+                      \(QueuedBackupAttachmentUpload.CodingKeys.state.rawValue) = ?
+                      AND \(QueuedBackupAttachmentUpload.CodingKeys.isFullsize.rawValue) = ?
                     ORDER BY
                         \(QueuedBackupAttachmentUpload.CodingKeys.maxOwnerTimestamp.rawValue) DESC NULLS FIRST
                     LIMIT ?
                     """,
-                arguments: [isFullsize, count]
+                arguments: [QueuedBackupAttachmentUpload.State.ready.rawValue, isFullsize, count]
             )
     }
 
     @discardableResult
-    public func removeQueuedUpload(
+    public func markUploadDone(
         for attachmentId: Attachment.IDType,
         fullsize: Bool,
         tx: DBWriteTransaction
     ) throws -> QueuedBackupAttachmentUpload? {
-        let record = try QueuedBackupAttachmentUpload
+        var record = try QueuedBackupAttachmentUpload
             .filter(Column(QueuedBackupAttachmentUpload.CodingKeys.attachmentRowId) == attachmentId)
             .filter(Column(QueuedBackupAttachmentUpload.CodingKeys.isFullsize) == fullsize)
             .fetchOne(tx.database)
-        try record?.delete(tx.database)
+        record?.state = .done
+        try record?.update(tx.database)
         return record
     }
 }

SignalServiceKit/Backups/Attachments/BackupAttachmentUploadStoreTests.swift

@@ -186,7 +186,7 @@ class BackupAttachmentUploadStoreTests: XCTestCase {
             )
 
             dequeuedRecords = try store.fetchNextUploads(
-                count: UInt(timestamps.count - 1),
+                count: UInt(timestamps.count),
                 isFullsize: true,
                 tx: tx
             )
@@ -203,11 +203,12 @@ class BackupAttachmentUploadStoreTests: XCTestCase {
         })
 
         // We should have gotten entries in timestamp order
-        XCTAssertEqual(dequeuedTimestamps, Array(sortedTimestamps.prefix(sortedTimestamps.count - 1)))
+        XCTAssertEqual(dequeuedTimestamps, Array(sortedTimestamps.prefix(sortedTimestamps.count)))
 
         try db.write { tx in
-            try dequeuedRecords.forEach { record in
-                try store.removeQueuedUpload(
+            // Finish all but one
+            try dequeuedRecords.prefix(timestamps.count - 1).forEach { record in
+                try store.markUploadDone(
                     for: record.attachmentRowId,
                     fullsize: true,
                     tx: tx
@@ -216,9 +217,26 @@ class BackupAttachmentUploadStoreTests: XCTestCase {
         }
 
         try db.read { tx in
-            // all rows but one should be deleted.
+            // Since not all rows are done, they should all stick around.
+            let records = try QueuedBackupAttachmentUpload.fetchAll(tx.database)
+            XCTAssertEqual(5, records.count)
+            XCTAssertEqual(4, records.filter({ $0.state == .done }).count)
+            XCTAssertEqual(1, records.filter({ $0.state == .ready }).count)
+        }
+
+        try db.write { tx in
+            // Finish the last one
+            _ = try store.markUploadDone(
+                for: dequeuedRecords.last!.attachmentRowId,
+                fullsize: true,
+                tx: tx
+            )
+        }
+
+        try db.read { tx in
+            // all rows but one should now be deleted.
             XCTAssertEqual(
-                1,
+                0,
                 try QueuedBackupAttachmentUpload.fetchCount(tx.database)
             )
         }
@@ -285,7 +303,7 @@ class BackupAttachmentUploadStoreTests: XCTestCase {
 
         try db.write { tx in
             try dequeuedRecords.forEach { record in
-                try store.removeQueuedUpload(
+                try store.markUploadDone(
                     for: record.attachmentRowId,
                     fullsize: record.isFullsize,
                     tx: tx
@@ -294,11 +312,12 @@ class BackupAttachmentUploadStoreTests: XCTestCase {
         }
 
         try db.read { tx in
-            // all fullsize rows should be deleted.
-            XCTAssertEqual(
-                4,
-                try QueuedBackupAttachmentUpload.fetchCount(tx.database)
-            )
+            // All fullsize rows should be done
+            let records = try QueuedBackupAttachmentUpload.fetchAll(tx.database)
+            XCTAssertEqual(8, records.count)
+            XCTAssertEqual(4, records.filter(\.isFullsize.negated).count)
+            XCTAssertEqual(4, records.filter(\.isFullsize).count)
+            XCTAssertEqual(4, records.filter({ $0.isFullsize && $0.state == .done }).count)
         }
     }
 

SignalServiceKit/Backups/Attachments/BackupListMediaManager.swift

@@ -879,16 +879,16 @@ public class BackupListMediaManagerImpl: BackupListMediaManager {
         // Since we now know this is uploaded, we can go ahead and remove
         // from the upload queue if present.
         if
-            let removedRecord = try backupAttachmentUploadStore.removeQueuedUpload(
+            let finishedRecord = try backupAttachmentUploadStore.markUploadDone(
                 for: attachment.id,
                 fullsize: isThumbnail.negated,
                 tx: tx
             )
         {
-            if removedRecord.isFullsize {
+            if finishedRecord.isFullsize {
                 Task {
                     await backupAttachmentUploadProgress.didFinishUploadOfFullsizeAttachment(
-                        uploadRecord: removedRecord
+                        uploadRecord: finishedRecord
                     )
                 }
             }

SignalServiceKit/Backups/Attachments/QueuedBackupAttachmentUpload.swift

@@ -64,6 +64,8 @@ public struct QueuedBackupAttachmentUpload: Codable, FetchableRecord, MutablePer
     /// Minimum timestamp at which this upload can be retried (if it failed in the past)
     public var minRetryTimestamp: UInt64
 
+    public var state: State
+
     public enum OwnerType {
         case threadWallpaper
         /// Timestamp of the newest message that owns this attachment.
@@ -80,6 +82,11 @@ public struct QueuedBackupAttachmentUpload: Codable, FetchableRecord, MutablePer
         }
     }
 
+    public enum State: Int, Codable {
+        case ready = 0
+        case done = 1
+    }
+
     public init(
         id: Int64? = nil,
         attachmentRowId: Attachment.IDType,
@@ -88,6 +95,7 @@ public struct QueuedBackupAttachmentUpload: Codable, FetchableRecord, MutablePer
         estimatedByteCount: UInt32,
         numRetries: UInt32 = 0,
         minRetryTimestamp: UInt64 = 0,
+        state: State = .ready,
     ) {
         self.id = id
         self.attachmentRowId = attachmentRowId
@@ -96,6 +104,7 @@ public struct QueuedBackupAttachmentUpload: Codable, FetchableRecord, MutablePer
         self.estimatedByteCount = estimatedByteCount
         self.numRetries = numRetries
         self.minRetryTimestamp = minRetryTimestamp
+        self.state = state
     }
 
     // MARK: FetchableRecord
@@ -124,6 +133,7 @@ public struct QueuedBackupAttachmentUpload: Codable, FetchableRecord, MutablePer
         case estimatedByteCount
         case numRetries
         case minRetryTimestamp
+        case state
     }
 
     public init(from decoder: any Decoder) throws {
@@ -140,6 +150,7 @@ public struct QueuedBackupAttachmentUpload: Codable, FetchableRecord, MutablePer
         self.estimatedByteCount = try container.decode(UInt32.self, forKey: .estimatedByteCount)
         self.numRetries = try container.decode(UInt32.self, forKey: .numRetries)
         self.minRetryTimestamp = try container.decode(UInt64.self, forKey: .minRetryTimestamp)
+        self.state = try container.decode(State.self, forKey: .state)
     }
 
     public func encode(to encoder: any Encoder) throws {
@@ -156,5 +167,6 @@ public struct QueuedBackupAttachmentUpload: Codable, FetchableRecord, MutablePer
         try container.encode(estimatedByteCount, forKey: .estimatedByteCount)
         try container.encode(numRetries, forKey: .numRetries)
         try container.encode(minRetryTimestamp, forKey: .minRetryTimestamp)
+        try container.encode(state, forKey: .state)
     }
 }

SignalServiceKit/Resources/schema.sql

@@ -2295,6 +2295,7 @@ CREATE
                 ,"isFullsize" BOOLEAN NOT NULL
                 ,"numRetries" INTEGER NOT NULL DEFAULT 0
                 ,"minRetryTimestamp" INTEGER NOT NULL DEFAULT 0
+                ,"state" INTEGER DEFAULT 0
 )
 ;
 
@@ -2393,13 +2394,6 @@ CREATE
         )
 ;
 
-CREATE
-    INDEX "index_BackupAttachmentUploadQueue_on_isFullsize_maxOwnerTimestamp"
-        ON "BackupAttachmentUploadQueue"("isFullsize"
-    ,"maxOwnerTimestamp"
-)
-;
-
 CREATE
     TABLE
         IF NOT EXISTS "Poll" (
@@ -2467,3 +2461,32 @@ CREATE
         ON "PollVote"("optionId"
 )
 ;
+
+CREATE
+    INDEX "index_BackupAttachmentUploadQueue_on_state_isFullsize_maxOwnerTimestamp"
+        ON "BackupAttachmentUploadQueue"("state"
+    ,"isFullsize"
+    ,"maxOwnerTimestamp"
+)
+;
+
+CREATE
+    TRIGGER __BackupAttachmentUploadQueue_au AFTER UPDATE
+            OF state
+                ON BackupAttachmentUploadQueue BEGIN DELETE
+                FROM
+                    BackupAttachmentUploadQueue
+                WHERE
+                    state = 1
+                    AND NOT EXISTS (
+                        SELECT
+                                id
+                            FROM
+                                BackupAttachmentUploadQueue
+                            WHERE
+                                state = 0
+                    )
+;
+
+END
+;