From 9cbc1e6a1b96dca0c9e1ab21578864fdf5a58620 Mon Sep 17 00:00:00 2001
From: Matthew Chen <charlesmchen@gmail.com>
Date: Thu, 6 Jul 2017 12:52:44 -0400
Subject: [PATCH] Block device from sleeping while Debug UI is visible and
 during database upgrades.

// FREEBIE
---
 Signal/src/AppDelegate.m                       |  8 ++++++++
 .../DebugUI/DebugUITableViewController.m       | 18 ++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/Signal/src/AppDelegate.m b/Signal/src/AppDelegate.m
index 3e5712ebc..ff438901f 100644
--- a/Signal/src/AppDelegate.m
+++ b/Signal/src/AppDelegate.m
@@ -110,6 +110,12 @@ static NSString *const kURLHostVerifyPrefix             = @"verify";
     // XXX - careful when moving this. It must happen before we initialize TSStorageManager.
     [self verifyDBKeysAvailableBeforeBackgroundLaunch];
 
+    // Prevent the device from sleeping during database view async registration
+    // (e.g. long database upgrades).
+    //
+    // This block will be cleared in databaseViewRegistrationComplete.
+    [DeviceSleepManager.sharedInstance addBlockWithBlockObject:self];
+
     [self setupEnvironment];
 
     [UIUtil applySignalAppearence];
@@ -783,6 +789,8 @@ static NSString *const kURLHostVerifyPrefix             = @"verify";
 {
     DDLogInfo(@"databaseViewRegistrationComplete");
 
+    [DeviceSleepManager.sharedInstance removeBlockWithBlockObject:self];
+
     [AppVersion.instance appLaunchDidComplete];
 
     [self ensureRootViewController];
diff --git a/Signal/src/ViewControllers/DebugUI/DebugUITableViewController.m b/Signal/src/ViewControllers/DebugUI/DebugUITableViewController.m
index 5ba7bef9e..d0af33ae1 100644
--- a/Signal/src/ViewControllers/DebugUI/DebugUITableViewController.m
+++ b/Signal/src/ViewControllers/DebugUI/DebugUITableViewController.m
@@ -28,6 +28,24 @@ NS_ASSUME_NONNULL_BEGIN
     return self.class.tag;
 }
 
+- (void)viewWillAppear:(BOOL)animated
+{
+    [super viewWillAppear:animated];
+
+    // Block device from sleeping while in the Debug UI.
+    //
+    // This is useful if you're using long-running actions in the
+    // Debug UI, like "send 1k messages", etc.
+    [DeviceSleepManager.sharedInstance addBlockWithBlockObject:self];
+}
+
+- (void)viewWillDisappear:(BOOL)animated
+{
+    [super viewWillDisappear:animated];
+
+    [DeviceSleepManager.sharedInstance removeBlockWithBlockObject:self];
+}
+
 #pragma mark - Factory Methods
 
 - (void)pushPageWithSection:(OWSTableSection *)section