Improve background task logic.

Signal/src/AppDelegate.m

@@ -445,7 +445,8 @@ static NSString *const kURLHostVerifyPrefix             = @"verify";
 - (void)applicationWillResignActive:(UIApplication *)application {
     DDLogWarn(@"%@ applicationWillResignActive.", self.logTag);
 
-    UIBackgroundTaskIdentifier bgTask = [application beginBackgroundTaskWithExpirationHandler:nil];
+    __block OWSBackgroundTask *backgroundTask = [OWSBackgroundTask backgroundTaskWithLabelStr:__PRETTY_FUNCTION__];
+
     dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
         if ([TSAccountManager isRegistered]) {
             dispatch_async(dispatch_get_main_queue(), ^{
@@ -454,8 +455,11 @@ static NSString *const kURLHostVerifyPrefix             = @"verify";
                     [self showScreenProtection];
                 }
                 [SignalApp.sharedApp.homeViewController updateInboxCountLabel];
-                [application endBackgroundTask:bgTask];
+
+                backgroundTask = nil;
             });
+        } else {
+            backgroundTask = nil;
         }
     });
 

Signal/src/Signal-Bridging-Header.h

@@ -73,6 +73,7 @@
 #import <SignalServiceKit/OWSAnalytics.h>
 #import <SignalServiceKit/OWSAnalyticsEvents.h>
 #import <SignalServiceKit/OWSAttachmentsProcessor.h>
+#import <SignalServiceKit/OWSBackgroundTask.h>
 #import <SignalServiceKit/OWSCallAnswerMessage.h>
 #import <SignalServiceKit/OWSCallBusyMessage.h>
 #import <SignalServiceKit/OWSCallHangupMessage.h>

Signal/src/call/CallService.swift

@@ -584,16 +584,19 @@ protocol CallServiceObserver: class {
 
         self.call = newCall
 
-        let backgroundTask = UIApplication.shared.beginBackgroundTask {
+        var backgroundTask = OWSBackgroundTask(label:"\(#function)", completionBlock: { [weak self] _ in
+            AssertIsOnMainThread()
+            guard let strongSelf = self else {
+                return
+            }
             let timeout = CallError.timeout(description: "background task time ran out before call connected.")
-            DispatchQueue.main.async {
+
-                guard self.call == newCall else {
+            guard strongSelf.call == newCall else {
-                    Logger.warn("\(self.logTag) ignoring obsolete call in \(#function)")
+                Logger.warn("\(strongSelf.logTag) ignoring obsolete call in \(#function)")
-                    return
+                return
-                }
-                self.handleFailedCall(failedCall: newCall, error: timeout)
             }
-        }
+            strongSelf.handleFailedCall(failedCall: newCall, error: timeout)
+        })
 
         let incomingCallPromise = firstly {
             return getIceServers()
@@ -674,7 +677,8 @@ protocol CallServiceObserver: class {
             }
         }.always {
             Logger.debug("\(self.logTag) ending background task awaiting inbound call connection")
-            UIApplication.shared.endBackgroundTask(backgroundTask)
+
+            backgroundTask = nil
         }
         incomingCallPromise.retainUntilComplete()
     }

SignalServiceKit/src/Messages/Attachments/OWSAttachmentsProcessor.m

@@ -157,7 +157,7 @@ static const CGFloat kAttachmentDownloadProgressTheta = 0.001f;
 {
     OWSAssert(transaction);
 
-    __block OWSBackgroundTask *backgroundTask = [[OWSBackgroundTask alloc] initWithLabelStr:__PRETTY_FUNCTION__];
+    __block OWSBackgroundTask *backgroundTask = [OWSBackgroundTask backgroundTaskWithLabelStr:__PRETTY_FUNCTION__];
 
     [self setAttachment:attachment isDownloadingInMessage:message transaction:transaction];
 

SignalServiceKit/src/Messages/OWSBatchMessageProcessor.m

@@ -336,7 +336,7 @@ NSString *const OWSMessageContentJobFinderExtensionGroup = @"OWSMessageContentJo
         return;
     }
 
-    OWSBackgroundTask *backgroundTask = [[OWSBackgroundTask alloc] initWithLabelStr:__PRETTY_FUNCTION__];
+    OWSBackgroundTask *backgroundTask = [OWSBackgroundTask backgroundTaskWithLabelStr:__PRETTY_FUNCTION__];
 
     [self processJobs:jobs];
 

SignalServiceKit/src/Messages/OWSMessageReceiver.m

@@ -309,7 +309,7 @@ NSString *const OWSMessageDecryptJobFinderExtensionGroup = @"OWSMessageProcessin
         return;
     }
 
