Add local cache of backup fragment metadata.

Signal/src/ViewControllers/DebugUI/DebugUIBackup.m

@@ -48,6 +48,10 @@ NS_ASSUME_NONNULL_BEGIN
                                      actionBlock:^{
                                          [DebugUIBackup clearAllCloudKitRecords];
                                      }]];
+    [items addObject:[OWSTableItem itemWithTitle:@"Clear Backup Metadata Cache"
+                                     actionBlock:^{
+                                         [DebugUIBackup clearBackupMetadataCache];
+                                     }]];
 
     return [OWSTableSection sectionWithTitle:self.name items:items];
 }
@@ -161,6 +165,16 @@ NS_ASSUME_NONNULL_BEGIN
     [OWSBackup.sharedManager clearAllCloudKitRecords];
 }
 
++ (void)clearBackupMetadataCache
+{
+    DDLogInfo(@"%@ ClearBackupMetadataCache.", self.logTag);
+
+    [OWSPrimaryStorage.sharedManager.newDatabaseConnection
+        readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
+            [transaction removeAllObjectsInCollection:[OWSBackupManifestItem collection]];
+        }];
+}
+
 @end
 
 NS_ASSUME_NONNULL_END

Signal/src/util/OWSBackup.m

@@ -507,14 +507,14 @@ NS_ASSUME_NONNULL_BEGIN
                                          return;
                                      }
                                      TSAttachmentStream *attachmentStream = object;
-                                     if (attachmentStream.backupRestoreRecordName.length < 1) {
+                                     if (!attachmentStream.backupRestoreMetadata) {
                                          OWSProdLogAndFail(@"%@ Invalid object: %@ in collection:%@",
                                              self.logTag,
                                              [object class],
                                              collection);
                                          return;
                                      }
-                                     [recordNames addObject:attachmentStream.backupRestoreRecordName];
+                                     [recordNames addObject:attachmentStream.backupRestoreMetadata.recordName];
                                  }];
     }];
     return recordNames;
@@ -556,9 +556,13 @@ NS_ASSUME_NONNULL_BEGIN
         return completion(NO);
     }
 
-    NSString *_Nullable recordName = attachment.backupRestoreRecordName;
+    OWSBackupManifestItem *_Nullable backupRestoreMetadata = attachment.backupRestoreMetadata;
-    NSData *_Nullable encryptionKey = attachment.backupRestoreEncryptionKey;
+    if (!backupRestoreMetadata) {
-    if (recordName.length < 1 || encryptionKey.length < 1) {
+        DDLogWarn(@"%@ Attachment missing lazy restore metadata.", self.logTag);
+        return completion(NO);
+    }
+    if (backupRestoreMetadata.recordName.length < 1 || backupRestoreMetadata.encryptionKey.length < 1) {
+        DDLogError(@"%@ Incomplete attachment metadata.", self.logTag);
         return completion(NO);
     }
 
@@ -571,14 +575,14 @@ NS_ASSUME_NONNULL_BEGIN
         return completion(NO);
     }
 
-    [OWSBackupAPI downloadFileFromCloudWithRecordName:recordName
+    [OWSBackupAPI downloadFileFromCloudWithRecordName:backupRestoreMetadata.recordName
         toFileUrl:[NSURL fileURLWithPath:tempFilePath]
         success:^{
             dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
                 [self lazyRestoreAttachment:attachment
                                    backupIO:backupIO
                           encryptedFilePath:tempFilePath
-                              encryptionKey:encryptionKey
+                              encryptionKey:backupRestoreMetadata.encryptionKey
                                  completion:completion];
             });
         }

Signal/src/util/OWSBackupExportJob.m

@@ -295,8 +295,7 @@ NS_ASSUME_NONNULL_BEGIN
 @property (nonatomic, nullable) OWSBackupExportItem *manifestItem;
 
 // If we are replacing an existing backup, we use some of its contents for continuity.
