Use background task during message processing.

SignalServiceKit/src/Messages/Attachments/OWSAttachmentsProcessor.m

@@ -7,6 +7,7 @@
 #import "Cryptography.h"
 #import "MIMETypeUtil.h"
 #import "NSNotificationCenter+OWS.h"
+#import "OWSBackgroundTask.h"
 #import "OWSError.h"
 #import "OWSSignalServiceProtos.pb.h"
 #import "TSAttachmentPointer.h"
@@ -156,26 +157,7 @@ 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;
-        });
-    };
+    __block OWSBackgroundTask *backgroundTask = [[OWSBackgroundTask alloc] initWithLabelStr:__PRETTY_FUNCTION__];
 
     [self setAttachment:attachment isDownloadingInMessage:message transaction:transaction];
 
@@ -185,7 +167,7 @@ static const CGFloat kAttachmentDownloadProgressTheta = 0.001f;
             [self setAttachment:attachment didFailInMessage:message error:error];
             failureHandler(error);
 
-            clearBackgroundTask();
+            backgroundTask = nil;
         });
     };
 
@@ -197,7 +179,7 @@ static const CGFloat kAttachmentDownloadProgressTheta = 0.001f;
                 [message touch];
             }
 
-            clearBackgroundTask();
+            backgroundTask = nil;
         });
     };
 

SignalServiceKit/src/Messages/OWSBatchMessageProcessor.m

@@ -5,6 +5,7 @@
 #import "OWSBatchMessageProcessor.h"
 #import "AppContext.h"
 #import "NSArray+OWS.h"
+#import "OWSBackgroundTask.h"
 #import "OWSMessageManager.h"
 #import "OWSQueues.h"
 #import "OWSSignalServiceProtos.pb.h"
@@ -335,10 +336,14 @@ NSString *const OWSMessageContentJobFinderExtensionGroup = @"OWSMessageContentJo
         return;
     }
 
+    OWSBackgroundTask *backgroundTask = [[OWSBackgroundTask alloc] initWithLabelStr:__PRETTY_FUNCTION__];
+
     [self processJobs:jobs];
 
     [self.finder removeJobsWithIds:jobs.uniqueIds];
 
+    backgroundTask = nil;
+
     DDLogVerbose(@"%@ completed %zd jobs. %zd jobs left.",
         self.logTag,
         jobs.count,

SignalServiceKit/src/Messages/OWSMessageReceiver.m

@@ -5,6 +5,7 @@
 #import "OWSMessageReceiver.h"
 #import "AppContext.h"
 #import "NSArray+OWS.h"
+#import "OWSBackgroundTask.h"
 #import "OWSBatchMessageProcessor.h"
 #import "OWSMessageDecrypter.h"
 #import "OWSQueues.h"
@@ -308,6 +309,8 @@ NSString *const OWSMessageDecryptJobFinderExtensionGroup = @"OWSMessageProcessin
         return;
     }
 
+    __block OWSBackgroundTask *backgroundTask = [[OWSBackgroundTask alloc] initWithLabelStr:__PRETTY_FUNCTION__];
+
     [self processJob:job
           completion:^(BOOL success) {
               [self.finder removeJobWithId:job.uniqueId];
@@ -316,6 +319,7 @@ NSString *const OWSMessageDecryptJobFinderExtensionGroup = @"OWSMessageProcessin
                   success ? @"decrypted" : @"failed to decrypt",
                   (unsigned long)[OWSMessageDecryptJob numberOfKeysInCollection]);
               [self drainQueueWorkStep];
+              backgroundTask = nil;
           }];
 }
 

SignalServiceKit/src/Network/WebSockets/TSSocketManager.m

@@ -7,6 +7,7 @@
 #import "Cryptography.h"
 #import "NSNotificationCenter+OWS.h"
 #import "NSTimer+OWS.h"
+#import "OWSBackgroundTask.h"
 #import "OWSMessageManager.h"
 #import "OWSMessageReceiver.h"
 #import "OWSSignalService.h"