-    __block OWSBackgroundTask *backgroundTask = [[OWSBackgroundTask alloc] initWithLabelStr:__PRETTY_FUNCTION__];
+    __block OWSBackgroundTask *backgroundTask = [OWSBackgroundTask backgroundTaskWithLabelStr:__PRETTY_FUNCTION__];
 
     [self processJob:job
           completion:^(BOOL success) {

SignalServiceKit/src/Messages/OWSMessageSender.m

@@ -7,6 +7,7 @@
 #import "ContactsUpdater.h"
 #import "NSData+keyVersionByte.h"
 #import "NSData+messagePadding.h"
+#import "OWSBackgroundTask.h"
 #import "OWSBlockingManager.h"
 #import "OWSDevice.h"
 #import "OWSDisappearingMessagesJob.h"
@@ -122,11 +123,6 @@ static void *kNSError_MessageSender_IsFatal = &kNSError_MessageSender_IsFatal;
                         success:(void (^)(void))successHandler
                         failure:(void (^)(NSError *_Nonnull error))failureHandler NS_DESIGNATED_INITIALIZER;
 
-#pragma mark - background task mgmt
-
-- (void)startBackgroundTask;
-- (void)endBackgroundTask;
-
 @end
 
 #pragma mark -
@@ -159,7 +155,7 @@ NSUInteger const OWSSendMessageOperationMaxRetries = 4;
 @property (nonatomic, readonly) void (^successHandler)(void);
 @property (nonatomic, readonly) void (^failureHandler)(NSError *_Nonnull error);
 @property (nonatomic) OWSSendMessageOperationState operationState;
-@property (nonatomic) UIBackgroundTaskIdentifier backgroundTaskIdentifier;
+@property (nonatomic) OWSBackgroundTask *backgroundTask;
 
 @end
 
@@ -178,7 +174,7 @@ NSUInteger const OWSSendMessageOperationMaxRetries = 4;
     }
 
     _operationState = OWSSendMessageOperationStateNew;
-    _backgroundTaskIdentifier = UIBackgroundTaskInvalid;
+    self.backgroundTask = [OWSBackgroundTask backgroundTaskWithLabelStr:__PRETTY_FUNCTION__];
 
     _message = message;
     _messageSender = messageSender;
@@ -216,42 +212,6 @@ NSUInteger const OWSSendMessageOperationMaxRetries = 4;
     return self;
 }
 
-#pragma mark - background task mgmt
-
-// We want to make sure to finish sending any in-flight messages when the app is backgrounded.
-// We have to call `startBackgroundTask` *before* the task is enqueued, since we can't guarantee when the operation will
-// be dequeued.
-- (void)startBackgroundTask
-{
-    AssertIsOnMainThread();
-    OWSAssert(self.backgroundTaskIdentifier == UIBackgroundTaskInvalid);
-
-    self.backgroundTaskIdentifier = [CurrentAppContext() beginBackgroundTaskWithExpirationHandler:^{
-        DDLogWarn(@"%@ Timed out while in background trying to send message: %@", self.logTag, self.message);
-        [self endBackgroundTask];
-    }];
-}
-
-- (void)endBackgroundTask
-{
-    if (self.backgroundTaskIdentifier == UIBackgroundTaskInvalid) {
-        return;
-    }
-    [CurrentAppContext() endBackgroundTask:self.backgroundTaskIdentifier];
-    self.backgroundTaskIdentifier = UIBackgroundTaskInvalid;
-}
-
-- (void)setBackgroundTaskIdentifier:(UIBackgroundTaskIdentifier)backgroundTaskIdentifier
-{
-    AssertIsOnMainThread();
-
-    // Should only be sent once per operation
-    OWSAssert(!CurrentAppContext().isMainApp || _backgroundTaskIdentifier == UIBackgroundTaskInvalid);
-    OWSAssert(!CurrentAppContext().isMainApp || backgroundTaskIdentifier != UIBackgroundTaskInvalid);
-
-    _backgroundTaskIdentifier = backgroundTaskIdentifier;
-}
-
 #pragma mark - NSOperation overrides
 
 - (BOOL)isExecuting
