From df8ded90b8ee722d25d96002dc43c0fe3f4bad41 Mon Sep 17 00:00:00 2001
From: Matthew Chen <charlesmchen@gmail.com>
Date: Fri, 15 Dec 2017 11:37:37 -0500
Subject: [PATCH] Use background task during attachment downloads.

---
 .../Attachments/OWSAttachmentsProcessor.m     | 26 +++++++++++++++++++
 .../src/Messages/OWSMessageSender.m           |  4 +++
 2 files changed, 30 insertions(+)

diff --git a/SignalServiceKit/src/Messages/Attachments/OWSAttachmentsProcessor.m b/SignalServiceKit/src/Messages/Attachments/OWSAttachmentsProcessor.m
index 42af6c7587..afdacf9627 100644
--- a/SignalServiceKit/src/Messages/Attachments/OWSAttachmentsProcessor.m
+++ b/SignalServiceKit/src/Messages/Attachments/OWSAttachmentsProcessor.m
@@ -3,6 +3,7 @@
 //
 
 #import "OWSAttachmentsProcessor.h"
+#import "AppContext.h"
 #import "Cryptography.h"
 #import "MIMETypeUtil.h"
 #import "NSNotificationCenter+OWS.h"
@@ -155,6 +156,27 @@ static const CGFloat kAttachmentDownloadProgressTheta = 0.001f;
 {
     OWSAssert(transaction);
 
+    __block UIBackgroundTaskIdentifier backgroundTaskId =
+        [CurrentAppContext() beginBackgroundTaskWithExpirationHandler:^{
+            OWSAssert([NSThread isMainThread]);
+
+            if (backgroundTaskId == UIBackgroundTaskInvalid) {
+                return;
+            }
+            DDLogInfo(@"%@ %s background task expired", self.logTag, __PRETTY_FUNCTION__);
+            backgroundTaskId = UIBackgroundTaskInvalid;
+        }];
+    void (^clearBackgroundTask)() = ^{
+        dispatch_async(dispatch_get_main_queue(), ^{
+            if (backgroundTaskId == UIBackgroundTaskInvalid) {
+                return;
+            }
+            DDLogInfo(@"%@ %s background task completed", self.logTag, __PRETTY_FUNCTION__);
+            [CurrentAppContext() endBackgroundTask:backgroundTaskId];
+            backgroundTaskId = UIBackgroundTaskInvalid;
+        });
+    };
+
     [self setAttachment:attachment isDownloadingInMessage:message transaction:transaction];
 
     void (^markAndHandleFailure)(NSError *) = ^(NSError *error) {
@@ -162,6 +184,8 @@ static const CGFloat kAttachmentDownloadProgressTheta = 0.001f;
         dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
             [self setAttachment:attachment didFailInMessage:message error:error];
             failureHandler(error);
+
+            clearBackgroundTask();
         });
     };
 
@@ -172,6 +196,8 @@ static const CGFloat kAttachmentDownloadProgressTheta = 0.001f;
             if (message) {
                 [message touch];
             }
+
+            clearBackgroundTask();
         });
     };
 
diff --git a/SignalServiceKit/src/Messages/OWSMessageSender.m b/SignalServiceKit/src/Messages/OWSMessageSender.m
index d65ef72d87..325d93eb77 100644
--- a/SignalServiceKit/src/Messages/OWSMessageSender.m
+++ b/SignalServiceKit/src/Messages/OWSMessageSender.m
@@ -234,7 +234,11 @@ NSUInteger const OWSSendMessageOperationMaxRetries = 4;
 
 - (void)endBackgroundTask
 {
+    if (self.backgroundTaskIdentifier == UIBackgroundTaskInvalid) {
+        return;
+    }
     [CurrentAppContext() endBackgroundTask:self.backgroundTaskIdentifier];
+    self.backgroundTaskIdentifier = UIBackgroundTaskInvalid;
 }
 
 - (void)setBackgroundTaskIdentifier:(UIBackgroundTaskIdentifier)backgroundTaskIdentifier