From 0224af74656572b79da7b76ae378cd26b034400f Mon Sep 17 00:00:00 2001
From: Matthew Chen <charlesmchen@gmail.com>
Date: Fri, 14 Apr 2017 09:14:52 -0400
Subject: [PATCH] De-bounce the prekey checks.

// FREEBIE
---
 src/Account/TSPreKeyManager.h    |  2 +-
 src/Account/TSPreKeyManager.m    | 50 +++++++++++++++++++++++++++-----
 src/Messages/TSMessagesManager.m |  2 +-
 3 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/src/Account/TSPreKeyManager.h b/src/Account/TSPreKeyManager.h
index e6ce9a716..8f628c384 100644
--- a/src/Account/TSPreKeyManager.h
+++ b/src/Account/TSPreKeyManager.h
@@ -26,7 +26,7 @@ typedef NS_ENUM(NSInteger, RefreshPreKeysMode) {
                         success:(void (^)())successHandler
                         failure:(void (^)(NSError *error))failureHandler;
 
-+ (void)refreshPreKeys;
++ (void)checkPreKeys;
 
 + (void)checkPreKeysIfNecessary;
 
diff --git a/src/Account/TSPreKeyManager.m b/src/Account/TSPreKeyManager.m
index a81b94ce1..781cf9923 100644
--- a/src/Account/TSPreKeyManager.m
+++ b/src/Account/TSPreKeyManager.m
@@ -95,9 +95,21 @@ static const CGFloat kSignedPreKeyUpdateFailureMaxFailureDuration = 10 * 24 * 60
         BOOL shouldCheck = (lastPreKeyCheckTimestamp == nil
             || fabs([lastPreKeyCheckTimestamp timeIntervalSinceNow]) >= kPreKeyCheckFrequencySeconds);
         if (shouldCheck) {
+            // Optimistically mark the prekeys as checked. This
+            // de-bounces prekey checks.
+            //
+            // If the check or key registration fails, the prekeys
+            // will be marked as _NOT_ checked.
+            //
+            // Note: [TSPreKeyManager checkPreKeys] will also
+            //       optimistically mark them as checked. This
+            //       redundancy is fine and precludes a race
+            //       condition.
+            lastPreKeyCheckTimestamp = [NSDate date];
+
             [[TSAccountManager sharedInstance] ifRegistered:YES
                                                    runAsync:^{
-                                                       [TSPreKeyManager refreshPreKeys];
+                                                       [TSPreKeyManager checkPreKeys];
                                                    }];
         }
     });
@@ -110,6 +122,9 @@ static const CGFloat kSignedPreKeyUpdateFailureMaxFailureDuration = 10 * 24 * 60
     // We use prekeyQueue to serialize this logic and ensure that only
     // one thread is "registering" or "clearing" prekeys at a time.
     dispatch_async(TSPreKeyManager.prekeyQueue, ^{
+        // Mark the prekeys as checked every time we update.
+        lastPreKeyCheckTimestamp = [NSDate date];
+
         RefreshPreKeysMode modeCopy = mode;
         TSStorageManager *storageManager = [TSStorageManager sharedManager];
         ECKeyPair *identityKeyPair = [storageManager identityKeyPair];
@@ -165,6 +180,9 @@ static const CGFloat kSignedPreKeyUpdateFailureMaxFailureDuration = 10 * 24 * 60
             failure:^(NSURLSessionDataTask *task, NSError *error) {
                 OWSAnalyticsError(@"Prekey update failed (%@): %@", description, error);
 
+                // Mark the prekeys as _NOT_ checked on failure.
+                [self markPreKeysAsNotChecked];
+
                 failureHandler(error);
 
                 NSInteger statusCode = 0;
@@ -181,7 +199,17 @@ static const CGFloat kSignedPreKeyUpdateFailureMaxFailureDuration = 10 * 24 * 60
     });
 }
 
-+ (void)refreshPreKeys {
++ (void)checkPreKeys
+{
+    // Optimistically mark the prekeys as checked. This
+    // de-bounces prekey checks.
+    //
+    // If the check or key registration fails, the prekeys
+    // will be marked as _NOT_ checked.
+    dispatch_async(TSPreKeyManager.prekeyQueue, ^{
+        lastPreKeyCheckTimestamp = [NSDate date];
+    });
+
     // We want to update prekeys if either the one-time or signed prekeys need an update, so
     // we check the status of both.
     //
@@ -242,11 +270,6 @@ static const CGFloat kSignedPreKeyUpdateFailureMaxFailureDuration = 10 * 24 * 60
                     DDLogDebug(@"%@ Not updating prekeys.", self.tag);
                 }
             }
-            
-            // Update the prekey check timestamp on success.
-            dispatch_async(TSPreKeyManager.prekeyQueue, ^{
-                lastPreKeyCheckTimestamp = [NSDate date];
-            });
 
             if (!didUpdatePreKeys) {
                 // If we didn't update the prekeys, our local "current signed key" state should
@@ -273,14 +296,27 @@ static const CGFloat kSignedPreKeyUpdateFailureMaxFailureDuration = 10 * 24 * 60
                     }
                     failure:^(NSURLSessionDataTask *task, NSError *error) {
                         DDLogWarn(@"%@ Could not retrieve current signed key from the service.", self.tag);
+
+                        // Mark the prekeys as _NOT_ checked on failure.
+                        [self markPreKeysAsNotChecked];
                     }];
             }
         }
         failure:^(NSURLSessionDataTask *task, NSError *error) {
             DDLogError(@"%@ Failed to retrieve the number of available prekeys.", self.tag);
+
+            // Mark the prekeys as _NOT_ checked on failure.
+            [self markPreKeysAsNotChecked];
         }];
 }
 
++ (void)markPreKeysAsNotChecked
+{
+    dispatch_async(TSPreKeyManager.prekeyQueue, ^{
+        lastPreKeyCheckTimestamp = nil;
+    });
+}
+
 + (void)clearSignedPreKeyRecords {
     TSStorageManager *storageManager = [TSStorageManager sharedManager];
     NSNumber *currentSignedPrekeyId = [storageManager currentSignedPrekeyId];
diff --git a/src/Messages/TSMessagesManager.m b/src/Messages/TSMessagesManager.m
index c9681bf29..a15e5340d 100644
--- a/src/Messages/TSMessagesManager.m
+++ b/src/Messages/TSMessagesManager.m
@@ -304,7 +304,7 @@ NS_ASSUME_NONNULL_BEGIN
             NSData *plaintextData;
             @try {
                 // Check whether we need to refresh our PreKeys every time we receive a PreKeyWhisperMessage.
-                [TSPreKeyManager refreshPreKeys];
+                [TSPreKeyManager checkPreKeys];
 
                 PreKeyWhisperMessage *message = [[PreKeyWhisperMessage alloc] initWithData:encryptedData];
                 SessionCipher *cipher = [[SessionCipher alloc] initWithSessionStore:storageManager