@@ -266,11 +226,6 @@ NSUInteger const OWSSendMessageOperationMaxRetries = 4;
 
 - (void)start
 {
-    // Should call `startBackgroundTask` before enqueuing the operation
-    // to ensure we don't get suspended before the operation completes.
-
-    OWSAssert(!CurrentAppContext().isMainApp || self.backgroundTaskIdentifier != UIBackgroundTaskInvalid);
-
     [self willChangeValueForKey:OWSSendMessageOperationKeyIsExecuting];
     self.operationState = OWSSendMessageOperationStateExecuting;
     [self didChangeValueForKey:OWSSendMessageOperationKeyIsExecuting];
@@ -343,8 +298,6 @@ NSUInteger const OWSSendMessageOperationMaxRetries = 4;
 
     [self didChangeValueForKey:OWSSendMessageOperationKeyIsExecuting];
     [self didChangeValueForKey:OWSSendMessageOperationKeyIsFinished];
-
-    [self endBackgroundTask];
 }
 
 @end
@@ -456,17 +409,9 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException";
                                                      success:successHandler
                                                      failure:failureHandler];
 
-        // startBackgroundTask must be called on the main thread.
+        dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
-        dispatch_async(dispatch_get_main_queue(), ^{
+            NSOperationQueue *sendingQueue = [self sendingQueueForMessage:message];
-            // We call `startBackgroundTask` here to prevent our app from suspending while being backgrounded
+            [sendingQueue addOperation:sendMessageOperation];
-            // until the operation is completed - at which point the OWSSendMessageOperation ends it's background task.
-
-            [sendMessageOperation startBackgroundTask];
-
-            dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
-                NSOperationQueue *sendingQueue = [self sendingQueueForMessage:message];
-                [sendingQueue addOperation:sendMessageOperation];
-            });
         });
     });
 }

SignalServiceKit/src/Network/WebSockets/TSSocketManager.m

@@ -73,7 +73,7 @@ NSString *const kNSNotification_SocketManagerStateDidChange = @"kNSNotification_
 @property (nonatomic) NSTimer *backgroundKeepAliveTimer;
 // This is used to manage the iOS "background task" used to
 // keep the app alive in the background.
-@property (nonatomic) UIBackgroundTaskIdentifier fetchingTaskIdentifier;
+@property (nonatomic) OWSBackgroundTask *backgroundTask;
 
 // We cache this value instead of consulting [UIApplication sharedApplication].applicationState,
 // because UIKit only provides a "will resign active" notification, not a "did resign active"
@@ -101,7 +101,6 @@ NSString *const kNSNotification_SocketManagerStateDidChange = @"kNSNotification_
     _signalService = [OWSSignalService sharedInstance];
     _messageReceiver = [OWSMessageReceiver sharedInstance];
     _state = SocketManagerStateClosed;
-    _fetchingTaskIdentifier = UIBackgroundTaskInvalid;
 
     OWSSingletonAssert();
 
