Merge branch 'charlesmchen/failedAttachmentDownloads'

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];

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

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

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;

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

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 <YapDatabase/YapDatabaseConnection.h>
+#import <YapDatabase/YapDatabaseQuery.h>
+#import <YapDatabase/YapDatabaseSecondaryIndex.h>
+
+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<NSString *> *)fetchAttemptingOutAttachmentIds:(YapDatabaseConnection *)dbConnection
+{
+    NSMutableArray<NSString *> *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

src/Messages/OWSFailedMessagesJob.m

@@ -21,6 +21,8 @@ static NSString *const OWSFailedMessagesJobMessageStateIndex = @"index_outoing_m
 
 @end
 
+#pragma mark -
+
 @implementation OWSFailedMessagesJob
 
 - (instancetype)initWithStorageManager:(TSStorageManager *)storageManager

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 {