Fix handler for multiple calls of -[FIRInstanceID instanceIDWithHandler:] (#2445)

Example/InstanceID/Tests/FIRInstanceIDTest.m

@@ -811,6 +811,182 @@ - (void)testDefaultToken_maxRetries {
   XCTAssertEqual(newTokenFetchCount, [FIRInstanceID maxRetryCountForDefaultToken]);
 }
 
+- (void)testInstanceIDWithHandler_WhileRequesting_Success {
+  [self stubKeyPairStoreToReturnValidKeypair];
+  [self mockAuthServiceToAlwaysReturnValidCheckin];
+
+  // Expect `fetchNewTokenWithAuthorizedEntity` to be called once
+  XCTestExpectation *fetchNewTokenExpectation =
+      [self expectationWithDescription:@"fetchNewTokenExpectation"];
+  __block FIRInstanceIDTokenHandler tokenHandler;
+
+  [[[self.mockTokenManager stub] andDo:^(NSInvocation *invocation) {
+    [invocation getArgument:&tokenHandler atIndex:6];
+    [fetchNewTokenExpectation fulfill];
+  }] fetchNewTokenWithAuthorizedEntity:kAuthorizedEntity
+                                 scope:kFIRInstanceIDDefaultTokenScope
+                               keyPair:[OCMArg any]
+                               options:[OCMArg any]
+                               handler:[OCMArg any]];
+
+  // Make 1st call
+  XCTestExpectation *handlerExpectation1 = [self expectationWithDescription:@"handlerExpectation1"];
+  FIRInstanceIDResultHandler handler1 =
+      ^(FIRInstanceIDResult *_Nullable result, NSError *_Nullable error) {
+        [handlerExpectation1 fulfill];
+        XCTAssertNotNil(result);
+        XCTAssertEqual(result.token, kToken);
+        XCTAssertNil(error);
+      };
+
+  [self.mockInstanceID instanceIDWithHandler:handler1];
+
+  // Make 2nd call
+  XCTestExpectation *handlerExpectation2 = [self expectationWithDescription:@"handlerExpectation1"];
+  FIRInstanceIDResultHandler handler2 =
+      ^(FIRInstanceIDResult *_Nullable result, NSError *_Nullable error) {
+        [handlerExpectation2 fulfill];
+        XCTAssertNotNil(result);
+        XCTAssertEqual(result.token, kToken);
+        XCTAssertNil(error);
+      };
+
+  [self.mockInstanceID instanceIDWithHandler:handler2];
+
+  // Wait for `fetchNewTokenWithAuthorizedEntity` to be performed
+  [self waitForExpectations:@[ fetchNewTokenExpectation ] timeout:1 enforceOrder:false];
+  // Finish token fetch request
+  tokenHandler(kToken, nil);
+
+  // Wait for completion handlers for both calls to be performed
+  [self waitForExpectationsWithTimeout:1 handler:NULL];
+}
+
+- (void)testInstanceIDWithHandler_WhileRequesting_RetrySuccess {
+  [self stubKeyPairStoreToReturnValidKeypair];
+  [self mockAuthServiceToAlwaysReturnValidCheckin];
+
+  // Expect `fetchNewTokenWithAuthorizedEntity` to be called twice
+  XCTestExpectation *fetchNewTokenExpectation1 =
+      [self expectationWithDescription:@"fetchNewTokenExpectation1"];
+  XCTestExpectation *fetchNewTokenExpectation2 =
+      [self expectationWithDescription:@"fetchNewTokenExpectation2"];
+  NSArray *fetchNewTokenExpectations = @[ fetchNewTokenExpectation1, fetchNewTokenExpectation2 ];
+
+  __block NSInteger fetchNewTokenCallCount = 0;
+  __block FIRInstanceIDTokenHandler tokenHandler;
+
+  [[[self.mockTokenManager stub] andDo:^(NSInvocation *invocation) {
+    [invocation getArgument:&tokenHandler atIndex:6];
+    [fetchNewTokenExpectations[fetchNewTokenCallCount] fulfill];
+    fetchNewTokenCallCount += 1;
+  }] fetchNewTokenWithAuthorizedEntity:kAuthorizedEntity
+                                 scope:kFIRInstanceIDDefaultTokenScope
+                               keyPair:[OCMArg any]
+                               options:[OCMArg any]
+                               handler:[OCMArg any]];
+
+  // Mock Instance ID's retry interval to 0, to vastly speed up this test.
+  [[[self.mockInstanceID stub] andReturnValue:@(0)] retryIntervalToFetchDefaultToken];
+
+  // Make 1st call
+  XCTestExpectation *handlerExpectation1 = [self expectationWithDescription:@"handlerExpectation1"];
+  FIRInstanceIDResultHandler handler1 =
+      ^(FIRInstanceIDResult *_Nullable result, NSError *_Nullable error) {
+        [handlerExpectation1 fulfill];
+        XCTAssertNotNil(result);
+        XCTAssertEqual(result.token, kToken);
+        XCTAssertNil(error);
+      };
+
+  [self.mockInstanceID instanceIDWithHandler:handler1];
+
+  // Make 2nd call
+  XCTestExpectation *handlerExpectation2 = [self expectationWithDescription:@"handlerExpectation1"];
+  FIRInstanceIDResultHandler handler2 =
+      ^(FIRInstanceIDResult *_Nullable result, NSError *_Nullable error) {
+        [handlerExpectation2 fulfill];
+        XCTAssertNotNil(result);
+        XCTAssertEqual(result.token, kToken);
+        XCTAssertNil(error);
+      };
+
+  [self.mockInstanceID instanceIDWithHandler:handler2];
+
+  // Wait for the 1st `fetchNewTokenWithAuthorizedEntity` to be performed
+  [self waitForExpectations:@[ fetchNewTokenExpectation1 ] timeout:1 enforceOrder:false];
+  // Fail for the 1st time
+  tokenHandler(nil, [NSError errorWithFIRInstanceIDErrorCode:kFIRInstanceIDErrorCodeUnknown]);
+
+  // Wait for the 2nd token feth
+  [self waitForExpectations:@[ fetchNewTokenExpectation2 ] timeout:1 enforceOrder:false];
+  // Finish with success
+  tokenHandler(kToken, nil);
+
+  // Wait for completion handlers for both calls to be performed
+  [self waitForExpectationsWithTimeout:1 handler:NULL];
+}
+
+- (void)testInstanceIDWithHandler_WhileRequesting_RetryFailure {
+  [self stubKeyPairStoreToReturnValidKeypair];
+  [self mockAuthServiceToAlwaysReturnValidCheckin];
+
+  // Expect `fetchNewTokenWithAuthorizedEntity` to be called once
+  NSMutableArray<XCTestExpectation *> *fetchNewTokenExpectations = [NSMutableArray array];
+  for (NSInteger i = 0; i < [[self.instanceID class] maxRetryCountForDefaultToken]; ++i) {
+    NSString *name = [NSString stringWithFormat:@"fetchNewTokenExpectation-%ld", (long)i];
+    [fetchNewTokenExpectations addObject:[self expectationWithDescription:name]];
+  }
+
+  __block NSInteger fetchNewTokenCallCount = 0;
+  __block FIRInstanceIDTokenHandler tokenHandler;
+
+  [[[self.mockTokenManager stub] andDo:^(NSInvocation *invocation) {
+    [invocation getArgument:&tokenHandler atIndex:6];
+    [fetchNewTokenExpectations[fetchNewTokenCallCount] fulfill];
+    fetchNewTokenCallCount += 1;
+  }] fetchNewTokenWithAuthorizedEntity:kAuthorizedEntity
+                                 scope:kFIRInstanceIDDefaultTokenScope
+                               keyPair:[OCMArg any]
+                               options:[OCMArg any]
+                               handler:[OCMArg any]];
+
+  // Mock Instance ID's retry interval to 0, to vastly speed up this test.
+  [[[self.mockInstanceID stub] andReturnValue:@(0)] retryIntervalToFetchDefaultToken];
+
+  // Make 1st call
+  XCTestExpectation *handlerExpectation1 = [self expectationWithDescription:@"handlerExpectation1"];
+  FIRInstanceIDResultHandler handler1 =
+      ^(FIRInstanceIDResult *_Nullable result, NSError *_Nullable error) {
+        [handlerExpectation1 fulfill];
+        XCTAssertNil(result);
+        XCTAssertNotNil(error);
+      };
+
+  [self.mockInstanceID instanceIDWithHandler:handler1];
+
+  // Make 2nd call
+  XCTestExpectation *handlerExpectation2 = [self expectationWithDescription:@"handlerExpectation1"];
+  FIRInstanceIDResultHandler handler2 =
+      ^(FIRInstanceIDResult *_Nullable result, NSError *_Nullable error) {
+        [handlerExpectation2 fulfill];
+        XCTAssertNil(result);
+        XCTAssertNotNil(error);
+      };
+
+  [self.mockInstanceID instanceIDWithHandler:handler2];
+
+  for (NSInteger i = 0; i < [[self.instanceID class] maxRetryCountForDefaultToken]; ++i) {
+    // Wait for the i `fetchNewTokenWithAuthorizedEntity` to be performed
+    [self waitForExpectations:@[ fetchNewTokenExpectations[i] ] timeout:1 enforceOrder:false];
+    // Fail for the i time
+    tokenHandler(nil, [NSError errorWithFIRInstanceIDErrorCode:kFIRInstanceIDErrorCodeUnknown]);
+  }
+
+  // Wait for completion handlers for both calls to be performed
+  [self waitForExpectationsWithTimeout:1 handler:NULL];
+}
+
 /**
  *  Tests a Keychain read failure while we try to fetch a new InstanceID token. If the Keychain
  *  read fails we won't be able to fetch the public key which is required while fetching a new

Firebase/InstanceID/FIRInstanceID.m

@@ -24,6 +24,7 @@
 #import <GoogleUtilities/GULAppEnvironmentUtil.h>
 #import "FIRInstanceID+Private.h"
 #import "FIRInstanceIDAuthService.h"
+#import "FIRInstanceIDCombinedHandler.h"
 #import "FIRInstanceIDConstants.h"
 #import "FIRInstanceIDDefines.h"
 #import "FIRInstanceIDKeyPairStore.h"
@@ -114,9 +115,9 @@ @interface FIRInstanceID ()
 @property(nonatomic, readwrite, strong) FIRInstanceIDKeyPairStore *keyPairStore;
 
 // backoff and retry for default token
-@property(atomic, readwrite, assign) BOOL isFetchingDefaultToken;
-@property(atomic, readwrite, assign) BOOL isDefaultTokenFetchScheduled;
 @property(nonatomic, readwrite, assign) NSInteger retryCountForDefaultToken;
+@property(atomic, strong, nullable)
+    FIRInstanceIDCombinedHandler<NSString *> *defaultTokenFetchHandler;
 
 @end
 
@@ -831,10 +832,30 @@ - (NSInteger)retryIntervalToFetchDefaultToken {
       kMaxRetryIntervalForDefaultTokenInSeconds);
 }
 
-- (void)defaultTokenWithHandler:(FIRInstanceIDTokenHandler)handler {
-  if (self.isFetchingDefaultToken || self.isDefaultTokenFetchScheduled) {
+- (void)defaultTokenWithHandler:(nullable FIRInstanceIDTokenHandler)aHandler {
+  [self defaultTokenWithRetry:NO handler:aHandler];
+}
+
+/**
+ * @param retry Indicates if the method is called to perform a retry after a failed attempt.
+ * If `YES`, then actual token request will be performed even if `self.defaultTokenFetchHandler !=
+ * nil`
+ */
+- (void)defaultTokenWithRetry:(BOOL)retry handler:(nullable FIRInstanceIDTokenHandler)aHandler {
+  BOOL shouldPerformRequest = retry || self.defaultTokenFetchHandler == nil;
+
+  if (!self.defaultTokenFetchHandler) {
+    self.defaultTokenFetchHandler = [[FIRInstanceIDCombinedHandler<NSString *> alloc] init];
+  }
+
+  if (aHandler) {
+    [self.defaultTokenFetchHandler addHandler:aHandler];
+  }
+
+  if (!shouldPerformRequest) {
     return;
   }
+
   NSDictionary *instanceIDOptions = @{};
   BOOL hasFirebaseMessaging = NSClassFromString(kFIRInstanceIDFCMSDKClassString) != nil;
   if (hasFirebaseMessaging && self.apnsTokenData) {
@@ -851,7 +872,6 @@ - (void)defaultTokenWithHandler:(FIRInstanceIDTokenHandler)handler {
   FIRInstanceID_WEAKIFY(self);
   FIRInstanceIDTokenHandler newHandler = ^void(NSString *token, NSError *error) {
     FIRInstanceID_STRONGIFY(self);
-    self.isFetchingDefaultToken = NO;
 
     if (error) {
       FIRInstanceIDLoggerError(kFIRInstanceIDMessageCodeInstanceID009,
@@ -871,21 +891,12 @@ - (void)defaultTokenWithHandler:(FIRInstanceIDTokenHandler)handler {
       // Do not retry beyond the maximum limit.
       if (self.retryCountForDefaultToken < [[self class] maxRetryCountForDefaultToken]) {
         NSInteger retryInterval = [self retryIntervalToFetchDefaultToken];
-        FIRInstanceID_WEAKIFY(self);
-        self.isDefaultTokenFetchScheduled = YES;
-        dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(retryInterval * NSEC_PER_SEC)),
-                       dispatch_get_main_queue(), ^{
-                         FIRInstanceID_STRONGIFY(self);
-                         self.isDefaultTokenFetchScheduled = NO;
-                         [self defaultTokenWithHandler:handler];
-                       });
+        [self retryGetDefaultTokenAfter:retryInterval];
       } else {
         FIRInstanceIDLoggerError(kFIRInstanceIDMessageCodeInstanceID007,
                                  @"Failed to retrieve the default FCM token after %ld retries",
                                  (long)self.retryCountForDefaultToken);
-        if (handler) {
-          handler(nil, error);
-        }
+        [self performDefaultTokenHandlerWithToken:nil error:error];
       }
     } else {
       // If somebody updated IID with APNS token while our initial request did not have it
@@ -904,13 +915,7 @@ - (void)defaultTokenWithHandler:(FIRInstanceIDTokenHandler)handler {
       if (!APNSRemainedSameDuringFetch && hasFirebaseMessaging) {
         // APNs value did change mid-fetch, so the token should be re-fetched with the current APNs
         // value.
-        self.isDefaultTokenFetchScheduled = YES;
-        FIRInstanceID_WEAKIFY(self);
-        dispatch_async(dispatch_get_main_queue(), ^{
-          FIRInstanceID_STRONGIFY(self);
-          self.isDefaultTokenFetchScheduled = NO;
-          [self defaultTokenWithHandler:handler];
-        });
+        [self retryGetDefaultTokenAfter:0];
         FIRInstanceIDLoggerDebug(kFIRInstanceIDMessageCodeRefetchingTokenForAPNS,
                                  @"Received APNS token while fetching default token. "
                                  @"Refetching default token.");
@@ -934,20 +939,41 @@ - (void)defaultTokenWithHandler:(FIRInstanceIDTokenHandler)handler {
                                           object:[self.defaultFCMToken copy]];
         [[NSNotificationQueue defaultQueue] enqueueNotification:tokenRefreshNotification
                                                    postingStyle:NSPostASAP];
-      }
-      if (handler) {
-        handler(token, nil);
+
+        [self performDefaultTokenHandlerWithToken:token error:nil];
       }
     }
   };
 
-  self.isFetchingDefaultToken = YES;
   [self tokenWithAuthorizedEntity:self.fcmSenderID
                             scope:kFIRInstanceIDDefaultTokenScope
                           options:instanceIDOptions
                           handler:newHandler];
 }
 
+/**
+ *
+ */
+- (void)performDefaultTokenHandlerWithToken:(NSString *)token error:(NSError *)error {
+  if (!self.defaultTokenFetchHandler) {
+    return;
+  }
+
+  [self.defaultTokenFetchHandler combinedHandler](token, error);
+  self.defaultTokenFetchHandler = nil;
+}
+
+- (void)retryGetDefaultTokenAfter:(NSTimeInterval)retryInterval {
+  FIRInstanceID_WEAKIFY(self);
+  dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(retryInterval * NSEC_PER_SEC)),
+                 dispatch_get_main_queue(), ^{
+                   FIRInstanceID_STRONGIFY(self);
+                   // Pass nil: no new handlers to be added, currently existing handlers
+                   // will be called
+                   [self defaultTokenWithRetry:YES handler:nil];
+                 });
+}
+
 #pragma mark - APNS Token
 // This should only be triggered from FCM.
 - (void)notifyAPNSTokenIsSet:(NSNotification *)notification {

Firebase/InstanceID/FIRInstanceIDCombinedHandler.h

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2019 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import <Foundation/Foundation.h>
+
+NS_ASSUME_NONNULL_BEGIN
+
+/**
+ * A generic class to combine several handler blocks into a single block in a thread-safe manner
+ */
+@interface FIRInstanceIDCombinedHandler<ResultType> : NSObject
+
+- (void)addHandler:(void (^)(ResultType _Nullable result, NSError* _Nullable error))handler;
+- (void (^)(ResultType _Nullable result, NSError* _Nullable error))combinedHandler;
+
+@end
+
+NS_ASSUME_NONNULL_END