@@ -382,7 +381,7 @@ NSString *const kNSNotification_SocketManagerStateDidChange = @"kNSNotification_
 
     if ([message.path isEqualToString:@"/api/v1/message"] && [message.verb isEqualToString:@"PUT"]) {
 
-        __block OWSBackgroundTask *backgroundTask = [[OWSBackgroundTask alloc] initWithLabelStr:__PRETTY_FUNCTION__];
+        __block OWSBackgroundTask *backgroundTask = [OWSBackgroundTask backgroundTaskWithLabelStr:__PRETTY_FUNCTION__];
 
         dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
 
@@ -518,7 +517,6 @@ NSString *const kNSNotification_SocketManagerStateDidChange = @"kNSNotification_
         return YES;
     } else if (self.backgroundKeepAliveUntilDate && [self.backgroundKeepAliveUntilDate timeIntervalSinceNow] > 0.f) {
         OWSAssert(self.backgroundKeepAliveTimer);
-        OWSAssert(self.fetchingTaskIdentifier != UIBackgroundTaskInvalid);
         // If app is doing any work in the background, keep web socket alive.
         return YES;
     } else {
@@ -537,7 +535,6 @@ NSString *const kNSNotification_SocketManagerStateDidChange = @"kNSNotification_
     } else if (!self.backgroundKeepAliveUntilDate) {
         OWSAssert(!self.backgroundKeepAliveUntilDate);
         OWSAssert(!self.backgroundKeepAliveTimer);
-        OWSAssert(self.fetchingTaskIdentifier == UIBackgroundTaskInvalid);
 
         DDLogInfo(@"%s activating socket in the background", __PRETTY_FUNCTION__);
 
@@ -556,19 +553,25 @@ NSString *const kNSNotification_SocketManagerStateDidChange = @"kNSNotification_
         // Additionally, we want the reconnect timer to work in the background too.
         [[NSRunLoop mainRunLoop] addTimer:self.backgroundKeepAliveTimer forMode:NSDefaultRunLoopMode];
 
-        self.fetchingTaskIdentifier = [CurrentAppContext() beginBackgroundTaskWithExpirationHandler:^{
+        __weak typeof(self) weakSelf = self;
-            OWSAssert([NSThread isMainThread]);
+        self.backgroundTask =
-
+            [OWSBackgroundTask backgroundTaskWithLabelStr:__PRETTY_FUNCTION__
-            DDLogInfo(@"%s background task expired", __PRETTY_FUNCTION__);
+                                          completionBlock:^(BackgroundTaskState backgroundTaskState) {
-
+                                              OWSAssert([NSThread isMainThread]);
-            [self clearBackgroundState];
+                                              __strong typeof(self) strongSelf = weakSelf;
-            [self applyDesiredSocketState];
+                                              if (!strongSelf) {
-        }];
+                                                  return;
+                                              }
+
+                                              if (backgroundTaskState == BackgroundTaskState_Expired) {
+                                                  [strongSelf clearBackgroundState];
+                                              }
+                                              [strongSelf applyDesiredSocketState];
+                                          }];
     } else {
         OWSAssert(self.backgroundKeepAliveUntilDate);
         OWSAssert(self.backgroundKeepAliveTimer);
         OWSAssert([self.backgroundKeepAliveTimer isValid]);
-        OWSAssert(self.fetchingTaskIdentifier != UIBackgroundTaskInvalid);
 
         if ([self.backgroundKeepAliveUntilDate timeIntervalSinceNow] < durationSeconds) {
             // Update state used to keep socket alive in background.
@@ -628,11 +631,7 @@ NSString *const kNSNotification_SocketManagerStateDidChange = @"kNSNotification_
     self.backgroundKeepAliveUntilDate = nil;
     [self.backgroundKeepAliveTimer invalidate];
     self.backgroundKeepAliveTimer = nil;
-
+    self.backgroundTask = nil;
-    if (self.fetchingTaskIdentifier != UIBackgroundTaskInvalid) {
-        [CurrentAppContext() endBackgroundTask:self.fetchingTaskIdentifier];
-        self.fetchingTaskIdentifier = UIBackgroundTaskInvalid;
-    }
 }
 
 #pragma mark - Reconnect

SignalServiceKit/src/Util/OWSAnalytics.m

@@ -4,6 +4,7 @@
 
 #import "OWSAnalytics.h"
 #import "AppContext.h"
+#import "OWSBackgroundTask.h"
 #import "OWSQueues.h"
 #import "TSStorageManager.h"
 #import <CocoaLumberjack/CocoaLumberjack.h>
@@ -231,22 +232,19 @@ NSString *NSStringForOWSAnalyticsSeverity(OWSAnalyticsSeverity severity)
 
     DDLogDebug(@"%@ submitting: %@", self.logTag, eventKey);
 
-    __block UIBackgroundTaskIdentifier task;
+    __block OWSBackgroundTask *backgroundTask =
-    task = [CurrentAppContext() beginBackgroundTaskWithExpirationHandler:^{
+        [OWSBackgroundTask backgroundTaskWithLabelStr:__PRETTY_FUNCTION__
-        failureBlock();
+                                      completionBlock:^(BackgroundTaskState backgroundTaskState) {
-
+                                          if (backgroundTaskState == BackgroundTaskState_Success) {
-        [CurrentAppContext() endBackgroundTask:task];
+                                              successBlock();
-    }];
+                                          } else {
+                                              failureBlock();
+                                          }
+                                      }];
 
     // Until we integrate with an analytics platform, behave as though all event delivery succeeds.
     dispatch_async(self.serialQueue, ^{
-        BOOL success = YES;
+        backgroundTask = nil;
-        if (success) {
-            successBlock();
-        } else {
-            failureBlock();
-        }
-        [CurrentAppContext() endBackgroundTask:task];
     });
 }
 