@@ -381,6 +382,8 @@ NSString *const kNSNotification_SocketManagerStateDidChange = @"kNSNotification_
 
     if ([message.path isEqualToString:@"/api/v1/message"] && [message.verb isEqualToString:@"PUT"]) {
 
+        __block OWSBackgroundTask *backgroundTask = [[OWSBackgroundTask alloc] initWithLabelStr:__PRETTY_FUNCTION__];
+
         dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
 
             NSData *decryptedPayload =
@@ -389,6 +392,7 @@ NSString *const kNSNotification_SocketManagerStateDidChange = @"kNSNotification_
             if (!decryptedPayload) {
                 DDLogWarn(@"%@ Failed to decrypt incoming payload or bad HMAC", self.logTag);
                 [self sendWebSocketMessageAcknowledgement:message];
+                backgroundTask = nil;
                 return;
             }
 
@@ -398,6 +402,7 @@ NSString *const kNSNotification_SocketManagerStateDidChange = @"kNSNotification_
 
             dispatch_async(dispatch_get_main_queue(), ^{
                 [self sendWebSocketMessageAcknowledgement:message];
+                backgroundTask = nil;
             });
         });
     } else {

SignalServiceKit/src/Util/OWSBackgroundTask.h

@@ -0,0 +1,10 @@
+//
+//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
+//
+
+@interface OWSBackgroundTask : NSObject
+
+- (instancetype)init NS_UNAVAILABLE;
+- (instancetype)initWithLabelStr:(const char *)labelStr;
+
+@end

SignalServiceKit/src/Util/OWSBackgroundTask.m

@@ -0,0 +1,85 @@
+//
+//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
+//
+
+#import "OWSBackgroundTask.h"
+#import "AppContext.h"
+
+@interface OWSBackgroundTask ()
+
+@property (nonatomic) UIBackgroundTaskIdentifier backgroundTaskId;
+@property (nonatomic, readonly) NSString *label;
+
+@end
+
+#pragma mark -
+
+@implementation OWSBackgroundTask
+
+- (instancetype)initWithLabelStr:(const char *)labelStr
+{
+    self = [super init];
+
+    if (!self) {
+        return self;
+    }
+
+    OWSAssert(labelStr);
+
+    _label = [NSString stringWithFormat:@"%s", labelStr];
+
+    [self startBackgroundTask];
+
+    return self;
+}
+
+- (void)dealloc
+{
+    [self endBackgroundTask];
+}
+
+- (void)startBackgroundTask
+{
+    @synchronized(self)
+    {
+        __weak typeof(self) weakSelf = self;
+
+        self.backgroundTaskId = [CurrentAppContext() beginBackgroundTaskWithExpirationHandler:^{
+            OWSAssert([NSThread isMainThread]);
+            __strong typeof(self) strongSelf = weakSelf;
+            if (!strongSelf) {
+                return;
+            }
+            @synchronized(strongSelf)
+            {
+                if (strongSelf.backgroundTaskId == UIBackgroundTaskInvalid) {
+                    return;
+                }
+                DDLogInfo(@"%@ %@ background task expired", strongSelf.logTag, strongSelf.label);
+                strongSelf.backgroundTaskId = UIBackgroundTaskInvalid;
+            }
+        }];
+    }
+}
+
+- (void)endBackgroundTask
+{
+    __weak typeof(self) weakSelf = self;
+    dispatch_async(dispatch_get_main_queue(), ^{
+        __strong typeof(self) strongSelf = weakSelf;
+        if (!strongSelf) {
+            return;
+        }
+        @synchronized(strongSelf)
+        {
+            if (strongSelf.backgroundTaskId == UIBackgroundTaskInvalid) {
+                return;
+            }
+            DDLogInfo(@"%@ %@ background task completed", strongSelf.logTag, strongSelf.label);
+            [CurrentAppContext() endBackgroundTask:strongSelf.backgroundTaskId];
+            strongSelf.backgroundTaskId = UIBackgroundTaskInvalid;
+        }
+    });
+}
+
+@end