-@property (nonatomic, nullable) NSDictionary<NSString *, OWSBackupManifestItem *> *lastManifestItemMap;
+@property (nonatomic, nullable) NSSet<NSString *> *lastValidRecordNames;
-@property (nonatomic, nullable) NSSet<NSString *> *lastRecordNames;
 
 @end
 
@@ -346,7 +345,7 @@ NS_ASSUME_NONNULL_BEGIN
         if (self.isComplete) {
             return;
         }
-        [self tryToFetchManifestWithCompletion:^(BOOL tryToFetchManifestSuccess) {
+        [self fetchAllRecordsWithCompletion:^(BOOL tryToFetchManifestSuccess) {
             if (!tryToFetchManifestSuccess) {
                 [self failWithErrorDescription:
                           NSLocalizedString(@"BACKUP_EXPORT_ERROR_COULD_NOT_EXPORT",
@@ -374,7 +373,7 @@ NS_ASSUME_NONNULL_BEGIN
                     [weakSelf failWithError:saveError];
                     return;
                 }
-                [self cleanUpCloudWithCompletion:^(NSError *_Nullable cleanUpError) {
+                [self cleanUpWithCompletion:^(NSError *_Nullable cleanUpError) {
                     if (cleanUpError) {
                         [weakSelf failWithError:cleanUpError];
                         return;
@@ -422,69 +421,6 @@ NS_ASSUME_NONNULL_BEGIN
         }];
 }
 
-- (void)tryToFetchManifestWithCompletion:(OWSBackupJobBoolCompletion)completion
-{
-    OWSAssert(completion);
-
-    if (self.isComplete) {
-        return;
-    }
-
-    DDLogVerbose(@"%@ %s", self.logTag, __PRETTY_FUNCTION__);
-
-    [self updateProgressWithDescription:NSLocalizedString(@"BACKUP_IMPORT_PHASE_CHECK_BACKUP",
-                                            @"Indicates that the backup import is checking for an existing backup.")
-                               progress:nil];
-
-    __weak OWSBackupExportJob *weakSelf = self;
-
-    [OWSBackupAPI checkForManifestInCloudWithSuccess:^(BOOL value) {
-        dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
-            if (value) {
-                [weakSelf fetchManifestWithCompletion:completion];
-            } else {
-                // There is no existing manifest; continue.
-                completion(YES);
-            }
-        });
-    }
-        failure:^(NSError *error) {
-            dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
-                completion(NO);
-            });
-        }];
-}
-
-- (void)fetchManifestWithCompletion:(OWSBackupJobBoolCompletion)completion
-{
-    OWSAssert(completion);
-
-    if (self.isComplete) {
-        return;
-    }
-
-    DDLogVerbose(@"%@ %s", self.logTag, __PRETTY_FUNCTION__);
-
-    __weak OWSBackupExportJob *weakSelf = self;
-    [weakSelf downloadAndProcessManifestWithSuccess:^(OWSBackupManifestContents *manifest) {
-        OWSBackupExportJob *strongSelf = weakSelf;
-        if (!strongSelf) {
-            return;
-        }
-        if (strongSelf.isComplete) {
-            return;
-        }
-        OWSCAssert(manifest.databaseItems.count > 0);
-        OWSCAssert(manifest.attachmentsItems);
-        [strongSelf processLastManifest:manifest];
-        [strongSelf fetchAllRecordsWithCompletion:completion];
-    }
-        failure:^(NSError *manifestError) {
-            completion(NO);
-        }
-        backupIO:self.backupIO];
-}
-
 - (void)fetchAllRecordsWithCompletion:(OWSBackupJobBoolCompletion)completion
 {
     OWSAssert(completion);
@@ -505,7 +441,7 @@ NS_ASSUME_NONNULL_BEGIN
             if (strongSelf.isComplete) {
                 return;
             }
-            strongSelf.lastRecordNames = [NSSet setWithArray:recordNames];
+            strongSelf.lastValidRecordNames = [NSSet setWithArray:recordNames];
             completion(YES);
         });
     }
@@ -516,17 +452,6 @@ NS_ASSUME_NONNULL_BEGIN
         }];
 }
 