SignalServiceKit/src/Util/OWSBackgroundTask.h

@@ -2,9 +2,28 @@
 //  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
 //
 
+typedef NS_ENUM(NSUInteger, BackgroundTaskState) {
+    BackgroundTaskState_Success,
+    BackgroundTaskState_CouldNotStart,
+    BackgroundTaskState_Expired,
+};
+
+typedef void (^BackgroundTaskCompletionBlock)(BackgroundTaskState backgroundTaskState);
+
 @interface OWSBackgroundTask : NSObject
 
 - (instancetype)init NS_UNAVAILABLE;
-- (instancetype)initWithLabelStr:(const char *)labelStr;
+
++ (OWSBackgroundTask *)backgroundTaskWithLabelStr:(const char *)labelStr;
+
+// completionBlock will be called exactly once on the main thread.
++ (OWSBackgroundTask *)backgroundTaskWithLabelStr:(const char *)labelStr
+                                  completionBlock:(BackgroundTaskCompletionBlock)completionBlock;
+
++ (OWSBackgroundTask *)backgroundTaskWithLabel:(NSString *)label;
+
+// completionBlock will be called exactly once on the main thread.
++ (OWSBackgroundTask *)backgroundTaskWithLabel:(NSString *)label
+                               completionBlock:(BackgroundTaskCompletionBlock)completionBlock;
 
 @end

SignalServiceKit/src/Util/OWSBackgroundTask.m

@@ -7,16 +7,49 @@
 
 @interface OWSBackgroundTask ()
 
-@property (nonatomic) UIBackgroundTaskIdentifier backgroundTaskId;
 @property (nonatomic, readonly) NSString *label;
 
+// This property should only be accessed while synchronized on this instance.
+@property (nonatomic) UIBackgroundTaskIdentifier backgroundTaskId;
+
+@property (nonatomic, nullable) BackgroundTaskCompletionBlock completionBlock;
+
 @end
 
 #pragma mark -
 
 @implementation OWSBackgroundTask
 
-- (instancetype)initWithLabelStr:(const char *)labelStr
++ (OWSBackgroundTask *)backgroundTaskWithLabelStr:(const char *)labelStr
+{
+    OWSAssert(labelStr);
+
+    NSString *label = [NSString stringWithFormat:@"%s", labelStr];
+    return [[OWSBackgroundTask alloc] initWithLabel:label completionBlock:nil];
+}
+
++ (OWSBackgroundTask *)backgroundTaskWithLabelStr:(const char *)labelStr
+                                  completionBlock:(BackgroundTaskCompletionBlock)completionBlock
+{
+
+    OWSAssert(labelStr);
+
+    NSString *label = [NSString stringWithFormat:@"%s", labelStr];
+    return [[OWSBackgroundTask alloc] initWithLabel:label completionBlock:completionBlock];
+}
+
++ (OWSBackgroundTask *)backgroundTaskWithLabel:(NSString *)label
+{
+    return [[OWSBackgroundTask alloc] initWithLabel:label completionBlock:nil];
+}
+
++ (OWSBackgroundTask *)backgroundTaskWithLabel:(NSString *)label
+                               completionBlock:(BackgroundTaskCompletionBlock)completionBlock
+{
+    return [[OWSBackgroundTask alloc] initWithLabel:label completionBlock:completionBlock];
+}
+
+- (instancetype)initWithLabel:(NSString *)label completionBlock:(BackgroundTaskCompletionBlock _Nullable)completionBlock
 {
     self = [super init];
 
@@ -24,9 +57,10 @@
         return self;
     }
 
