Add cancel method for PINRemoteImageManager (#509)

* Add cancel method for PINRemoteImageManager

* Add comment and test

* Add changelog

CHANGELOG.md

@@ -3,6 +3,7 @@
 - [new] Add PINRemoteImageManagerConfiguration configuration object. [#492](https://github.com/pinterest/PINRemoteImage/pull/492) [rqueue](https://github.com/rqueue)
 - [fixed] Fixes blending in animated WebP images. [#507](https://github.com/pinterest/PINRemoteImage/pull/507) [garrettmoon](https://github.com/garrettmoon)
 - [fixed] Fixes support in PINAnimatedImageView for WebP animated images. [#507](https://github.com/pinterest/PINRemoteImage/pull/507) [garrettmoon](https://github.com/garrettmoon)
+- [new] Add cancel method for PINRemoteImageManager. [#509](https://github.com/pinterest/PINRemoteImage/pull/509) [zhongwuzw](https://github.com/zhongwuzw)
 
 ## 3.0.0 Beta 14
 - [fixed] Re-enable warnings check [#506](https://github.com/pinterest/PINRemoteImage/pull/506) [garrettmoon](https://github.com/garrettmoon)

Source/Classes/PINRemoteImageManager.h

@@ -659,6 +659,16 @@ typedef void(^PINRemoteImageManagerMetrics)(NSURL  * __nonnull url, NSURLSession
  */
 - (void)cancelTaskWithUUID:(nonnull NSUUID *)UUID storeResumeData:(BOOL)storeResumeData;
 
+/**
+ Cancel all tasks.
+ */
+- (void)cancelAllTasks;
+
+/**
+ Cancel all tasks and store resume data if any task has.
+ */
+- (void)cancelAllTasksAndStoreResumeData:(BOOL)storeResumeData;
+
 /**
  Set the priority of a download task. Since there is only one task per download, the priority of the download task will always be the last priority this method was called with.
  

Source/Classes/PINRemoteImageManager.m

@@ -115,6 +115,7 @@ @interface PINRemoteImageManager () <PINURLSessionManagerDelegate>
 @property (nonatomic, strong) PINURLSessionManager *sessionManager;
 @property (nonatomic, strong) NSMutableDictionary <NSString *, __kindof PINRemoteImageTask *> *tasks;
 @property (nonatomic, strong) NSHashTable <NSUUID *> *canceledTasks;
+@property (nonatomic, strong) NSHashTable <NSUUID *> *UUIDs;
 @property (nonatomic, strong) NSArray <NSNumber *> *progressThresholds;
 @property (nonatomic, assign) BOOL shouldBlurProgressive;
 @property (nonatomic, assign) CGSize maxProgressiveRenderSize;
@@ -240,6 +241,7 @@ -(nonnull instancetype)initWithSessionConfiguration:(NSURLSessionConfiguration *
         _maxProgressiveRenderSize = configuration.maxProgressiveRenderSize;
         self.tasks = [[NSMutableDictionary alloc] init];
         self.canceledTasks = [[NSHashTable alloc] initWithOptions:NSHashTableWeakMemory capacity:5];
+        self.UUIDs = [NSHashTable weakObjectsHashTable];
 
         if (alternateRepProvider == nil) {
             _defaultAlternateRepresentationProvider = [[PINAlternateRepresentationProvider alloc] init];
@@ -695,6 +697,8 @@ - (NSUUID *)downloadImageWithURL:(NSURL *)url
             }
             [task addCallbacksWithCompletionBlock:completion progressImageBlock:progressImage progressDownloadBlock:progressDownload withUUID:UUID];
             [self.tasks setObject:task forKey:key];
+            // Relax :), task retain the UUID for us, it's ok to have a weak reference to UUID here.
+            [self.UUIDs addObject:UUID];
 
             NSAssert(taskClass == [task class], @"Task class should be the same!");
         [self unlock];
@@ -1057,6 +1061,23 @@ - (void)cancelTaskWithUUID:(nonnull NSUUID *)UUID storeResumeData:(BOOL)storeRes
     } withPriority:PINOperationQueuePriorityHigh];
 }
 
+- (void)cancelAllTasks
+{
+    [self cancelAllTasksAndStoreResumeData:NO];
+}
+
+- (void)cancelAllTasksAndStoreResumeData:(BOOL)storeResumeData
+{
+    [_concurrentOperationQueue scheduleOperation:^{
+        [self lock];
+            NSArray<NSUUID *> *uuidToTask = [self.UUIDs allObjects];
+        [self unlock];
+        for (NSUUID *uuid in uuidToTask) {
+            [self cancelTaskWithUUID:uuid storeResumeData:storeResumeData];
+        }
+    } withPriority:PINOperationQueuePriorityHigh];
+}
+
 - (void)setPriority:(PINRemoteImageManagerPriority)priority ofTaskWithUUID:(NSUUID *)UUID
 {
     if (UUID == nil) {

Source/Classes/PINRemoteImageTask.m

@@ -13,7 +13,7 @@
 
 @interface PINRemoteImageTask ()
 {
-    NSMutableDictionary<NSUUID *, PINRemoteImageCallbacks *> *_callbackBlocks;
+    NSMutableDictionary<NSUUID *, PINRemoteImageCallbacks *> *_callbackBlocks; // We need to copy/retain `NSUUID`, because `PINRemoteImageManager` has a weak table `UUIDs` to store all UUIDs.
 }
 
 @end

Tests/PINRemoteImageTests.m

@@ -24,6 +24,7 @@
 #import <objc/runtime.h>
 
 static BOOL requestRetried = NO;
+static NSInteger canceledCount = 0;
 
 static inline BOOL PINImageAlphaInfoIsOpaque(CGImageAlphaInfo info) {
 	switch (info) {
@@ -68,6 +69,12 @@ - (NSString *)resumeCacheKeyForURL:(NSURL *)url;
 
 @end
 
+@interface PINRemoteImageManager (Swizzled)
+
+- (void)swizzled_cancelTaskWithUUID:(nonnull NSUUID *)UUID storeResumeData:(BOOL)storeResumeData;
+
+@end
+
 @interface PINURLSessionManager ()
 
 @property (nonatomic, strong) NSURLSession *session;
@@ -1398,6 +1405,58 @@ - (void)testRetry
   method_exchangeImplementations(originalMethod, swizzledMethod);
 }
 
+- (void)testCancelAllTasks
+{
+    XCTestExpectation *expectation = [self expectationWithDescription:@"Cancel all tasks"];
+    dispatch_group_t group = dispatch_group_create();
+
+    dispatch_group_enter(group);
+    [self.imageManager downloadImageWithURL:[self transparentPNGURL] options:0 progressDownload:^(int64_t completedBytes, int64_t totalBytes) {
+        static dispatch_once_t onceToken;
+        dispatch_once(&onceToken, ^{
+            dispatch_group_leave(group);
+        });
+    } completion:nil];
+
+    dispatch_group_enter(group);
+    [self.imageManager downloadImageWithURL:[self progressiveURL] options:0 progressDownload:^(int64_t completedBytes, int64_t totalBytes) {
+        static dispatch_once_t onceToken;
+        dispatch_once(&onceToken, ^{
+            dispatch_group_leave(group);
+        });
+    } completion:nil];
+
+    dispatch_group_notify(group, dispatch_get_main_queue(), ^{
+        SEL originalSelector = @selector(cancelTaskWithUUID:storeResumeData:);
+        SEL swizzledSelector = @selector(swizzled_cancelTaskWithUUID:storeResumeData:);
+
+        Method originalMethod = class_getInstanceMethod([PINRemoteImageManager class], originalSelector);
+        Method swizzledMethod = class_getInstanceMethod([PINRemoteImageManager class], swizzledSelector);
+
+        method_exchangeImplementations(originalMethod, swizzledMethod);
+
+        [self.imageManager cancelAllTasks];
+
+        // Give manager 2 seconds to cancel all tasks
+        dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(2 * NSEC_PER_SEC)), dispatch_get_main_queue(), ^{
+            XCTAssert(canceledCount == 2);
+            [expectation fulfill];
+        });
+    });
+    dispatch_group_wait(group, [self timeout]);
+    [self waitForExpectationsWithTimeout:[self timeoutTimeInterval] handler:nil];
+}
+
+@end
+
+@implementation PINRemoteImageManager (Swizzled)
+
+- (void)swizzled_cancelTaskWithUUID:(nonnull NSUUID *)UUID storeResumeData:(BOOL)storeResumeData
+{
+    canceledCount++;
+    [self swizzled_cancelTaskWithUUID:UUID storeResumeData:storeResumeData];
+}
+
 @end
 
 @implementation PINRemoteImageDownloadTask (Swizzled)