From bdde3c73c5fa75cd209ef29c9454e6a610bb42f7 Mon Sep 17 00:00:00 2001 From: Matthew Chen Date: Tue, 28 Mar 2017 17:47:43 -0400 Subject: [PATCH 1/2] Improve handling of incomplete and failed attachment downloads. // FREEBIE --- .../Attachments/OWSAttachmentsProcessor.m | 5 ++--- src/Messages/Attachments/TSAttachment.h | 2 ++ .../Attachments/TSAttachmentPointer.h | 9 +++++++-- .../Attachments/TSAttachmentPointer.m | 20 +++++++++++++++++-- src/Messages/OWSFailedMessagesJob.m | 2 ++ 5 files changed, 31 insertions(+), 7 deletions(-) diff --git a/src/Messages/Attachments/OWSAttachmentsProcessor.m b/src/Messages/Attachments/OWSAttachmentsProcessor.m index 3ca0748c6..720d32a10 100644 --- a/src/Messages/Attachments/OWSAttachmentsProcessor.m +++ b/src/Messages/Attachments/OWSAttachmentsProcessor.m @@ -331,7 +331,7 @@ NS_ASSUME_NONNULL_BEGIN - (void)setAttachment:(TSAttachmentPointer *)pointer isDownloadingInMessage:(nullable TSMessage *)message { - pointer.downloading = YES; + pointer.state = TSAttachmentPointerStateDownloading; [pointer save]; if (message) { [message touch]; @@ -340,8 +340,7 @@ NS_ASSUME_NONNULL_BEGIN - (void)setAttachment:(TSAttachmentPointer *)pointer didFailInMessage:(nullable TSMessage *)message { - pointer.downloading = NO; - pointer.failed = YES; + pointer.state = TSAttachmentPointerStateFailed; [pointer save]; if (message) { [message touch]; diff --git a/src/Messages/Attachments/TSAttachment.h b/src/Messages/Attachments/TSAttachment.h index 537ae0a33..9fd8d3f29 100644 --- a/src/Messages/Attachments/TSAttachment.h +++ b/src/Messages/Attachments/TSAttachment.h @@ -39,6 +39,8 @@ NS_ASSUME_NONNULL_BEGIN // that represent downloaded incoming attachments. - (instancetype)initWithPointer:(TSAttachment *)pointer; +- (nullable instancetype)initWithCoder:(NSCoder *)coder; + - (void)upgradeFromAttachmentSchemaVersion:(NSUInteger)attachmentSchemaVersion; @end diff --git a/src/Messages/Attachments/TSAttachmentPointer.h b/src/Messages/Attachments/TSAttachmentPointer.h index a6a825b5c..cb8143fdb 100644 --- a/src/Messages/Attachments/TSAttachmentPointer.h +++ b/src/Messages/Attachments/TSAttachmentPointer.h @@ -6,6 +6,12 @@ NS_ASSUME_NONNULL_BEGIN +typedef NS_ENUM(NSUInteger, TSAttachmentPointerState) { + TSAttachmentPointerStateEnqueued = 0, + TSAttachmentPointerStateDownloading = 1, + TSAttachmentPointerStateFailed = 2, +}; + /** * A TSAttachmentPointer is a yet-to-be-downloaded attachment. */ @@ -18,8 +24,7 @@ NS_ASSUME_NONNULL_BEGIN relay:(NSString *)relay NS_DESIGNATED_INITIALIZER; @property (nonatomic, readonly) NSString *relay; -@property (atomic, readwrite, getter=isDownloading) BOOL downloading; -@property (atomic, readwrite, getter=hasFailed) BOOL failed; +@property (atomic) TSAttachmentPointerState state; // Though now required, `digest` may be null for pre-existing records or from // messages received from other clients diff --git a/src/Messages/Attachments/TSAttachmentPointer.m b/src/Messages/Attachments/TSAttachmentPointer.m index 40a930830..73ca07a53 100644 --- a/src/Messages/Attachments/TSAttachmentPointer.m +++ b/src/Messages/Attachments/TSAttachmentPointer.m @@ -8,6 +8,23 @@ NS_ASSUME_NONNULL_BEGIN @implementation TSAttachmentPointer +- (nullable instancetype)initWithCoder:(NSCoder *)coder +{ + self = [super initWithCoder:coder]; + if (!self) { + return self; + } + + // A TSAttachmentPointer is a yet-to-be-downloaded attachment. + // If this is an old TSAttachmentPointer from another session, + // we know that it failed to complete before the session completed. + if (![coder containsValueForKey:@"state"]) { + _state = TSAttachmentPointerStateFailed; + } + + return self; +} + - (instancetype)initWithServerId:(UInt64)serverId key:(NSData *)key digest:(nullable NSData *)digest @@ -20,8 +37,7 @@ NS_ASSUME_NONNULL_BEGIN } _digest = digest; - _failed = NO; - _downloading = NO; + _state = TSAttachmentPointerStateEnqueued; _relay = relay; return self; diff --git a/src/Messages/OWSFailedMessagesJob.m b/src/Messages/OWSFailedMessagesJob.m index d787961e9..b296cb59d 100644 --- a/src/Messages/OWSFailedMessagesJob.m +++ b/src/Messages/OWSFailedMessagesJob.m @@ -21,6 +21,8 @@ static NSString *const OWSFailedMessagesJobMessageStateIndex = @"index_outoing_m @end +#pragma mark - + @implementation OWSFailedMessagesJob - (instancetype)initWithStorageManager:(TSStorageManager *)storageManager From 49a24a4e6a8056866233478ff99e1a7d310ef263 Mon Sep 17 00:00:00 2001 From: Matthew Chen Date: Tue, 28 Mar 2017 18:01:09 -0400 Subject: [PATCH 2/2] Improve handling of incomplete and failed attachment downloads. // FREEBIE --- .../OWSFailedAttachmentDownloadsJob.h | 28 ++++ .../OWSFailedAttachmentDownloadsJob.m | 148 ++++++++++++++++++ src/Storage/TSStorageManager.m | 4 + 3 files changed, 180 insertions(+) create mode 100644 src/Messages/OWSFailedAttachmentDownloadsJob.h create mode 100644 src/Messages/OWSFailedAttachmentDownloadsJob.m diff --git a/src/Messages/OWSFailedAttachmentDownloadsJob.h b/src/Messages/OWSFailedAttachmentDownloadsJob.h new file mode 100644 index 000000000..7ad6b20ee --- /dev/null +++ b/src/Messages/OWSFailedAttachmentDownloadsJob.h @@ -0,0 +1,28 @@ +// +// Copyright (c) 2017 Open Whisper Systems. All rights reserved. +// + +NS_ASSUME_NONNULL_BEGIN + +@class TSStorageManager; + +@interface OWSFailedAttachmentDownloadsJob : NSObject + +- (instancetype)init NS_UNAVAILABLE; +- (instancetype)initWithStorageManager:(TSStorageManager *)storageManager NS_DESIGNATED_INITIALIZER; + +- (void)run; + +/** + * Database extensions required for class to work. + */ +- (void)asyncRegisterDatabaseExtensions; + +/** + * Only use the sync version for testing, generally we'll want to register extensions async + */ +- (void)blockingRegisterDatabaseExtensions; + +@end + +NS_ASSUME_NONNULL_END diff --git a/src/Messages/OWSFailedAttachmentDownloadsJob.m b/src/Messages/OWSFailedAttachmentDownloadsJob.m new file mode 100644 index 000000000..ccd215513 --- /dev/null +++ b/src/Messages/OWSFailedAttachmentDownloadsJob.m @@ -0,0 +1,148 @@ +// +// Copyright (c) 2017 Open Whisper Systems. All rights reserved. +// + +#import "OWSFailedAttachmentDownloadsJob.h" +#import "TSAttachmentPointer.h" +#import "TSStorageManager.h" +#import +#import +#import + +NS_ASSUME_NONNULL_BEGIN + +static NSString *const OWSFailedAttachmentDownloadsJobAttachmentStateColumn = @"state"; +static NSString *const OWSFailedAttachmentDownloadsJobAttachmentStateIndex = @"index_attachment_downloads_on_state"; + +@interface OWSFailedAttachmentDownloadsJob () + +@property (nonatomic, readonly) TSStorageManager *storageManager; + +@end + +#pragma mark - + +@implementation OWSFailedAttachmentDownloadsJob + +- (instancetype)initWithStorageManager:(TSStorageManager *)storageManager +{ + self = [super init]; + if (!self) { + return self; + } + + _storageManager = storageManager; + + return self; +} + +- (NSArray *)fetchAttemptingOutAttachmentIds:(YapDatabaseConnection *)dbConnection +{ + NSMutableArray *attachmentIds = [NSMutableArray new]; + + NSString *formattedString = [NSString stringWithFormat:@"WHERE %@ != %d", + OWSFailedAttachmentDownloadsJobAttachmentStateColumn, + (int)TSAttachmentPointerStateFailed]; + YapDatabaseQuery *query = [YapDatabaseQuery queryWithFormat:formattedString]; + [dbConnection readWithBlock:^(YapDatabaseReadTransaction *_Nonnull transaction) { + [[transaction ext:OWSFailedAttachmentDownloadsJobAttachmentStateIndex] + enumerateKeysMatchingQuery:query + usingBlock:^void(NSString *collection, NSString *key, BOOL *stop) { + [attachmentIds addObject:key]; + }]; + }]; + + return [attachmentIds copy]; +} + +- (void)enumerateAttemptingOutAttachmentsWithBlock:(void (^_Nonnull)(TSAttachmentPointer *attachment))block +{ + YapDatabaseConnection *dbConnection = [self.storageManager newDatabaseConnection]; + + // Since we can't directly mutate the enumerated attachments, we store only their ids in hopes + // of saving a little memory and then enumerate the (larger) TSAttachment objects one at a time. + for (NSString *attachmentId in [self fetchAttemptingOutAttachmentIds:dbConnection]) { + TSAttachmentPointer *_Nullable attachment = [TSAttachmentPointer fetchObjectWithUniqueID:attachmentId]; + if ([attachment isKindOfClass:[TSAttachmentPointer class]]) { + block(attachment); + } else { + DDLogError(@"%@ unexpected object: %@", self.tag, attachment); + } + } +} + +- (void)run +{ + __block uint count = 0; + [self enumerateAttemptingOutAttachmentsWithBlock:^(TSAttachmentPointer *attachment) { + // sanity check + if (attachment.state != TSAttachmentPointerStateFailed) { + DDLogDebug(@"%@ marking attachment as failed", self.tag); + attachment.state = TSAttachmentPointerStateFailed; + [attachment save]; + count++; + } + }]; + + DDLogDebug(@"%@ Marked %u attachments as unsent", self.tag, count); +} + +#pragma mark - YapDatabaseExtension + +- (YapDatabaseSecondaryIndex *)indexDatabaseExtension +{ + YapDatabaseSecondaryIndexSetup *setup = [YapDatabaseSecondaryIndexSetup new]; + [setup addColumn:OWSFailedAttachmentDownloadsJobAttachmentStateColumn + withType:YapDatabaseSecondaryIndexTypeInteger]; + + YapDatabaseSecondaryIndexHandler *handler = + [YapDatabaseSecondaryIndexHandler withObjectBlock:^(YapDatabaseReadTransaction *transaction, + NSMutableDictionary *dict, + NSString *collection, + NSString *key, + id object) { + if (![object isKindOfClass:[TSAttachmentPointer class]]) { + return; + } + TSAttachmentPointer *attachment = (TSAttachmentPointer *)object; + dict[OWSFailedAttachmentDownloadsJobAttachmentStateColumn] = @(attachment.state); + }]; + + return [[YapDatabaseSecondaryIndex alloc] initWithSetup:setup handler:handler]; +} + +// Useful for tests, don't use in app startup path because it's slow. +- (void)blockingRegisterDatabaseExtensions +{ + [self.storageManager.database registerExtension:[self indexDatabaseExtension] + withName:OWSFailedAttachmentDownloadsJobAttachmentStateIndex]; +} + +- (void)asyncRegisterDatabaseExtensions +{ + [self.storageManager.database asyncRegisterExtension:[self indexDatabaseExtension] + withName:OWSFailedAttachmentDownloadsJobAttachmentStateIndex + completionBlock:^(BOOL ready) { + if (ready) { + DDLogDebug(@"%@ completed registering extension async.", self.tag); + } else { + DDLogError(@"%@ failed registering extension async.", self.tag); + } + }]; +} + +#pragma mark - Logging + ++ (NSString *)tag +{ + return [NSString stringWithFormat:@"[%@]", self.class]; +} + +- (NSString *)tag +{ + return self.class.tag; +} + +@end + +NS_ASSUME_NONNULL_END diff --git a/src/Storage/TSStorageManager.m b/src/Storage/TSStorageManager.m index c3d9a613d..eaca489b8 100644 --- a/src/Storage/TSStorageManager.m +++ b/src/Storage/TSStorageManager.m @@ -6,6 +6,7 @@ #import "NSData+Base64.h" #import "OWSAnalytics.h" #import "OWSDisappearingMessagesFinder.h" +#import "OWSFailedAttachmentDownloadsJob.h" #import "OWSFailedMessagesJob.h" #import "OWSIncomingMessageFinder.h" #import "OWSReadReceipt.h" @@ -205,6 +206,9 @@ static NSString *keychainDBPassAccount = @"TSDatabasePass"; [finder asyncRegisterDatabaseExtensions]; OWSFailedMessagesJob *failedMessagesJob = [[OWSFailedMessagesJob alloc] initWithStorageManager:self]; [failedMessagesJob asyncRegisterDatabaseExtensions]; + OWSFailedAttachmentDownloadsJob *failedAttachmentDownloadsMessagesJob = + [[OWSFailedAttachmentDownloadsJob alloc] initWithStorageManager:self]; + [failedAttachmentDownloadsMessagesJob asyncRegisterDatabaseExtensions]; } - (void)protectSignalFiles {