-    OWSAssert(labelStr);
+    OWSAssert(label.length > 0);
 
-    _label = [NSString stringWithFormat:@"%s", labelStr];
+    _label = label;
+    self.completionBlock = completionBlock;
 
     [self startBackgroundTask];
 
@@ -40,10 +74,9 @@
 
 - (void)startBackgroundTask
 {
-    @synchronized(self)
+    __weak typeof(self) weakSelf = self;
-    {
+    // beginBackgroundTaskWithExpirationHandler must be called on the main thread.
-        __weak typeof(self) weakSelf = self;
+    dispatch_async(dispatch_get_main_queue(), ^{
-
         self.backgroundTaskId = [CurrentAppContext() beginBackgroundTaskWithExpirationHandler:^{
             OWSAssert([NSThread isMainThread]);
             __strong typeof(self) strongSelf = weakSelf;
@@ -57,27 +90,55 @@
                 }
                 DDLogInfo(@"%@ %@ background task expired", strongSelf.logTag, strongSelf.label);
                 strongSelf.backgroundTaskId = UIBackgroundTaskInvalid;
+
+                if (strongSelf.completionBlock) {
+                    strongSelf.completionBlock(BackgroundTaskState_Expired);
+                    strongSelf.completionBlock = nil;
+                }
             }
         }];
-    }
+
+        // If a background task could not be begun, call the completion block.
+        if (self.backgroundTaskId == UIBackgroundTaskInvalid) {
+            BackgroundTaskCompletionBlock _Nullable completionBlock;
+            @synchronized(self)
+            {
+                completionBlock = self.completionBlock;
+                self.completionBlock = nil;
+            }
+            if (completionBlock) {
+                dispatch_async(dispatch_get_main_queue(), ^{
+                    completionBlock(BackgroundTaskState_CouldNotStart);
+                });
+            }
+        }
+    });
 }
 
 - (void)endBackgroundTask
 {
-    __weak typeof(self) weakSelf = self;
+    // Make a local copy of this state, since this method is called by `dealloc`.
+    UIBackgroundTaskIdentifier backgroundTaskId;
+    NSString *logTag = self.logTag;
+    NSString *label = self.label;
+    BackgroundTaskCompletionBlock _Nullable completionBlock = self.completionBlock;
+
+    @synchronized(self)
+    {
+        backgroundTaskId = self.backgroundTaskId;
+    }
+
+    if (backgroundTaskId == UIBackgroundTaskInvalid) {
+        return;
+    }
+
+    // endBackgroundTask must be called on the main thread.
     dispatch_async(dispatch_get_main_queue(), ^{
-        __strong typeof(self) strongSelf = weakSelf;
+        DDLogInfo(@"%@ %@ background task completed", logTag, label);
-        if (!strongSelf) {
+        [CurrentAppContext() endBackgroundTask:backgroundTaskId];
-            return;
+
-        }
+        if (completionBlock) {
-        @synchronized(strongSelf)
+            completionBlock(BackgroundTaskState_Success);
-        {
-            if (strongSelf.backgroundTaskId == UIBackgroundTaskInvalid) {
-                return;
-            }
-            DDLogInfo(@"%@ %@ background task completed", strongSelf.logTag, strongSelf.label);
-            [CurrentAppContext() endBackgroundTask:strongSelf.backgroundTaskId];
-            strongSelf.backgroundTaskId = UIBackgroundTaskInvalid;
         }
     });
 }