The service host limit should call the completion handler for stopFetching when enabled

Sources/Core/GTMSessionFetcher.m

@@ -124,6 +124,19 @@
 #define GTM_SWIFT_DISABLE_ASYNC
 #endif
 
+// Internal tracking of the delayed state.
+//
+// This is currently not complete, it is being worked in to better track things
+// to ensure `stopFetching` always triggers callback when requested.
+typedef NS_ENUM(NSUInteger, GTMSessionFetcherDelayState) {
+  kDelayStateNotDelayed = 0,
+  kDelayStateServiceDelayed,
+  kDelayStateCalculatingUA,
+  kDelayStateAuthorizing,
+  // TODO(thomasvl): More to be added as needed.
+  // ApplyingDecorators?
+};
+
 @interface GTMSessionFetcher ()
 
 @property(atomic, strong, readwrite, nullable) NSData *downloadedData;
@@ -240,6 +253,8 @@ @implementation GTMSessionFetcher {
   NSDate *_initialRequestDate;     // date of first request to the target server (ignoring auth)
   BOOL _hasAttemptedAuthRefresh;   // accessed only in shouldRetryNowForStatus:
 
+  GTMSessionFetcherDelayState _delayState;
+
   NSString *_comment;  // comment for log
   NSString *_log;
 #if !STRIP_GTM_FETCH_LOGGING
@@ -487,6 +502,11 @@ - (void)beginFetchMayDelay:(BOOL)mayDelay
   // This is the internal entry point for re-starting fetches.
   GTMSessionCheckNotSynchronized(self);
 
+  // Reset the delayed state since things are starting over.
+  @synchronized(self) {
+    _delayState = kDelayStateNotDelayed;
+  }
+
   NSMutableURLRequest *fetchRequest =
       _request;  // The request property is now externally immutable.
   NSURL *fetchRequestURL = fetchRequest.URL;
@@ -686,17 +706,46 @@ - (void)beginFetchMayDelay:(BOOL)mayDelay
     mayDelay = NO;
   }
   if (mayDelay && _service) {
+    BOOL savedStoppedState;
+    @synchronized(self) {
+      savedStoppedState = _userStoppedFetching;
+      // Set the delayed state so there can't be a race incase between it getting queued in
+      // the `fetcherShouldBeginFetching:` call and some other thread completing a different
+      // fetch and thus starting this one. If we were trying to set the state based on the
+      // return result, there would be a small window for that race.
+      _delayState = kDelayStateServiceDelayed;
+    }
+
     BOOL shouldFetchNow = [_service fetcherShouldBeginFetching:self];
     if (!shouldFetchNow) {
       // The fetch is deferred, but will happen later.
       //
       // If this session is held by the fetcher service, clear the session now so that we don't
       // assume it's still valid after the fetcher is restarted.
+      //
+      // NOTE: In hindsight, this could be a race, some other thread could be starting it while
+      // doing these checks/work, but it also doesn't seem safe to put the whole block in a
+      // since `@synchronized(self)` block.
       if (self.canShareSession) {
         self.session = nil;
       }
       return;
     }
+
+    @synchronized(self) {
+      // Per comment above, correct state since it wasn't delayed.
+      _delayState = kDelayStateNotDelayed;
+
+      // If a `-stopFetching` came in while the service check was made, then the side effect of the
+      // state setting caused the handler (if needed) to already be made, so there we want to just
+      // exit and not continue the fetch.
+      //
+      // TODO(thomasvl): If `savedStoppedState` was already `YES`, then it's a more general problem
+      // which will be handled elsewhere/later.
+      if (!savedStoppedState && savedStoppedState != _userStoppedFetching) {
+        return;
+      }
+    }
   }
 
   if ([fetchRequest valueForHTTPHeaderField:@"User-Agent"] == nil) {
@@ -1016,6 +1065,9 @@ - (void)updateUserAgentAsynchronouslyForRequest:(NSMutableURLRequest *)fetchRequ
                                     mayDecorate:(BOOL)mayDecorate {
   GTMSESSION_LOG_DEBUG_VERBOSE(
       @"GTMSessionFetcher fetching User-Agent from GTMUserAgentProvider %@...", _userAgentProvider);
+  @synchronized(self) {
+    _delayState = kDelayStateCalculatingUA;
+  }
   __weak __typeof__(self) weakSelf = self;
   dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
     __strong __typeof__(self) strongSelf = weakSelf;
@@ -1736,6 +1788,10 @@ - (void)endBackgroundTask {
 - (void)authorizeRequest {
   GTMSessionCheckNotSynchronized(self);
 
+  @synchronized(self) {
+    _delayState = kDelayStateAuthorizing;
+  }
+
   id authorizer = self.authorizer;
   // Prefer the block-based implementation. This *is* a change in behavior, but if authorizers
   // previously provided this method they would presumably assume they can be used for the same
@@ -2013,14 +2069,28 @@ - (void)forgetSessionIdentifierForFetcherWithoutSyncCheck {
 
 // External stop method
 - (void)stopFetching {
+  BOOL triggerCallback;
   @synchronized(self) {
     GTMSessionMonitorSynchronized(self);
 
     // Prevent enqueued callbacks from executing. The completion handler will still execute if
     // the property `stopFetchingTriggersCompletionHandler` is `YES`.
     _userStoppedFetching = YES;
+
+    // `stopFetchReleasingCallbacks:` will dequeue it if there is a sevice throttled
+    // delay, so the canceled callback needs to be directly triggered since the serivce
+    // won't attempt to restart it.
+    triggerCallback = _delayState == kDelayStateServiceDelayed && self.stopFetchingTriggersCompletionHandler;
   }  // @synchronized(self)
-  [self stopFetchReleasingCallbacks:!self.stopFetchingTriggersCompletionHandler];
+
+  if (triggerCallback) {
+    NSError *error = [NSError errorWithDomain:kGTMSessionFetcherErrorDomain
+                                         code:GTMSessionFetcherErrorUserCancelled
+                                     userInfo:nil];
+    [self finishWithError:error shouldRetry:NO];
+  } else {
+    [self stopFetchReleasingCallbacks:!self.stopFetchingTriggersCompletionHandler];
+  }
 }
 
 // Cancel the fetch of the URL that's currently in progress.

UnitTests/GTMSessionFetcherServiceTest.m

@@ -341,21 +341,23 @@ - (void)testFetcherService {
   NSMutableArray *urlArray = [NSMutableArray array];
   for (int idx = 1; idx <= 4; idx++) [urlArray addObject:validFileURL];
   [urlArray addObject:invalidFileURL];
+  // Ensure things will get queued due to the per host limit.
+  XCTAssertGreaterThan(urlArray.count, kMaxRunningFetchersPerHost);
   for (int idx = 1; idx <= 5; idx++) [urlArray addObject:validFileURL];
   for (int idx = 1; idx <= 3; idx++) [urlArray addObject:stopSilentFileURL];
   for (int idx = 1; idx <= 5; idx++) [urlArray addObject:altValidURL];
-  // TODO(thomasvl) currently fail: The current code paths remove these from
-  // pending, but doesn't actually get the completions ever called as requested.
-  // Will be fixed in follow up PRs.
-  // for (int idx = 1; idx <= 3; idx++) [urlArray addObject:stopCallbackFileURL];
+  for (int idx = 1; idx <= 3; idx++) [urlArray addObject:stopCallbackFileURL];
   for (int idx = 1; idx <= 5; idx++) [urlArray addObject:validFileURL];
+  for (int idx = 1; idx <= 2; idx++) [urlArray addObject:stopCallbackFileURL];
   for (int idx = 1; idx <= 2; idx++) [urlArray addObject:stopSilentFileURL];
   NSUInteger totalNumberOfFetchers = urlArray.count;
 
+  // These are used to track the fetchers via notifications.
   __block NSMutableArray *pending = [NSMutableArray array];
   __block NSMutableArray *running = [NSMutableArray array];
   __block NSMutableArray *completed = [NSMutableArray array];
-  NSMutableArray *stoppedNoCallback = [NSMutableArray array];
+  // Fetchers that will get stopped.
+  NSMutableArray *stopped = [NSMutableArray array];
 
   NSInteger priorityVal = 0;
 
@@ -376,16 +378,15 @@ - (void)testFetcherService {
                     object:fetcher
                      queue:nil
                 usingBlock:^(NSNotification *note) {
+                  XCTAssertTrue([pending containsObject:fetcher]);
+
                   // Verify that we have at most two fetchers running for this fetcher's host.
                   [running addObject:fetcher];
                   [pending removeObject:fetcher];
 
                   NSURLRequest *fetcherReq = fetcher.request;
                   NSURL *fetcherReqURL = fetcherReq.URL;
 
-                  // A fetcher that gets stopped should not trigger any notifications.
-                  XCTAssertFalse([fetcherReqURL.query hasPrefix:@"stop="]);
-
                   NSString *host = fetcherReqURL.host;
                   NSUInteger numberRunning = FetchersPerHost(running, host);
                   XCTAssertTrue(numberRunning > 0, @"count error");
@@ -418,6 +419,8 @@ - (void)testFetcherService {
                     object:fetcher
                      queue:nil
                 usingBlock:^(NSNotification *note) {
+                  XCTAssertTrue([running containsObject:fetcher]);
+
                   // Verify that we only have two fetchers running.
                   [completed addObject:fetcher];
                   [running removeObject:fetcher];
@@ -430,7 +433,7 @@ - (void)testFetcherService {
                   NSUInteger numberRunning = FetchersPerHost(running, host);
                   NSUInteger numberPending = FetchersPerHost(pending, host);
                   NSUInteger numberCompleted = FetchersPerHost(completed, host);
-                  NSUInteger numberStopped = FetchersPerHost(stoppedNoCallback, host);
+                  NSUInteger numberStopped = FetchersPerHost(stopped, host);
 
                   XCTAssertLessThanOrEqual(numberRunning, kMaxRunningFetchersPerHost,
                                            @"too many running");
@@ -447,8 +450,6 @@ - (void)testFetcherService {
                 }];
     [observers addObject:stopObserver];
 
-    [pending addObject:fetcher];
-
     // Set the fetch priority to a value that cycles 0, 1, -1, 0, ...
     priorityVal++;
     if (priorityVal > 1) priorityVal = -1;
@@ -460,6 +461,17 @@ - (void)testFetcherService {
       fetcher.stopFetchingTriggersCompletionHandler = YES;
     }
 
+    if (stopFetch) {
+      // All the fetchers that will be stopped should end up queued, as we're
+      // also testing that they are handled correctly for callbacks and the
+      // general per host limits. So make sure we've got atleast the max per
+      // host pending already.
+      NSUInteger numberPending = FetchersPerHost(pending, fetcher.request.URL.host);
+      XCTAssertGreaterThanOrEqual(numberPending, kMaxRunningFetchersPerHost);
+    } else {
+      [pending addObject:fetcher];
+    }
+
     // Start this fetcher.
     [fetchersInFlight addObject:fetcher];
     [fetcher beginFetchWithCompletionHandler:^(NSData *fetchData, NSError *fetchError) {
@@ -489,9 +501,10 @@ - (void)testFetcherService {
     }];
     if (stopFetch) {
       [fetcher stopFetching];
-      [pending removeObject:fetcher];
-      if (!stopWithCallback) {
-        [stoppedNoCallback addObject:fetcher];
+      [stopped addObject:fetcher];
+      if (stopWithCallback) {
+        // The completion will remove it from inflight and complete the expectation.
+      } else {
         [fetchersInFlight removeObject:fetcher];
         [expectation fulfill];
       }
@@ -500,10 +513,12 @@ - (void)testFetcherService {
 
   [self waitForExpectations:expectations timeout:_timeoutInterval];
 
+  // Check the notification counts.
   XCTAssertEqual(pending.count, (NSUInteger)0, @"still pending: %@", pending);
   XCTAssertEqual(running.count, (NSUInteger)0, @"still running: %@", running);
-  XCTAssertEqual(completed.count + stoppedNoCallback.count, (NSUInteger)totalNumberOfFetchers,
+  XCTAssertEqual(completed.count + stopped.count, (NSUInteger)totalNumberOfFetchers,
                  @"incomplete");
+  // All the expected callbacks happened.
   XCTAssertEqual(fetchersInFlight.count, (NSUInteger)0, @"Uncompleted: %@", fetchersInFlight);
 
   XCTAssertEqual([service numberOfFetchers], (NSUInteger)0, @"service non-empty");
@@ -555,6 +570,10 @@ - (void)testStopAllFetchersSeparately {
   [self stopAllFetchersHelperUseStopAllAPI:NO callbacksAfterStop:NO];
 }
 
+- (void)testStopAllFetchersWithCallbacks {
+  [self stopAllFetchersHelperUseStopAllAPI:YES callbacksAfterStop:YES];
+}
+
 - (void)testStopAllFetchersSeparatelyWithCallbacks {
   [self stopAllFetchersHelperUseStopAllAPI:NO callbacksAfterStop:YES];
 }
@@ -582,9 +601,10 @@ - (void)stopAllFetchersHelperUseStopAllAPI:(BOOL)useStopAllAPI
   XCTestExpectation *fetcherCallbackExpectation =
       [[XCTNSNotificationExpectation alloc] initWithName:@"callback"];
   if (doStopCallbacks) {
-    fetcherCallbackExpectation.expectedFulfillmentCount = 4;
+    fetcherCallbackExpectation.expectedFulfillmentCount = urlArray.count;
     fetcherCallbackExpectation.assertForOverFulfill = YES;
   } else {
+    // No callbacks, just complete this expectation now.
     [fetcherCallbackExpectation fulfill];
   }