-- (void)processLastManifest:(OWSBackupManifestContents *)manifest
-{
-    OWSAssert(manifest);
-
-    NSMutableDictionary<NSString *, OWSBackupManifestItem *> *lastManifestItemMap = [NSMutableDictionary new];
-    for (OWSBackupManifestItem *manifestItem in manifest.attachmentsItems) {
-        lastManifestItemMap[manifestItem.recordName] = manifestItem;
-    }
-    self.lastManifestItemMap = [lastManifestItemMap copy];
-}
-
 - (BOOL)exportDatabase
 {
     OWSAssert(self.backupIO);
@@ -806,7 +731,7 @@ NS_ASSUME_NONNULL_BEGIN
     OWSAttachmentExport *attachmentExport = self.unsavedAttachmentExports.lastObject;
     [self.unsavedAttachmentExports removeLastObject];
 
-    if (self.lastManifestItemMap && self.lastRecordNames) {
+    if (self.lastValidRecordNames) {
         // Wherever possible, we do incremental backups and re-use fragments of the last backup.
         // Recycling fragments doesn't just reduce redundant network activity,
         // it allows us to skip the local export work, i.e. encryption.
@@ -818,8 +743,9 @@ NS_ASSUME_NONNULL_BEGIN
         //   this record's metadata.
         // * That this record does in fact exist in our CloudKit database.
         NSString *lastRecordName = [OWSBackupAPI recordNameForPersistentFileWithFileId:attachmentExport.attachmentId];
-        OWSBackupManifestItem *_Nullable lastManifestItem = self.lastManifestItemMap[lastRecordName];
+        OWSBackupManifestItem *_Nullable lastManifestItem =
-        if (lastManifestItem && [self.lastRecordNames containsObject:lastRecordName]) {
+            [OWSBackupManifestItem fetchObjectWithUniqueID:lastRecordName];
+        if (lastManifestItem && [self.lastValidRecordNames containsObject:lastRecordName]) {
             OWSAssert(lastManifestItem.encryptionKey.length > 0);
             OWSAssert(lastManifestItem.relativeFilePath.length > 0);
 
@@ -887,6 +813,14 @@ NS_ASSUME_NONNULL_BEGIN
                 exportItem.attachmentExport = attachmentExport;
                 [strongSelf.savedAttachmentItems addObject:exportItem];
 
+                // Immediately save the record metadata to facilitate export resume.
+                OWSBackupManifestItem *backupRestoreMetadata = [OWSBackupManifestItem new];
+                backupRestoreMetadata.recordName = recordName;
+                backupRestoreMetadata.encryptionKey = exportItem.encryptedItem.encryptionKey;
+                backupRestoreMetadata.relativeFilePath = attachmentExport.relativeFilePath;
+                backupRestoreMetadata.uncompressedDataLength = exportItem.uncompressedDataLength;
+                [backupRestoreMetadata save];
+
                 DDLogVerbose(@"%@ saved attachment: %@ as %@",
                     self.logTag,
                     attachmentExport.attachmentFilePath,
@@ -999,10 +933,15 @@ NS_ASSUME_NONNULL_BEGIN
     return result;
 }
 
-- (void)cleanUpCloudWithCompletion:(OWSBackupJobCompletion)completion
+- (void)cleanUpWithCompletion:(OWSBackupJobCompletion)completion
 {
     OWSAssert(completion);
 
+    if (self.isComplete) {
+        // Job was aborted.
+        return completion(nil);
+    }
+
     DDLogVerbose(@"%@ %s", self.logTag, __PRETTY_FUNCTION__);
 
     [self updateProgressWithDescription:NSLocalizedString(@"BACKUP_EXPORT_PHASE_CLEAN_UP",
@@ -1034,6 +973,44 @@ NS_ASSUME_NONNULL_BEGIN
     NSArray<NSString *> *restoringRecordNames = [OWSBackup.sharedManager attachmentRecordNamesForLazyRestore];
     [activeRecordNames addObjectsFromArray:restoringRecordNames];
 
+    [self cleanUpMetadataCacheWithActiveRecordNames:activeRecordNames];
+
+    [self cleanUpCloudWithActiveRecordNames:activeRecordNames completion:completion];
+}
+
+- (void)cleanUpMetadataCacheWithActiveRecordNames:(NSSet<NSString *> *)activeRecordNames
+{
+    OWSAssert(activeRecordNames.count > 0);
+
+    if (self.isComplete) {
+        // Job was aborted.
+        return;
+    }
+
+    // After every successful backup export, we can (and should) cull metadata
+    // for any backup fragment (i.e. CloudKit record) that wasn't involved in
+    // the latest backup export.
+    [self.primaryStorage.newDatabaseConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
+        NSMutableSet<NSString *> *obsoleteRecordNames = [NSMutableSet new];
+        [obsoleteRecordNames addObjectsFromArray:[transaction allKeysInCollection:[OWSBackupManifestItem collection]]];
+        [obsoleteRecordNames minusSet:activeRecordNames];
+
+        [transaction removeObjectsForKeys:obsoleteRecordNames.allObjects
+                             inCollection:[OWSBackupManifestItem collection]];
+    }];
+}
+
+- (void)cleanUpCloudWithActiveRecordNames:(NSSet<NSString *> *)activeRecordNames
+                               completion:(OWSBackupJobCompletion)completion
+{
+    OWSAssert(activeRecordNames.count > 0);
+    OWSAssert(completion);
+
+    if (self.isComplete) {
+        // Job was aborted.
+        return completion(nil);
+    }
+
     __weak OWSBackupExportJob *weakSelf = self;
     [OWSBackupAPI fetchAllRecordNamesWithSuccess:^(NSArray<NSString *> *recordNames) {
         // Ensure that we continue to work off the main thread.

Signal/src/util/OWSBackupImportJob.m

@@ -109,6 +109,13 @@ NSString *const kOWSBackup_ImportDatabaseKeySpec = @"kOWSBackup_ImportDatabaseKe
     [allItems addObjectsFromArray:self.databaseItems];
     [allItems addObjectsFromArray:self.attachmentsItems];
 
+    // Record metadata for all items, so that we can re-use them in incremental backups after the restore.
+    [self.primaryStorage.newDatabaseConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
+        for (OWSBackupManifestItem *item in allItems) {
+            [item saveWithTransaction:transaction];
+        }
+    }];
+
     __weak OWSBackupImportJob *weakSelf = self;
     [weakSelf
         downloadFilesFromCloud:allItems

Signal/src/util/OWSBackupJob.h

@@ -2,6 +2,9 @@
 //  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
 //
 
+#import "TSYapDatabaseObject.h"
+#import <SignalServiceKit/OWSBackupManifestItem.h>
+
 NS_ASSUME_NONNULL_BEGIN
 
 extern NSString *const kOWSBackup_ManifestKey_DatabaseFiles;
@@ -20,27 +23,6 @@ typedef void (^OWSBackupJobCompletion)(NSError *_Nullable error);
 typedef void (^OWSBackupJobManifestSuccess)(OWSBackupManifestContents *manifest);
 typedef void (^OWSBackupJobManifestFailure)(NSError *error);
 
-@interface OWSBackupManifestItem : NSObject
-
-@property (nonatomic) NSString *recordName;
-
-@property (nonatomic) NSData *encryptionKey;
-
-// This property is only set for certain types of manifest item,
-// namely attachments where we need to know where the attachment's
-// file should reside relative to the attachments folder.
-@property (nonatomic, nullable) NSString *relativeFilePath;
-
-// This property is only set if the manifest item is downloaded.
-@property (nonatomic, nullable) NSString *downloadFilePath;
-
-// This property is only set if the manifest item is compressed.
-@property (nonatomic, nullable) NSNumber *uncompressedDataLength;
-
-@end
-
-#pragma mark -
-
 @interface OWSBackupManifestContents : NSObject
 
 @property (nonatomic) NSArray<OWSBackupManifestItem *> *databaseItems;

Signal/src/util/OWSBackupJob.m

@@ -20,12 +20,6 @@ NSString *const kOWSBackup_ManifestKey_DataSize = @"data_size";
 
 NSString *const kOWSBackup_KeychainService = @"kOWSBackup_KeychainService";
 
-@implementation OWSBackupManifestItem
-
-@end
-
-#pragma mark -
-
 @implementation OWSBackupManifestContents
 
 @end

SignalServiceKit/src/Messages/Attachments/TSAttachmentStream.h

@@ -3,6 +3,7 @@
 //
 
 #import "DataSource.h"
+#import "OWSBackupManifestItem.h"
 #import "TSAttachment.h"
 
 #if TARGET_OS_IPHONE
@@ -34,10 +35,8 @@ NS_ASSUME_NONNULL_BEGIN
 
 @property (nonatomic, readonly) NSDate *creationTimestamp;
 
-// Optional properties.  Set only for attachments which
+// Optional property.  Only set for attachments which need "lazy backup restore."
-// need "lazy backup restore."
+@property (nonatomic, readonly, nullable) NSString *backupRestoreMetadataId;
-@property (nonatomic, readonly, nullable) NSString *backupRestoreRecordName;
-@property (nonatomic, readonly, nullable) NSData *backupRestoreEncryptionKey;
 
 #if TARGET_OS_IPHONE
 - (nullable UIImage *)image;
@@ -67,10 +66,12 @@ NS_ASSUME_NONNULL_BEGIN
 
 + (nullable NSError *)migrateToSharedData;
 
+- (nullable OWSBackupManifestItem *)backupRestoreMetadata;
+
 #pragma mark - Update With... Methods
 
 // Marks attachment as needing "lazy backup restore."
-- (void)updateWithBackupRestoreRecordName:(NSString *)recordName encryptionKey:(NSData *)encryptionKey;
+- (void)updateWithBackupRestoreMetadata:(OWSBackupManifestItem *)backupRestoreMetadata;
 // Marks attachment as having completed "lazy backup restore."
 - (void)updateWithBackupRestoreComplete;
 

SignalServiceKit/src/Messages/Attachments/TSAttachmentStream.m

@@ -26,8 +26,7 @@ NS_ASSUME_NONNULL_BEGIN
 // This property should only be accessed on the main thread.
 @property (nullable, nonatomic) NSNumber *cachedAudioDurationSeconds;
 
-@property (nonatomic, nullable) NSString *backupRestoreRecordName;
+@property (nonatomic, nullable) NSString *backupRestoreMetadataId;
-@property (nonatomic, nullable) NSData *backupRestoreEncryptionKey;
 
 @end
 
@@ -613,18 +612,30 @@ NS_ASSUME_NONNULL_BEGIN
     return audioDurationSeconds;
 }
 
+- (nullable OWSBackupManifestItem *)backupRestoreMetadata
+{
+    if (!self.backupRestoreMetadataId) {
+        return nil;
+    }
+    return [OWSBackupManifestItem fetchObjectWithUniqueID:self.backupRestoreMetadataId];
+}
+
 #pragma mark - Update With... Methods
 
-- (void)updateWithBackupRestoreRecordName:(NSString *)recordName encryptionKey:(NSData *)encryptionKey
+- (void)updateWithBackupRestoreMetadata:(OWSBackupManifestItem *)backupRestoreMetadata
 {
-    OWSAssert(recordName.length > 0);
+    OWSAssert(backupRestoreMetadata);
-    OWSAssert(encryptionKey.length > 0);
 
     [self.dbReadWriteConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
+        if (!backupRestoreMetadata.uniqueId) {
+            // If metadata hasn't been saved yet, save now.
+            [backupRestoreMetadata saveWithTransaction:transaction];
+
+            OWSAssert(backupRestoreMetadata.uniqueId);
+        }
         [self applyChangeToSelfAndLatestCopy:transaction
                                  changeBlock:^(TSAttachmentStream *attachment) {
-                                     [attachment setBackupRestoreRecordName:recordName];
+                                     [attachment setBackupRestoreMetadataId:backupRestoreMetadata.uniqueId];
-                                     [attachment setBackupRestoreEncryptionKey:encryptionKey];
                                  }];
     }];
 }
@@ -634,8 +645,7 @@ NS_ASSUME_NONNULL_BEGIN
     [self.dbReadWriteConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
         [self applyChangeToSelfAndLatestCopy:transaction
                                  changeBlock:^(TSAttachmentStream *attachment) {
-                                     [attachment setBackupRestoreRecordName:nil];
+                                     [attachment setBackupRestoreMetadataId:nil];
-                                     [attachment setBackupRestoreEncryptionKey:nil];
                                  }];
     }];
 }

SignalServiceKit/src/Storage/TSDatabaseView.m

@@ -354,8 +354,7 @@ NSString *const TSLazyRestoreAttachmentsGroup = @"TSLazyRestoreAttachmentsGroup"
             return nil;
         }
         TSAttachmentStream *attachmentStream = (TSAttachmentStream *)object;
-        if (attachmentStream.backupRestoreRecordName.length > 0
+        if (attachmentStream.backupRestoreMetadata) {
-            && attachmentStream.backupRestoreEncryptionKey.length > 0) {
             return TSLazyRestoreAttachmentsGroup;
         } else {
             return nil;

SignalServiceKit/src/Util/OWSBackupManifestItem.h

@@ -0,0 +1,39 @@
+//
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
+//
+
+#import "TSYapDatabaseObject.h"
+
+NS_ASSUME_NONNULL_BEGIN
+
+// We store metadata for known backup fragments (i.e.  CloudKit record) in
+// the database.  We might learn about them from:
+//
+// * Past backup exports.
+// * An import downloading and parsing the manifest of the last complete backup.
+//
+// Storing this data in the database provides continuity.
+//
+// * Backup exports can reuse fragments from previous Backup exports even if they
+//   don't complete (i.e. backup export resume).
+// * Backup exports can reuse fragments from the backup import, if any.
+@interface OWSBackupManifestItem : TSYapDatabaseObject
+
+@property (nonatomic) NSString *recordName;
+
+@property (nonatomic) NSData *encryptionKey;
+
+// This property is only set for certain types of manifest item,
+// namely attachments where we need to know where the attachment's
+// file should reside relative to the attachments folder.
+@property (nonatomic, nullable) NSString *relativeFilePath;
+
+// This property is only set if the manifest item is downloaded.
+@property (nonatomic, nullable) NSString *downloadFilePath;
+
+// This property is only set if the manifest item is compressed.
+@property (nonatomic, nullable) NSNumber *uncompressedDataLength;
+
+@end
+
+NS_ASSUME_NONNULL_END

SignalServiceKit/src/Util/OWSBackupManifestItem.m

@@ -0,0 +1,23 @@
+//
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
+//
+
+#import "OWSBackupManifestItem.h"
+
+NS_ASSUME_NONNULL_BEGIN
+
+@implementation OWSBackupManifestItem
+
+- (void)saveWithTransaction:(YapDatabaseReadWriteTransaction *)transaction
+{
+    OWSAssert(self.recordName.length > 0);
+
+    if (!self.uniqueId) {
+        self.uniqueId = self.recordName;
+    }
+    [super saveWithTransaction:transaction];
+}
+
+@end
+
+NS_ASSUME_NONNULL_END